From b8beabb47e525a132320b2af8c3ad2412f652a31 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Sun, 26 Jul 2026 19:00:14 -0700 Subject: [PATCH 1/6] feat(grpc): restack native vLLM sidecar protocol onto current main --- rust/README.md | 15 + rust/proto/control.proto | 57 +- rust/proto/inference.proto | 84 +- rust/src/chat/src/lib.rs | 25 +- rust/src/chat/src/multimodal.rs | 2 +- rust/src/chat/src/output/default/unified.rs | 1 + rust/src/chat/src/output/harmony/tests.rs | 1 + rust/src/chat/tests/roundtrip.rs | 2 + rust/src/cmd/src/cli/tests.rs | 29 + rust/src/engine-core-client/src/client.rs | 187 ++-- rust/src/engine-core-client/src/client/imp.rs | 30 +- .../engine-core-client/src/client/state.rs | 5 + rust/src/engine-core-client/src/error.rs | 2 + .../src/engine-core-client/src/mock_engine.rs | 8 +- .../src/protocol/handshake.rs | 32 +- .../engine-core-client/src/protocol/output.rs | 143 +++- .../src/protocol/sampling.rs | 3 + rust/src/engine-core-client/src/test_utils.rs | 60 +- .../engine-core-client/src/tests/client.rs | 333 +++++--- .../src/tests/client/utility.rs | 212 +++++ .../src/tests/python_compat.py | 70 +- rust/src/engine-core-client/src/transport.rs | 327 ++++++- rust/src/llm/src/output.rs | 17 +- rust/src/server/Cargo.toml | 2 +- rust/src/server/src/grpc/control.rs | 200 ++++- rust/src/server/src/grpc/convert.rs | 799 +++++++++++++----- .../src/server/src/grpc/convert/multimodal.rs | 540 ++++++++++++ rust/src/server/src/grpc/convert/sampling.rs | 163 ++++ rust/src/server/src/grpc/convert/xargs.rs | 77 ++ rust/src/server/src/grpc/inference.rs | 86 +- rust/src/server/src/grpc/lora_rpc.rs | 172 ++++ rust/src/server/src/grpc/mod.rs | 51 ++ rust/src/server/src/grpc/struct_json.rs | 84 ++ rust/src/server/src/grpc/tests.rs | 14 +- rust/src/server/src/grpc/tests/lora.rs | 192 +++++ rust/src/server/src/lib.rs | 171 ++-- rust/src/server/src/lora.rs | 659 ++++++++++++++- rust/src/server/src/lora/tests/inplace.rs | 129 +++ rust/src/server/src/lora_path.rs | 248 ++++++ rust/src/server/src/routes.rs | 2 +- .../server/src/routes/inference/generate.rs | 13 +- .../src/routes/inference/generate/convert.rs | 1 + rust/src/server/src/routes/lora.rs | 192 +---- .../src/routes/openai/chat_completions.rs | 10 +- .../routes/openai/chat_completions/convert.rs | 2 + .../server/src/routes/openai/completions.rs | 16 +- .../src/routes/openai/completions/convert.rs | 2 + rust/src/server/src/routes/openai/models.rs | 7 +- rust/src/server/src/routes/tests.rs | 218 +++++ .../src/routes/tests/lora_concurrency.rs | 306 +++++++ rust/src/server/src/state.rs | 25 +- rust/src/text/src/lower.rs | 121 +-- rust/src/text/src/output/decoded.rs | 12 +- rust/src/text/src/output/mod.rs | 7 + rust/src/text/src/request.rs | 3 + .../token_in_token_out/test_mm_serde.py | 13 + tests/v1/engine/test_engine_core.py | 43 +- tests/v1/engine/test_engine_core_client.py | 3 + vllm/entrypoints/serve/disagg/__init__.py | 3 + vllm/entrypoints/serve/disagg/mm_serde.py | 10 + vllm/v1/engine/__init__.py | 11 +- vllm/v1/engine/core.py | 178 ++-- 62 files changed, 5460 insertions(+), 970 deletions(-) create mode 100644 rust/src/engine-core-client/src/tests/client/utility.rs create mode 100644 rust/src/server/src/grpc/convert/multimodal.rs create mode 100644 rust/src/server/src/grpc/convert/sampling.rs create mode 100644 rust/src/server/src/grpc/convert/xargs.rs create mode 100644 rust/src/server/src/grpc/lora_rpc.rs create mode 100644 rust/src/server/src/grpc/struct_json.rs create mode 100644 rust/src/server/src/grpc/tests/lora.rs create mode 100644 rust/src/server/src/lora/tests/inplace.rs create mode 100644 rust/src/server/src/lora_path.rs create mode 100644 rust/src/server/src/routes/tests/lora_concurrency.rs create mode 100644 vllm/entrypoints/serve/disagg/__init__.py create mode 100644 vllm/entrypoints/serve/disagg/mm_serde.py diff --git a/rust/README.md b/rust/README.md index b14aba3fae19..1934992ee7b0 100644 --- a/rust/README.md +++ b/rust/README.md @@ -87,3 +87,18 @@ curl http://127.0.0.1:8000/v1/chat/completions \ "stream": true }' ``` + +### Runtime LoRA adapters + +Dynamic LoRA loading requires `VLLM_ALLOW_RUNTIME_LORA_UPDATING=1`. When an adapter is loaded from +the local filesystem, `VLLM_RUNTIME_LORA_ALLOWED_PATH_PREFIXES` is also required. Set it to a +platform-separated list of directories that may contain adapters; local paths outside those +directories are rejected. Hugging Face repository IDs do not require a local path prefix. + +For example: + +```bash +VLLM_ALLOW_RUNTIME_LORA_UPDATING=1 \ +VLLM_RUNTIME_LORA_ALLOWED_PATH_PREFIXES=/model-cache/loras:/workspace/loras \ +vllm serve Qwen/Qwen3-0.6B --enable-lora +``` diff --git a/rust/proto/control.proto b/rust/proto/control.proto index 0e25aea26474..21fe82477fb0 100644 --- a/rust/proto/control.proto +++ b/rust/proto/control.proto @@ -8,6 +8,12 @@ service Control { rpc GetServerInfo (GetServerInfoRequest) returns (ServerInfo) {} rpc GetModelInfo (GetModelInfoRequest) returns (ModelInfo) {} rpc Abort (AbortRequest) returns (AbortResponse) {} + rpc Drain (DrainRequest) returns (DrainResponse) {} + + rpc LoadLora (LoadLoraRequest) returns (LoadLoraResponse) {} + rpc UnloadLora (UnloadLoraRequest) returns (UnloadLoraResponse) {} + rpc ListLoras (ListLorasRequest) returns (ListLorasResponse) {} + rpc GetKvEventSources (GetKvEventSourcesRequest) returns (GetKvEventSourcesResponse) {} } @@ -23,6 +29,8 @@ message ServerInfo { uint64 total_kv_blocks = 7; uint64 max_running_requests = 8; uint64 max_batched_tokens = 9; + uint32 max_loras = 10; + repeated string capabilities = 11; } message ParallelismInfo { @@ -31,6 +39,8 @@ message ParallelismInfo { uint32 data_parallel_size = 3; uint32 data_parallel_rank = 4; uint32 decode_context_parallel_size = 5; + uint32 data_parallel_start_rank = 6; + uint32 managed_data_parallel_size = 7; } message GetModelInfoRequest {} @@ -39,9 +49,11 @@ message ModelInfo { string model_id = 1; string served_model_name = 2; repeated string served_model_aliases = 3; + repeated string tokenizer_modes = 4; bool supports_text_input = 20; bool supports_token_ids_input = 21; + bool supports_lora = 22; bool supports_multimodal = 23; string reasoning_parser = 24; string tool_call_parser = 25; @@ -53,6 +65,43 @@ message AbortRequest { message AbortResponse {} +message DrainRequest {} + +message DrainResponse { + DrainState state = 1; + uint32 in_flight_requests = 2; + string message = 3; +} + +enum DrainState { + DRAIN_STATE_UNSPECIFIED = 0; + DRAIN_STATE_IN_PROGRESS = 1; + DRAIN_STATE_COMPLETE = 2; +} + +// ====================================================================================== +// LoRA lifecycle +// ====================================================================================== + +message LoraAdapter { + int64 lora_id = 1; + string lora_name = 2; + string source_path = 3; +} + +message LoadLoraRequest { + LoraAdapter adapter = 1; + bool load_inplace = 2; +} +message LoadLoraResponse { + LoraAdapter adapter = 1; + bool already_loaded = 2; +} +message UnloadLoraRequest { string lora_name = 1; } +message UnloadLoraResponse { LoraAdapter adapter = 1; } +message ListLorasRequest {} +message ListLorasResponse { repeated LoraAdapter adapters = 1; } + // ====================================================================================== // KV discovery // ====================================================================================== @@ -60,9 +109,15 @@ message AbortResponse {} message GetKvEventSourcesRequest {} message GetKvEventSourcesResponse { repeated KvEventSource sources = 1; } +message KvEventEndpoint { + string host = 1; + uint32 port = 2; + string protocol = 3; +} + message KvEventSource { string transport = 1; - string endpoint = 2; + KvEventEndpoint endpoint_addr = 2; string topic = 3; string replay_endpoint = 4; optional uint32 data_parallel_rank = 5; diff --git a/rust/proto/inference.proto b/rust/proto/inference.proto index b1c08ae76e60..9493ee757d5a 100644 --- a/rust/proto/inference.proto +++ b/rust/proto/inference.proto @@ -1,6 +1,3 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project - syntax = "proto3"; package vllm; @@ -45,21 +42,32 @@ message GenerateRequest { uint32 truncate_prompt_tokens = 11; int32 priority = 12; + + // Multimodal inputs aligned with placeholder markers in token_ids. + repeated MediaItem media = 13; + // Loaded LoRA name; empty selects the base model. + string lora_name = 14; + // Lossless JSON object for engine-specific vLLM extensions. + optional bytes vllm_xargs_json = 15; + // Already-processed multimodal features from the token-in/token-out renderer. + repeated PreprocessedMultimodalFeature mm_features = 16; + // Skip this many prompt-token routing rows in the returned expert tensor. + uint32 routed_experts_prompt_start = 17; } message RandomSampling { uint32 num_sequences = 1; // "n", default (0) means 1 - uint32 top_k = 2; // 0 means default - float top_p = 3; // 0 means default - float min_p = 4; // 0 means default + optional int32 top_k = 2; + optional float top_p = 3; + optional float min_p = 4; optional int64 seed = 5; } message DecodingParameters { // Penalties - float presence_penalty = 1; // Default (0.0) means no penalty - float frequency_penalty = 2; // Default (0.0) means no penalty - float repetition_penalty = 3; // Default (0.0) means no penalty + optional float presence_penalty = 1; + optional float frequency_penalty = 2; + optional float repetition_penalty = 3; map logit_bias = 4; repeated uint32 allowed_token_ids = 5; @@ -76,6 +84,10 @@ message DecodingParameters { bool json_object = 10; string structural_tag = 11; } + bool structured_output_disable_any_whitespace = 12; + bool structured_output_disable_additional_properties = 13; + optional string structured_output_whitespace_pattern = 14; + repeated string bad_words = 15; } message StoppingCriteria { @@ -89,6 +101,7 @@ message StoppingCriteria { bool include_stop_strings = 5; bool ignore_eos = 6; + optional int64 thinking_token_budget = 7; } message ResponseOptions { @@ -110,9 +123,6 @@ message KVCacheParameters { // KV Connector transfer parameters google.protobuf.Struct kv_transfer_params = 3; - - // Encoder cache connector transfer parameters - google.protobuf.Struct ec_transfer_params = 4; } // Controls which extra candidate tokens at each position should be returned @@ -147,6 +157,13 @@ message SequenceOutput { // Only present in final output for this sequence optional FinishInfo finish_info = 8; + optional RoutedExpertsTensor routed_experts = 9; +} + +message RoutedExpertsTensor { + string dtype = 1; + repeated uint64 shape = 2; + bytes data = 3; } // Prompt info, returned in the first response @@ -179,7 +196,6 @@ message FinishInfo { google.protobuf.Struct kv_transfer_params = 6; //uint64 seed = 7; - google.protobuf.Struct ec_transfer_params = 8; } // Info for candidate tokens other than the input/sampled @@ -200,3 +216,45 @@ message CandidateTokenInfo { message TokenIds { repeated uint32 ids = 1; } + +// ====================================================================================== +// Media +// ====================================================================================== + +enum Modality { + MODALITY_UNSPECIFIED = 0; + MODALITY_IMAGE = 1; + MODALITY_VIDEO = 2; + MODALITY_AUDIO = 3; +} + +message MediaItem { + Modality modality = 1; + oneof source { + string url = 2; + string data_uri = 3; + bytes raw_bytes = 4; + } + string mime_type = 5; + string uuid = 6; +} + +message MultimodalPlaceholder { + uint64 offset = 1; + uint64 length = 2; + // Empty means every position in the placeholder is an embedding position. + repeated bool is_embed = 3; +} + +message PreprocessedMultimodalFeature { + string modality = 1; + // Canonical modality+kwargs identity. Kept equal to cache_identifier so + // Python vLLM cannot prefer an unverified renderer-provided cache key. + string mm_hash = 2; + MultimodalPlaceholder position = 3; + // Python vLLM's inline msgpack encoding. The current gRPC server requires it. + optional bytes kwargs_msgpack = 4; + // Canonical cache identity derived from modality and kwargs_msgpack. + string cache_identifier = 5; +} + diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index 0d538a24ba1a..a93fe72a2224 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -21,6 +21,7 @@ pub use event::{ AssistantToolCall, ChatEvent, }; use futures::{StreamExt, TryStreamExt as _}; +pub use llm_multimodal::MediaContentPart; pub use output::{ ChatOutputProcessor, DefaultChatOutputProcessor, DynChatOutputProcessor, HarmonyChatOutputProcessor, @@ -41,6 +42,7 @@ pub use request::{ ChatToolChoice, GenerationPromptMode, ReasoningEffort, SamplingParams, }; pub use stream::{ChatEventStream, ChatEventStreamTrait, CollectedAssistantMessage}; +pub use vllm_engine_core_client::protocol::multimodal::MmFeatures; pub use vllm_llm::FinishReason; mod backend; @@ -55,7 +57,6 @@ mod stream; use vllm_engine_core_client::EngineCoreClient; use vllm_engine_core_client::protocol::dtype::ModelDtype; -use vllm_engine_core_client::protocol::multimodal::MmFeatures; use vllm_engine_core_client::protocol::request::ReasoningParserKwargs; use vllm_llm::Llm; use vllm_text::{Prompt, TextLlm, TextRequest}; @@ -283,6 +284,28 @@ impl ChatLlm { } } + /// Prepare media for an already-tokenized request. + pub async fn prepare_media( + &self, + media: Vec, + token_ids: &mut Vec, + ) -> Result> { + if media.is_empty() { + return Ok(None); + } + let info = self + .processor + .backend + .multimodal_model_info() + .ok_or(Error::UnsupportedMultimodalRenderer)?; + let model_dtype = self + .processor + .model_dtype + .ok_or(Error::UnsupportedMultimodalRenderer)?; + let features = info.prepare_multimodal(media, token_ids, model_dtype).await?; + Ok(Some(features)) + } + /// Render, tokenize, and submit one chat request. pub async fn chat(&self, request: ChatRequest) -> Result { let (text_request, output_processor) = self diff --git a/rust/src/chat/src/multimodal.rs b/rust/src/chat/src/multimodal.rs index c0600a6329ba..edb4722d4824 100644 --- a/rust/src/chat/src/multimodal.rs +++ b/rust/src/chat/src/multimodal.rs @@ -696,7 +696,7 @@ impl MultimodalModelInfo { /// `prompt_token_ids` is mutated in place because placeholder expansion /// changes both the final prompt and the offsets recorded in /// `PlaceholderRange`. - async fn prepare_multimodal( + pub(crate) async fn prepare_multimodal( &self, media_parts: Vec, prompt_token_ids: &mut Vec, diff --git a/rust/src/chat/src/output/default/unified.rs b/rust/src/chat/src/output/default/unified.rs index f865e5a6ae0d..075526c771a4 100644 --- a/rust/src/chat/src/output/default/unified.rs +++ b/rust/src/chat/src/output/default/unified.rs @@ -392,6 +392,7 @@ mod tests { finish_reason: crate::FinishReason::Stop(None), kv_transfer_params: None, ec_transfer_params: None, + routed_experts: None, }), } } diff --git a/rust/src/chat/src/output/harmony/tests.rs b/rust/src/chat/src/output/harmony/tests.rs index 9b9544bdc6bd..37ab2e141687 100644 --- a/rust/src/chat/src/output/harmony/tests.rs +++ b/rust/src/chat/src/output/harmony/tests.rs @@ -56,6 +56,7 @@ fn finished() -> Finished { finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, + routed_experts: None, } } diff --git a/rust/src/chat/tests/roundtrip.rs b/rust/src/chat/tests/roundtrip.rs index 9cd670d29d08..8e18a325f532 100644 --- a/rust/src/chat/tests/roundtrip.rs +++ b/rust/src/chat/tests/roundtrip.rs @@ -722,6 +722,7 @@ fn decoded_completion_stream( finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, + routed_experts: None, }), } }); @@ -733,6 +734,7 @@ fn decoded_completion_stream( finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, + routed_experts: None, }); events.push(DecodedTextEvent::TextDelta { delta: chunk.delta, diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index e586477a897e..0a084ee53803 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -156,6 +156,35 @@ fn serve_args_auto_forward_enable_lora_to_python() { .assert_debug_eq(&args.managed_engine.python_args); } +#[test] +fn serve_args_forward_lora_capacity_to_python() { + let cli = Cli::try_parse_from([ + "vllm-rs", + "serve", + "Qwen/Qwen3-0.6B", + "--enable-lora", + "--max-loras", + "4", + "--max-lora-rank", + "64", + ]) + .unwrap(); + + let Command::Serve(args) = cli.command else { + panic!("expected serve args"); + }; + expect![[r#" + [ + "--enable-lora", + "--max-loras", + "4", + "--max-lora-rank", + "64", + ] + "#]] + .assert_debug_eq(&args.managed_engine.python_args); +} + #[test] fn serve_args_forward_shutdown_timeout_to_managed_engine() { let cli = Cli::try_parse_from([ diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index a5709e52d67f..9f872e07a931 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use std::time::Duration; -use futures::future::{join_all, try_join_all}; +use futures::future::join_all; use itertools::Itertools; use serde::Serialize; use tokio::sync::mpsc; @@ -32,14 +32,15 @@ pub use stream::{EngineCoreOutputStream, EngineCoreStreamOutput}; /// `EngineCoreProc`s. #[derive(Debug, Clone, PartialEq, Eq, Serialize)] pub enum TransportMode { - /// The Rust process owns the startup handshake and allocates or binds the - /// frontend transport addresses itself before replying to engine - /// `HELLO` messages. + /// The Rust process is the sole owner of the startup handshake and + /// allocates or binds the frontend transport addresses itself. Exactly one + /// owner must bind the handshake endpoint for the complete engine cohort. HandshakeOwner { /// Shared handshake endpoint that engines dial during startup. handshake_address: String, /// Host/IP that engines should use to connect back to the frontend - /// transport sockets. + /// transport sockets. It must be locally bindable and routable from + /// every engine; coordinated DP also advertises it as the master IP. advertised_host: String, /// Total number of engines expected to join this transport. engine_count: usize, @@ -202,6 +203,7 @@ pub struct EngineCoreClient { input_address: String, output_address: String, engines: Vec, + managed_data_parallel_span: transport::ManagedDataParallelSpan, inner: Arc, coordinator: Option, abort_tx: mpsc::UnboundedSender, @@ -216,6 +218,31 @@ pub struct EngineCoreClient { coordinator_task: Option>, } +/// Removes utility waiters if a call future is cancelled before completion. +struct UtilityCallGuard { + inner: Arc, + call_ids: Vec, +} + +impl UtilityCallGuard { + fn new(inner: Arc) -> Self { + Self { + inner, + call_ids: Vec::new(), + } + } + + fn track(&mut self, call_id: u64) { + self.call_ids.push(call_id); + } +} + +impl Drop for UtilityCallGuard { + fn drop(&mut self) { + self.inner.unregister_utility_calls(self.call_ids.drain(..)); + } +} + impl EngineCoreClient { /// Connect to Python `EngineCoreProc`s using the configured /// transport/coordinator modes. @@ -287,6 +314,7 @@ impl EngineCoreClient { let (output_tx, output_rx) = mpsc::channel(64); let (abort_tx, abort_rx) = mpsc::unbounded_channel(); let engines = connected.engines; + let managed_data_parallel_span = connected.managed_data_parallel_span; let runtime = build_zmq_runtime(); let inner = Arc::new(ClientInner::new( connected.input_send, @@ -346,6 +374,7 @@ impl EngineCoreClient { input_address: connected.input_address, output_address: connected.output_address, engines, + managed_data_parallel_span, inner, coordinator, abort_tx, @@ -375,6 +404,20 @@ impl EngineCoreClient { self.engines.len() } + /// Return the first rank and number of contiguous data-parallel ranks + /// validated during the engine transport handshake. + pub fn managed_data_parallel_span(&self) -> (u32, u32) { + ( + self.managed_data_parallel_span.start_rank, + self.managed_data_parallel_span.size, + ) + } + + #[cfg(test)] + pub(crate) fn pending_utility_call_count(&self) -> usize { + self.inner.pending_utility_call_count() + } + /// Return the engine-side indices connected to this client. pub fn engine_indices(&self) -> Vec { self.engines @@ -396,8 +439,11 @@ impl EngineCoreClient { /// Return the first engine's ready response. /// - /// Per-engine fields such as `data_parallel_rank` should be read through - /// [`ready_responses`](Self::ready_responses). + /// Engine topology, role, and limits (TP/PP size, KV block size, KV + /// connector role, scheduler caps) are uniform across the connected + /// engines, so the first engine's response is representative for discovery. + /// Per-engine values that differ (e.g. `data_parallel_rank`) should be read + /// via [`ready_responses`](Self::ready_responses) instead. pub fn ready_response(&self) -> &EngineCoreReadyResponse { &self .engines @@ -405,6 +451,7 @@ impl EngineCoreClient { .expect("engine core client requires at least one engine") .ready_response } + /// Return the engine-reported effective model dtype. pub fn model_dtype(&self) -> ModelDtype { self.engines @@ -558,14 +605,14 @@ impl EngineCoreClient { Ok(()) } - /// Call a typed utility method on all connected engines, returning one - /// decoded result per connected engine if all calls succeed or an error - /// if any call fails. - /// - /// Callers should pass utility arguments using Rust tuple semantics so the - /// encoded payload matches Python's `(client_index, call_id, - /// method_name, args)` contract: `()`, `(arg,)`, `(arg1, arg2)`, etc. - pub async fn call_utility(&self, method: &str, args: A) -> Result> + /// Call a typed utility method on all connected engines, preserving one + /// result per engine. The outer error is reserved for failures that happen + /// before any request can be dispatched. + pub async fn call_utility_per_engine( + &self, + method: &str, + args: A, + ) -> Result>> where T: serde::de::DeserializeOwned, A: serde::Serialize + std::fmt::Debug, @@ -577,20 +624,15 @@ impl EngineCoreClient { "sending utility request" ); - // Phase 1: allocate one call id per engine and build the per-engine - // request payloads up-front. Any failure here (registry closed, encode - // error) must roll back the call ids already allocated so they do not - // leak in the utility registry until shutdown. + let mut call_guard = UtilityCallGuard::new(self.inner.clone()); let mut pending_calls = Vec::with_capacity(self.engines.len()); let mut prepared_sends = Vec::with_capacity(self.engines.len()); for engine in &self.engines { let (call_id, rx) = match self.inner.allocate_and_register_utility_call() { Ok(pair) => pair, - Err(err) => { - self.inner.unregister_utility_calls(pending_calls.iter().map(|(id, _)| *id)); - return Err(err); - } + Err(err) => return Err(err), }; + call_guard.track(call_id); let request = match EngineCoreUtilityRequest::new( self.config.client_index, call_id, @@ -598,41 +640,86 @@ impl EngineCoreClient { &args, ) { Ok(request) => request, - Err(err) => { - self.inner.unregister_utility_calls( - pending_calls.iter().map(|(id, _)| *id).chain(std::iter::once(call_id)), - ); - return Err(err); - } + Err(err) => return Err(err), }; pending_calls.push((call_id, rx)); prepared_sends.push((&engine.engine_id, request)); } - // Phase 2: dispatch every utility request concurrently. `try_join_all` - // fails fast on the first transport error and drops the remaining send - // futures; any engines that already received the request will reply, - // but those replies are simply dropped because we roll back the call - // ids below. - let send_futures = prepared_sends.iter().map(|(engine_id, request)| { + let send_results = join_all(prepared_sends.iter().map(|(engine_id, request)| { self.inner.send_to_engine(engine_id, EngineCoreRequestType::Utility, request) - }); - if let Err(err) = try_join_all(send_futures).await { - self.inner.unregister_utility_calls(pending_calls.iter().map(|(id, _)| *id)); - return Err(err); + })) + .await; + let calls = pending_calls.into_iter().zip(send_results).map( + |((call_id, rx), send_result)| async move { + send_result?; + rx.await + .map_err(|_| Error::UtilityCallClosed { + method: method.to_string(), + call_id, + })?? + .into_typed_result(method) + }, + ); + Ok(join_all(calls).await) + } + + /// Call one utility method on the connected engine with this logical + /// data-parallel rank. + pub async fn call_utility_on_engine( + &self, + data_parallel_rank: u32, + method: &str, + args: A, + ) -> Result + where + T: serde::de::DeserializeOwned, + A: serde::Serialize + std::fmt::Debug, + { + let engine = self + .engines + .iter() + .find(|engine| engine.ready_response.data_parallel_rank == data_parallel_rank) + .ok_or(Error::InvalidDataParallelRank { + rank: data_parallel_rank, + num_engines: self.engines.len() as u32, + })?; + let mut call_guard = UtilityCallGuard::new(self.inner.clone()); + let (call_id, rx) = self.inner.allocate_and_register_utility_call()?; + call_guard.track(call_id); + let request = + match EngineCoreUtilityRequest::new(self.config.client_index, call_id, method, &args) { + Ok(request) => request, + Err(error) => return Err(error), + }; + if let Err(error) = self + .inner + .send_to_engine(&engine.engine_id, EngineCoreRequestType::Utility, &request) + .await + { + return Err(error); } + rx.await + .map_err(|_| Error::UtilityCallClosed { + method: method.to_string(), + call_id, + })?? + .into_typed_result(method) + } - // Phase 3: wait for all engines to respond and preserve the per-engine - // result list. - let futures = pending_calls.into_iter().map(|(call_id, rx)| async move { - rx.await - .map_err(|_| Error::UtilityCallClosed { - method: method.to_string(), - call_id, - })?? - .into_typed_result(method) - }); - try_join_all(futures).await + /// Call a typed utility method on all connected engines, returning one + /// decoded result per connected engine if all calls succeed or an error + /// if any call fails. + /// + /// Callers should pass utility arguments using Rust tuple semantics so the + /// encoded payload matches Python's `(client_index, call_id, + /// method_name, args)` contract: `()`, `(arg,)`, `(arg1, arg2)`, etc. + pub async fn call_utility(&self, method: &str, args: A) -> Result> + where + T: serde::de::DeserializeOwned, + A: serde::Serialize + std::fmt::Debug, + { + self.call_utility_per_engine(method, args).await?.into_iter().collect() } /// Call a utility method on all connected engines and return the shared diff --git a/rust/src/engine-core-client/src/client/imp.rs b/rust/src/engine-core-client/src/client/imp.rs index 9c539bf180e9..aaa969a73eba 100644 --- a/rust/src/engine-core-client/src/client/imp.rs +++ b/rust/src/engine-core-client/src/client/imp.rs @@ -118,6 +118,11 @@ impl ClientInner { self.utility_reg.lock().unregister_many(call_ids); } + #[cfg(test)] + pub fn pending_utility_call_count(&self) -> usize { + self.utility_reg.lock().len() + } + /// Undo a request registration when `add_request()` fails. pub fn rollback_request(&self, request_id: &str) { let _ = self.request_reg.lock().remove(request_id); @@ -315,14 +320,23 @@ impl ClientInner { /// recorded for this client. Later failures do not overwrite the first /// one so `/health` and post-close callers observe a stable cause. fn record_health_error(&self, error: Arc) -> Arc { - if let Some(existing) = self.health_error.load_full() { - return existing; - } - self.health_error - .rcu(|current| current.clone().unwrap_or_else(|| error.clone())); - self.health_error - .load_full() - .expect("health error must be recorded before registries close") + let persistent_error = if let Some(existing) = self.health_error.load_full() { + existing + } else { + self.health_error + .rcu(|current| current.clone().unwrap_or_else(|| error.clone())); + self.health_error + .load_full() + .expect("health error must be recorded before registries close") + }; + self.health_tx.send_if_modified(|healthy| { + if !*healthy { + return false; + } + *healthy = false; + true + }); + persistent_error } /// Publish the sticky healthy-to-unhealthy transition. diff --git a/rust/src/engine-core-client/src/client/state.rs b/rust/src/engine-core-client/src/client/state.rs index 28e701ba7320..a42e2f5a57ef 100644 --- a/rust/src/engine-core-client/src/client/state.rs +++ b/rust/src/engine-core-client/src/client/state.rs @@ -437,6 +437,11 @@ impl UtilityRegistry { self.utility_calls.contains_key(&call_id) } + #[cfg(test)] + pub fn len(&self) -> usize { + self.utility_calls.len() + } + pub fn is_closed(&self) -> bool { self.closed } diff --git a/rust/src/engine-core-client/src/error.rs b/rust/src/engine-core-client/src/error.rs index 94809a67860f..20a1b1238eae 100644 --- a/rust/src/engine-core-client/src/error.rs +++ b/rust/src/engine-core-client/src/error.rs @@ -90,6 +90,8 @@ pub enum Error { }, #[error("utility call `{method}` closed unexpectedly (call_id={call_id})")] UtilityCallClosed { method: String, call_id: u64 }, + #[error("utility call `{method}` timed out after {timeout:?}")] + UtilityCallTimeout { method: String, timeout: Duration }, #[error("utility call `{method}` returned inconsistent results across engines: {values}")] InconsistentUtilityResults { method: String, values: String }, diff --git a/rust/src/engine-core-client/src/mock_engine.rs b/rust/src/engine-core-client/src/mock_engine.rs index ee066635ff0c..c5184998cca1 100644 --- a/rust/src/engine-core-client/src/mock_engine.rs +++ b/rust/src/engine-core-client/src/mock_engine.rs @@ -66,7 +66,13 @@ pub fn default_ready_response() -> EngineCoreReadyResponse { instance_id: "test-instance".to_string(), kv_cache_size_tokens: None, kv_cache_max_concurrency: None, - kv_events_config: None, + kv_role: None, + kv_events_publisher: None, + kv_events_endpoint: None, + kv_events_topic: None, + kv_event_block_size: DEFAULT_MOCK_BLOCK_SIZE, + supports_lora: false, + max_loras: 0, } } diff --git a/rust/src/engine-core-client/src/protocol/handshake.rs b/rust/src/engine-core-client/src/protocol/handshake.rs index a4e1cebde67d..acc2503efd2a 100644 --- a/rust/src/engine-core-client/src/protocol/handshake.rs +++ b/rust/src/engine-core-client/src/protocol/handshake.rs @@ -24,19 +24,6 @@ pub struct ReadyMessage { pub parallel_config_hash: Option, } -/// KV-event publisher configuration reported by EngineCore. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct KvEventsConfig { - pub enable_kv_cache_events: bool, - pub publisher: String, - pub endpoint: String, - pub replay_endpoint: Option, - pub buffer_steps: u32, - pub hwm: u32, - pub max_queue_size: u32, - pub topic: String, -} - /// Post-initialization configuration sent from each engine on the input socket /// registration message, after the handshake completes. /// @@ -84,9 +71,24 @@ pub struct EngineCoreReadyResponse { pub kv_cache_size_tokens: Option, /// Maximum achievable request concurrency given the KV cache, if reported. pub kv_cache_max_concurrency: Option, - /// KV-event publisher configuration, if configured. + /// KV transfer role (`kv_producer` / `kv_consumer` / `kv_both`), if any. + #[serde(default)] + pub kv_role: Option, + /// KV-event publisher backend (`null` / `zmq`), if configured. + #[serde(default)] + pub kv_events_publisher: Option, + /// ZMQ endpoint the engine publishes KV events on, if configured. + #[serde(default)] + pub kv_events_endpoint: Option, + /// Topic the KV-event publisher tags events with, if configured. #[serde(default)] - pub kv_events_config: Option, + pub kv_events_topic: Option, + /// Main-attention block size used by published KV events. + pub kv_event_block_size: u64, + /// Whether the engine was started with LoRA support enabled. + pub supports_lora: bool, + /// Maximum number of LoRA adapters the engine may keep active. + pub max_loras: u32, } /// Frontend-owned ZMQ addresses that are sent to the engine during startup diff --git a/rust/src/engine-core-client/src/protocol/output.rs b/rust/src/engine-core-client/src/protocol/output.rs index 78b9ff2daf99..c48fc215ab32 100644 --- a/rust/src/engine-core-client/src/protocol/output.rs +++ b/rust/src/engine-core-client/src/protocol/output.rs @@ -13,6 +13,7 @@ use super::utility::UtilityOutput; use crate::error::{Error, Result, ext_value_decode}; use crate::protocol::logprobs::MaybeWireLogprobs; use crate::protocol::stats::{PrefillStats, SchedulerStats}; +use crate::protocol::tensor::{ShapeExt, WireArrayData, WireNdArray}; use crate::protocol::{OpaqueValue, decode_msgpack}; /// The stop reason associated with a finished output. @@ -108,7 +109,7 @@ pub struct EngineCoreOutput { #[serde(default)] pub prefill_stats: Option, #[serde(default)] - pub routed_experts: Option, + pub routed_experts: Option, /// Number of NaNs seen in logits. Values above zero indicate corruption. #[serde(default)] pub num_nans_in_logits: u32, @@ -132,10 +133,111 @@ impl EngineCoreOutput { self.new_prompt_logprobs_tensors = (self.new_prompt_logprobs_tensors.take()) .map(|value| value.resolve(frames, "new_prompt_logprobs_tensors")) .transpose()?; + self.routed_experts = (self.routed_experts.take()) + .map(|value| resolve_routed_experts(value, frames)) + .transpose()?; Ok(()) } } +fn resolve_routed_experts(value: WireNdArray, frames: &[Frame]) -> Result +where + Frame: AsRef<[u8]>, +{ + let WireNdArray { dtype, shape, data } = value; + if shape.len() != 3 { + return Err(Error::Decode { + target_type: "EngineCoreOutput.routed_experts", + message: format!("expected rank 3, got shape {shape:?}"), + }); + } + let element_width = match dtype.as_str() { + "uint8" | "|u1" => 1, + "uint16" | "u2" | "=u2" => 2, + _ => { + return Err(Error::Decode { + target_type: "EngineCoreOutput.routed_experts", + message: format!("expected uint8 or uint16 dtype, got {dtype:?}"), + }); + } + }; + let bytes = match data { + WireArrayData::RawView(bytes) => bytes, + WireArrayData::AuxIndex(index) => frames + .get(index) + .ok_or_else(|| Error::Decode { + target_type: "EngineCoreOutput.routed_experts", + message: format!( + "aux frame index {index} out of range for {} frames", + frames.len() + ), + })? + .as_ref() + .to_vec(), + }; + let expected = shape + .checked_numel() + .and_then(|elements| elements.checked_mul(element_width)) + .ok_or_else(|| Error::Decode { + target_type: "EngineCoreOutput.routed_experts", + message: "shape byte length overflows usize".to_string(), + })?; + if bytes.len() != expected { + return Err(Error::Decode { + target_type: "EngineCoreOutput.routed_experts", + message: format!( + "byte length mismatch: expected {expected}, got {}", + bytes.len() + ), + }); + } + Ok(WireNdArray { + dtype, + shape, + data: WireArrayData::RawView(bytes), + }) +} + +/// Concatenate resolved routed-expert chunks along their token axis. +pub fn concatenate_routed_experts( + chunks: impl IntoIterator, +) -> Result> { + let mut chunks = chunks.into_iter(); + let Some(mut combined) = chunks.next() else { + return Ok(None); + }; + let WireArrayData::RawView(combined_data) = &mut combined.data else { + return Err(Error::Decode { + target_type: "EngineCoreOutput.routed_experts", + message: "unresolved auxiliary payload".to_string(), + }); + }; + for chunk in chunks { + if chunk.dtype != combined.dtype + || chunk.shape.len() != combined.shape.len() + || chunk.shape[1..] != combined.shape[1..] + { + return Err(Error::Decode { + target_type: "EngineCoreOutput.routed_experts", + message: "incompatible routed-expert chunk dtype or shape".to_string(), + }); + } + combined.shape[0] = + combined.shape[0].checked_add(chunk.shape[0]).ok_or_else(|| Error::Decode { + target_type: "EngineCoreOutput.routed_experts", + message: "concatenated token dimension overflows usize".to_string(), + })?; + let WireArrayData::RawView(data) = chunk.data else { + return Err(Error::Decode { + target_type: "EngineCoreOutput.routed_experts", + message: "unresolved auxiliary payload".to_string(), + }); + }; + combined_data.extend(data); + } + Ok(Some(combined)) +} + /// Raw Python/msgpack engine-core output envelope. /// /// Original Python definition: @@ -403,6 +505,45 @@ mod tests { ); } + #[test] + fn routed_experts_resolve_from_aux_frame() { + let outputs = WireEngineCoreOutputs { + outputs: vec![EngineCoreOutput { + request_id: "req-routed".to_string(), + new_token_ids: vec![42], + routed_experts: Some(WireNdArray { + dtype: "|u1".to_string(), + shape: vec![1, 2, 2], + data: WireArrayData::AuxIndex(1), + }), + ..Default::default() + }], + ..Default::default() + }; + let primary = encode_msgpack(&outputs).unwrap(); + let frames = vec![primary, vec![1, 2, 3, 4]]; + + let decoded = decode_engine_core_outputs(&frames).unwrap(); + let routed = + decoded.as_request_batch().unwrap().outputs[0].routed_experts.as_ref().unwrap(); + assert_eq!(routed.dtype, "|u1"); + assert_eq!(routed.shape, vec![1, 2, 2]); + assert_eq!(routed.data.as_raw_view().unwrap(), &[1, 2, 3, 4]); + } + + #[test] + fn routed_experts_chunks_concatenate_on_token_axis() { + let combined = concatenate_routed_experts([ + WireNdArray::from_raw("|u1", vec![1, 2, 1], vec![1, 2]), + WireNdArray::from_raw("|u1", vec![2, 2, 1], vec![3, 4, 5, 6]), + ]) + .unwrap() + .unwrap(); + + assert_eq!(combined.shape, vec![3, 2, 1]); + assert_eq!(combined.data.as_raw_view().unwrap(), &[1, 2, 3, 4, 5, 6]); + } + #[test] fn engine_core_outputs_classify_request_batch() { let outputs = WireEngineCoreOutputs { diff --git a/rust/src/engine-core-client/src/protocol/sampling.rs b/rust/src/engine-core-client/src/protocol/sampling.rs index d5c6a852081f..222a33d8ae43 100644 --- a/rust/src/engine-core-client/src/protocol/sampling.rs +++ b/rust/src/engine-core-client/src/protocol/sampling.rs @@ -86,6 +86,8 @@ pub struct EngineCoreSamplingParams { /// reaching this DTO, so only non-negative values are sent. Enforced /// engine-side (and only when a reasoning parser is configured). pub thinking_token_budget: Option, + /// Number of prompt-token rows to omit from returned routed-expert data. + pub routed_experts_prompt_start: u32, /// Number of log probabilities to return per generated token. /// /// `None` disables sample logprobs. `-1` requests the full vocabulary. @@ -158,6 +160,7 @@ impl EngineCoreSamplingParams { max_tokens: 65536, min_tokens: 0, thinking_token_budget: None, + routed_experts_prompt_start: 0, logprobs: None, prompt_logprobs: None, min_p: 0.0, diff --git a/rust/src/engine-core-client/src/test_utils.rs b/rust/src/engine-core-client/src/test_utils.rs index 13d473f59790..b6e39420ccb5 100644 --- a/rust/src/engine-core-client/src/test_utils.rs +++ b/rust/src/engine-core-client/src/test_utils.rs @@ -85,18 +85,49 @@ pub async fn setup_mock_engine_sockets( .expect("connect mock engine") } +/// Complete the handshake with an explicit engine-ready response. +pub async fn setup_mock_engine_sockets_with_ready( + engine_handshake: String, + engine_id: impl Into, + ready_response: EngineCoreReadyResponse, +) -> MockEngineSockets { + connect_to_frontend( + engine_handshake, + engine_id, + test_mock_engine_config_with_ready(ready_response), + ) + .await + .expect("connect mock engine") +} + /// Connect one mock engine directly to already-bootstrapped frontend /// input/output sockets. pub async fn setup_bootstrapped_mock_engine( input_address: String, output_address: String, engine_id: impl Into, +) -> (DealerSocket, PushSocket) { + setup_bootstrapped_mock_engine_with_ready( + input_address, + output_address, + engine_id, + default_ready_response(), + ) + .await +} + +/// Connect a bootstrapped mock engine with an explicit ready response. +pub async fn setup_bootstrapped_mock_engine_with_ready( + input_address: String, + output_address: String, + engine_id: impl Into, + ready_response: EngineCoreReadyResponse, ) -> (DealerSocket, PushSocket) { connect_to_bootstrapped_frontend( input_address, output_address, engine_id, - test_mock_engine_config(), + test_mock_engine_config_with_ready(ready_response), ) .await .expect("connect bootstrapped mock engine") @@ -142,6 +173,24 @@ pub fn spawn_mock_engine_task( engine_id: impl Into, run: F, ) -> (oneshot::Sender<()>, tokio::task::JoinHandle<()>) +where + F: for<'a> FnOnce( + &'a mut DealerSocket, + &'a mut PushSocket, + ) -> Pin + Send + 'a>> + + Send + + 'static, +{ + spawn_mock_engine_task_with_config(engine_handshake, engine_id, test_mock_engine_config(), run) +} + +/// Variant of [`spawn_mock_engine_task`] with an explicit startup response. +pub fn spawn_mock_engine_task_with_config( + engine_handshake: String, + engine_id: impl Into, + config: MockEngineConfig, + run: F, +) -> (oneshot::Sender<()>, tokio::task::JoinHandle<()>) where F: for<'a> FnOnce( &'a mut DealerSocket, @@ -153,7 +202,14 @@ where let (shutdown_tx, shutdown_rx) = oneshot::channel(); let engine_id = engine_id.into(); let engine_task = tokio::spawn(async move { - let (mut dealer, mut push) = setup_mock_engine(engine_handshake, engine_id).await; + let MockEngineSockets { data_sockets, .. } = + connect_to_frontend(engine_handshake, engine_id, config) + .await + .expect("connect mock engine"); + let MockEngineDataSockets { + mut dealer, + mut push, + } = data_sockets.into_iter().next().expect("mock engine data socket"); run(&mut dealer, &mut push).await; let _ = shutdown_rx.await; }); diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 11d08db646ba..c66235a63b57 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -32,11 +32,12 @@ use crate::protocol::output::{ use crate::protocol::request::{EngineCoreRequest, EngineCoreRequestType}; use crate::protocol::sampling::EngineCoreSamplingParams; use crate::protocol::stats::SchedulerStats; -use crate::protocol::tensor::{WireArrayData, WireTensor}; +use crate::protocol::tensor::WireTensor; use crate::protocol::utility::{UtilityOutput, UtilityResultEnvelope}; use crate::test_utils::{ - IpcNamespace, setup_bootstrapped_mock_engine, setup_mock_engine_sockets, - setup_mock_engine_with_init, spawn_mock_engine_task, + IpcNamespace, setup_bootstrapped_mock_engine, setup_bootstrapped_mock_engine_with_ready, + setup_mock_engine_sockets, setup_mock_engine_sockets_with_ready, spawn_mock_engine_task, + spawn_mock_engine_task_with_ready, }; use crate::{ CoordinatorMode, ENGINE_CORE_DEAD_SENTINEL, EngineCoreClient, EngineCoreClientConfig, EngineId, @@ -45,6 +46,8 @@ use crate::{ static TRACING: Once = Once::new(); +mod utility; + fn expect_sample_logprobs(actual: &MaybeWireLogprobs) { expect_test::expect![[r#" Logprobs { @@ -162,6 +165,18 @@ fn sample_request_with_id(request_id: &str) -> EngineCoreRequest { ..EngineCoreSamplingParams::for_test() }), arrival_time: 42.5, + lora_request: Some(crate::protocol::lora::LoraRequest { + lora_name: "adapter-a".to_string(), + lora_int_id: 17, + lora_path: "/models/adapter-a".to_string(), + base_model_name: Some("Qwen/Qwen3-0.6B".to_string()), + tensorizer_config_dict: Some(rmpv::Value::Map(vec![( + rmpv::Value::from("format"), + rmpv::Value::from("safetensors"), + )])), + load_inplace: true, + is_3d_lora_weight: true, + }), ..EngineCoreRequest::default() } } @@ -320,6 +335,40 @@ fn bootstrapped_test_config( } } +fn two_rank_ready(rank: u32) -> EngineCoreReadyResponse { + EngineCoreReadyResponse { + data_parallel_size: 2, + data_parallel_rank: rank, + ..crate::mock_engine::default_ready_response() + } +} + +async fn setup_two_rank_mock_engine_sockets( + engine_handshake: String, + engine_id: impl Into, + rank: u32, +) -> crate::mock_engine::MockEngineSockets { + setup_mock_engine_sockets_with_ready(engine_handshake, engine_id, two_rank_ready(rank)).await +} + +fn spawn_two_rank_mock_engine_task( + engine_handshake: String, + engine_id: impl Into, + rank: u32, + run: F, +) -> (oneshot::Sender<()>, tokio::task::JoinHandle<()>) +where + F: for<'a> FnOnce( + &'a mut DealerSocket, + &'a mut PushSocket, + ) + -> std::pin::Pin + Send + 'a>> + + Send + + 'static, +{ + spawn_mock_engine_task_with_ready(engine_handshake, engine_id, two_rank_ready(rank), run) +} + fn bootstrapped_test_config_with_start_index( input_address: String, output_address: String, @@ -376,6 +425,7 @@ async fn send_external_coordinator_publish( fn spawn_mock_engine_task_with_init( engine_handshake: String, engine_id: impl Into, + rank: u32, run: F, ) -> ( oneshot::Receiver, @@ -395,10 +445,11 @@ where let (init_tx, init_rx) = oneshot::channel(); let engine_id = engine_id.into(); let engine_task = tokio::spawn(async move { - let (init, mut dealer, mut push) = - setup_mock_engine_with_init(engine_handshake, engine_id).await; - let _ = init_tx.send(init); - run(&mut dealer, &mut push).await; + let mut sockets = + setup_two_rank_mock_engine_sockets(engine_handshake, engine_id, rank).await; + let _ = init_tx.send(sockets.init); + let data_socket = sockets.data_sockets.first_mut().expect("mock engine data socket"); + run(&mut data_socket.dealer, &mut data_socket.push).await; let _ = shutdown_rx.await; }); (init_rx, shutdown_tx, engine_task) @@ -535,12 +586,131 @@ async fn coordinator_handshake_includes_engine_control_addresses() { assert!(init.addresses.coordinator_input.is_some()); assert!(init.addresses.coordinator_output.is_some()); assert!(init.addresses.frontend_stats_publish_address.is_none()); + assert!(init.parallel_config.is_empty()); let _ = shutdown_tx.send(()); engine_task.await.unwrap(); client.shutdown().await.unwrap(); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn non_coordinator_handshake_omits_data_parallel_bootstrap() { + init_tracing(); + let ipc = IpcNamespace::new().unwrap(); + let handshake_address = ipc.handshake_endpoint(); + + let (init0_rx, shutdown0_tx, engine0_task) = + spawn_mock_engine_task_with_init(handshake_address.clone(), &[0x00, 0x00], 0, |_, _| { + Box::pin(async {}) + }); + let (init1_rx, shutdown1_tx, engine1_task) = + spawn_mock_engine_task_with_init(handshake_address, &[0x01, 0x00], 1, |_, _| { + Box::pin(async {}) + }); + + let client = connect_client_with_ipc( + handshake_test_config( + ipc.handshake_endpoint(), + 2, + "test-model", + Duration::from_secs(2), + 0, + None, + ), + &ipc, + ) + .await; + + assert!(init0_rx.await.unwrap().parallel_config.is_empty()); + assert!(init1_rx.await.unwrap().parallel_config.is_empty()); + + let _ = shutdown0_tx.send(()); + let _ = shutdown1_tx.send(()); + engine0_task.await.unwrap(); + engine1_task.await.unwrap(); + client.shutdown().await.unwrap(); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn coordinator_handshake_shares_data_parallel_bootstrap() { + init_tracing(); + let ipc = IpcNamespace::new().unwrap(); + let handshake_address = ipc.handshake_endpoint(); + let advertised_host = "127.0.0.2"; + + let (init0_rx, shutdown0_tx, engine0_task) = + spawn_mock_engine_task_with_init(handshake_address.clone(), &[0x00, 0x00], 0, |_, _| { + Box::pin(async {}) + }); + let (init1_rx, shutdown1_tx, engine1_task) = + spawn_mock_engine_task_with_init(handshake_address, &[0x01, 0x00], 1, |_, _| { + Box::pin(async {}) + }); + + let mut config = handshake_test_config( + ipc.handshake_endpoint(), + 2, + "test-model", + Duration::from_secs(2), + 0, + Some(CoordinatorMode::InProc), + ); + let TransportMode::HandshakeOwner { + advertised_host: config_host, + .. + } = &mut config.transport_mode + else { + unreachable!("handshake_test_config returns handshake-owned transport") + }; + *config_host = advertised_host.to_string(); + + let client = connect_client_with_ipc(config, &ipc).await; + let init0 = init0_rx.await.unwrap(); + let init1 = init1_rx.await.unwrap(); + + assert_eq!(init0.parallel_config, init1.parallel_config); + assert_eq!( + init0.parallel_config.keys().map(String::as_str).collect::>(), + vec![ + "_data_parallel_master_port_list", + "data_parallel_master_ip", + "data_parallel_master_port", + "data_parallel_size", + ] + ); + assert_eq!( + init0.parallel_config["data_parallel_master_ip"].as_str(), + Some(advertised_host) + ); + assert_eq!( + init0.parallel_config["data_parallel_size"].as_u64(), + Some(2) + ); + + let master_port = init0.parallel_config["data_parallel_master_port"] + .as_u64() + .expect("data_parallel_master_port must be an integer"); + let remaining_ports = init0.parallel_config["_data_parallel_master_port_list"] + .as_array() + .expect("_data_parallel_master_port_list must be an array"); + let ports = std::iter::once(master_port) + .chain( + remaining_ports + .iter() + .map(|port| port.as_u64().expect("bootstrap port must be an integer")), + ) + .map(|port| u16::try_from(port).expect("bootstrap port must fit in u16")) + .collect::>(); + assert_eq!(ports.len(), 5); + assert!(ports.iter().all(|port| *port != 0)); + + let _ = shutdown0_tx.send(()); + let _ = shutdown1_tx.send(()); + engine0_task.await.unwrap(); + engine1_task.await.unwrap(); + client.shutdown().await.unwrap(); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() { init_tracing(); @@ -551,7 +721,8 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() { let engine0_task = tokio::spawn({ let handshake_address = handshake_address.clone(); async move { - let mut engine = setup_mock_engine_sockets(handshake_address, &[0x00, 0x00]).await; + let mut engine = + setup_two_rank_mock_engine_sockets(handshake_address, &[0x00, 0x00], 0).await; let mut coordinator = engine.coordinator.take().expect("coordinator sockets should be present"); let data_socket = engine.data_sockets.first_mut().expect("data socket"); @@ -632,7 +803,8 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() { let engine1_task = tokio::spawn({ let handshake_address = handshake_address.clone(); async move { - let mut engine = setup_mock_engine_sockets(handshake_address, &[0x01, 0x00]).await; + let mut engine = + setup_two_rank_mock_engine_sockets(handshake_address, &[0x01, 0x00], 1).await; let mut coordinator = engine.coordinator.take().expect("coordinator sockets should be present"); let data_socket = engine.data_sockets.first_mut().expect("data socket"); @@ -752,7 +924,8 @@ async fn coordinator_rebroadcasts_engine_start_wave_control() { let engine0_task = tokio::spawn({ let handshake_address = handshake_address.clone(); async move { - let mut engine = setup_mock_engine_sockets(handshake_address, &[0x00, 0x00]).await; + let mut engine = + setup_two_rank_mock_engine_sockets(handshake_address, &[0x00, 0x00], 0).await; let mut coordinator = engine.coordinator.take().expect("coordinator sockets should be present"); @@ -767,7 +940,8 @@ async fn coordinator_rebroadcasts_engine_start_wave_control() { let engine1_task = tokio::spawn({ let handshake_address = handshake_address.clone(); async move { - let mut engine = setup_mock_engine_sockets(handshake_address, &[0x01, 0x00]).await; + let mut engine = + setup_two_rank_mock_engine_sockets(handshake_address, &[0x01, 0x00], 1).await; let mut coordinator = engine.coordinator.take().expect("coordinator sockets should be present"); @@ -1728,86 +1902,6 @@ async fn client_decodes_multipart_logprob_outputs() { client.shutdown().await.unwrap(); } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn client_sends_large_multimodal_tensor_as_aux_frame() { - init_tracing(); - let ipc = IpcNamespace::new().unwrap(); - let handshake_address = ipc.handshake_endpoint(); - let engine_id = b"engine-multimodal-aux".to_vec(); - let tensor_data = (0..64).map(|value| value as f32).collect::>(); - let expected_bytes = - tensor_data.iter().flat_map(|value| value.to_ne_bytes()).collect::>(); - let mut request = sample_multimodal_request(); - request.mm_features.as_mut().unwrap()[0] - .data - .as_mut() - .unwrap() - .get_mut("pixel_values") - .unwrap() - .data = Some(MmKwargValue::Tensor( - WireTensor::from_f32(vec![64], tensor_data).unwrap(), - )); - - let (shutdown_tx, engine_task) = spawn_mock_engine_task( - handshake_address.clone(), - engine_id.clone(), - move |dealer, push| { - Box::pin(async move { - let add = recv_engine_message(dealer).await; - assert_eq!(add.len(), 3); - assert_eq!(add[0].as_ref(), &[0x00]); - let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap(); - let MmKwargValue::Tensor(tensor) = - request.mm_features.as_ref().unwrap()[0].data.as_ref().unwrap()["pixel_values"] - .data - .as_ref() - .unwrap() - else { - panic!("expected tensor"); - }; - assert_eq!(tensor.data, WireArrayData::AuxIndex(1)); - assert_eq!(add[2].as_ref(), expected_bytes); - - send_outputs( - push, - RequestBatchOutputs { - outputs: vec![request_output( - "req-mm", - vec![], - Some(EngineCoreFinishReason::Length), - )], - finished_requests: Some(BTreeSet::from(["req-mm".to_string()])), - ..Default::default() - } - .into(), - ) - .await; - }) - }, - ); - - let client = connect_client_with_ipc( - handshake_test_config( - handshake_address, - 1, - "test-model", - Duration::from_secs(2), - 0, - None, - ), - &ipc, - ) - .await; - - let outputs = client.call(request).await.unwrap().collect::>().await; - assert_eq!(outputs.len(), 1); - assert!(outputs[0].is_ok()); - - let _ = shutdown_tx.send(()); - engine_task.await.unwrap(); - client.shutdown().await.unwrap(); -} - #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn multi_engine_client_shares_transport_and_routes_by_inflight_count() { init_tracing(); @@ -1822,6 +1916,7 @@ async fn multi_engine_client_shares_transport_and_routes_by_inflight_count() { let (init_rx_0, shutdown_tx_0, engine_task_0) = spawn_mock_engine_task_with_init( handshake_address.clone(), b"engine-0".to_vec(), + 0, |dealer, push| { Box::pin(async move { let add_1 = recv_engine_message(dealer).await; @@ -1871,6 +1966,7 @@ async fn multi_engine_client_shares_transport_and_routes_by_inflight_count() { let (init_rx_1, shutdown_tx_1, engine_task_1) = spawn_mock_engine_task_with_init( handshake_address.clone(), b"engine-1".to_vec(), + 1, |dealer, push| { Box::pin(async move { let add_2 = recv_engine_message(dealer).await; @@ -1992,9 +2088,10 @@ async fn multi_engine_abort_is_grouped_and_utility_fans_out_to_all_engines() { let ipc = IpcNamespace::new().unwrap(); let handshake_address = ipc.handshake_endpoint(); - let (shutdown_tx_0, engine_task_0) = spawn_mock_engine_task( + let (shutdown_tx_0, engine_task_0) = spawn_two_rank_mock_engine_task( handshake_address.clone(), EngineId::from_engine_index(0).into_frame().to_vec(), + 0, |dealer, push| { Box::pin(async move { let utility = recv_engine_message(dealer).await; @@ -2048,9 +2145,10 @@ async fn multi_engine_abort_is_grouped_and_utility_fans_out_to_all_engines() { }, ); tokio::time::sleep(Duration::from_millis(50)).await; - let (shutdown_tx_1, engine_task_1) = spawn_mock_engine_task( + let (shutdown_tx_1, engine_task_1) = spawn_two_rank_mock_engine_task( handshake_address.clone(), EngineId::from_engine_index(1).into_frame().to_vec(), + 1, |dealer, push| { Box::pin(async move { let utility = recv_engine_message(dealer).await; @@ -2161,9 +2259,10 @@ async fn collective_rpc_flattens_results_from_all_engines() { let ipc = IpcNamespace::new().unwrap(); let handshake_address = ipc.handshake_endpoint(); - let (shutdown_tx_0, engine_task_0) = spawn_mock_engine_task( + let (shutdown_tx_0, engine_task_0) = spawn_two_rank_mock_engine_task( handshake_address.clone(), b"engine-0".to_vec(), + 0, |dealer, push| { Box::pin(async move { let utility = recv_engine_message(dealer).await; @@ -2194,9 +2293,10 @@ async fn collective_rpc_flattens_results_from_all_engines() { }, ); tokio::time::sleep(Duration::from_millis(50)).await; - let (shutdown_tx_1, engine_task_1) = spawn_mock_engine_task( + let (shutdown_tx_1, engine_task_1) = spawn_two_rank_mock_engine_task( handshake_address.clone(), b"engine-1".to_vec(), + 1, |dealer, push| { Box::pin(async move { let utility = recv_engine_message(dealer).await; @@ -2269,6 +2369,7 @@ async fn collective_rpc_flattens_results_from_all_engines() { fn spawn_mock_utility_engine( handshake_address: String, engine_id: Vec, + rank: u32, expected_method: &'static str, expected_args: Value, result: bool, @@ -2276,7 +2377,7 @@ fn spawn_mock_utility_engine( tokio::sync::oneshot::Sender<()>, tokio::task::JoinHandle<()>, ) { - spawn_mock_engine_task(handshake_address, engine_id, move |dealer, push| { + spawn_two_rank_mock_engine_task(handshake_address, engine_id, rank, move |dealer, push| { Box::pin(async move { let utility = recv_engine_message(dealer).await; assert_eq!(utility[0].as_ref(), &[0x03]); @@ -2316,6 +2417,7 @@ async fn is_sleeping_returns_error_when_engines_disagree() { let (shutdown_tx_0, engine_task_0) = spawn_mock_utility_engine( handshake_address.clone(), b"engine-0".to_vec(), + 0, "is_sleeping", Value::Array(vec![]), true, @@ -2323,6 +2425,7 @@ async fn is_sleeping_returns_error_when_engines_disagree() { let (shutdown_tx_1, engine_task_1) = spawn_mock_utility_engine( handshake_address.clone(), b"engine-1".to_vec(), + 1, "is_sleeping", Value::Array(vec![]), false, @@ -2366,6 +2469,7 @@ async fn is_sleeping_returns_value_when_all_engines_agree() { let (shutdown_tx_0, engine_task_0) = spawn_mock_utility_engine( handshake_address.clone(), b"engine-0".to_vec(), + 0, "is_sleeping", Value::Array(vec![]), true, @@ -2373,6 +2477,7 @@ async fn is_sleeping_returns_value_when_all_engines_agree() { let (shutdown_tx_1, engine_task_1) = spawn_mock_utility_engine( handshake_address.clone(), b"engine-1".to_vec(), + 1, "is_sleeping", Value::Array(vec![]), true, @@ -2409,6 +2514,7 @@ async fn reset_prefix_cache_returns_true_when_all_engines_succeed() { let (shutdown_tx_0, engine_task_0) = spawn_mock_utility_engine( handshake_address.clone(), b"engine-0".to_vec(), + 0, "reset_prefix_cache", Value::Array(vec![Value::from(false), Value::from(false)]), true, @@ -2416,6 +2522,7 @@ async fn reset_prefix_cache_returns_true_when_all_engines_succeed() { let (shutdown_tx_1, engine_task_1) = spawn_mock_utility_engine( handshake_address.clone(), b"engine-1".to_vec(), + 1, "reset_prefix_cache", Value::Array(vec![Value::from(false), Value::from(false)]), true, @@ -2452,6 +2559,7 @@ async fn reset_prefix_cache_returns_false_when_any_engine_fails() { let (shutdown_tx_0, engine_task_0) = spawn_mock_utility_engine( handshake_address.clone(), b"engine-0".to_vec(), + 0, "reset_prefix_cache", Value::Array(vec![Value::from(false), Value::from(false)]), true, @@ -2459,6 +2567,7 @@ async fn reset_prefix_cache_returns_false_when_any_engine_fails() { let (shutdown_tx_1, engine_task_1) = spawn_mock_utility_engine( handshake_address.clone(), b"engine-1".to_vec(), + 1, "reset_prefix_cache", Value::Array(vec![Value::from(false), Value::from(false)]), false, @@ -2540,6 +2649,7 @@ fn python_msgpack_fixtures_match_rust_encoding() { max_tokens: 16, min_tokens: 0, thinking_token_budget: None, + routed_experts_prompt_start: 0, logprobs: None, prompt_logprobs: None, min_p: 0.0, @@ -2677,21 +2787,12 @@ fn python_msgpack_fixtures_match_rust_encoding() { rust_ready_keys, python_ready_keys, "EngineCoreReadyResponse drifted from the Python dataclass", ); - - let ready_response: EngineCoreReadyResponse = + let python_ready: crate::protocol::handshake::EngineCoreReadyResponse = rmp_serde::from_slice(&hex::decode(ready_response_hex).unwrap()).unwrap(); - let kv_events_config = ready_response.kv_events_config.expect("KV events config should decode"); - assert!(kv_events_config.enable_kv_cache_events); - assert_eq!(kv_events_config.publisher, "zmq"); - assert_eq!(kv_events_config.endpoint, "tcp://127.0.0.1:5557"); - assert_eq!( - kv_events_config.replay_endpoint.as_deref(), - Some("tcp://127.0.0.1:5558") - ); - assert_eq!(kv_events_config.topic, "kv"); - assert_eq!(kv_events_config.buffer_steps, 10_000); - assert_eq!(kv_events_config.hwm, 100_000); - assert_eq!(kv_events_config.max_queue_size, 100_000); + assert_eq!(python_ready.kv_event_block_size, 256); + assert_eq!(python_ready.kv_events_publisher.as_deref(), Some("zmq")); + assert!(python_ready.supports_lora); + assert_eq!(python_ready.max_loras, 8); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] @@ -2754,14 +2855,20 @@ async fn bootstrapped_connects_with_contiguous_engine_ids() { } }); - let (_dealer0, _push0) = setup_bootstrapped_mock_engine( + let (_dealer0, _push0) = setup_bootstrapped_mock_engine_with_ready( input_address.clone(), output_address.clone(), &[0x00, 0x00], + two_rank_ready(0), + ) + .await; + let (_dealer1, _push1) = setup_bootstrapped_mock_engine_with_ready( + input_address, + output_address, + &[0x01, 0x00], + two_rank_ready(1), ) .await; - let (_dealer1, _push1) = - setup_bootstrapped_mock_engine(input_address, output_address, &[0x01, 0x00]).await; let client = client_task.await.unwrap(); assert_eq!(client.engine_count(), 2); diff --git a/rust/src/engine-core-client/src/tests/client/utility.rs b/rust/src/engine-core-client/src/tests/client/utility.rs new file mode 100644 index 000000000000..76466243e910 --- /dev/null +++ b/rust/src/engine-core-client/src/tests/client/utility.rs @@ -0,0 +1,212 @@ +use super::*; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn cancelled_utility_call_unregisters_waiter() { + init_tracing(); + let ipc = IpcNamespace::new().unwrap(); + let handshake_address = ipc.handshake_endpoint(); + let (received_tx, received_rx) = oneshot::channel(); + let (_shutdown, engine_task) = spawn_mock_engine_task( + handshake_address.clone(), + vec![0x00, 0x00], + |dealer, _push| { + Box::pin(async move { + let _utility = recv_engine_message(dealer).await; + let _ = received_tx.send(()); + std::future::pending::<()>().await; + }) + }, + ); + let client = std::sync::Arc::new( + connect_client_with_ipc( + handshake_test_config( + handshake_address, + 1, + "test-model", + Duration::from_secs(2), + 0, + None, + ), + &ipc, + ) + .await, + ); + let call_client = client.clone(); + let call = tokio::spawn(async move { + call_client.call_utility_per_engine::("add_lora", ()).await + }); + + received_rx.await.unwrap(); + assert_eq!(client.pending_utility_call_count(), 1); + call.abort(); + assert!(call.await.unwrap_err().is_cancelled()); + assert_eq!(client.pending_utility_call_count(), 0); + + engine_task.abort(); + let client = match std::sync::Arc::try_unwrap(client) { + Ok(client) => client, + Err(_) => panic!("utility task retained the client after cancellation"), + }; + client.shutdown().await.unwrap(); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn call_utility_per_engine_preserves_partial_results() { + init_tracing(); + let ipc = IpcNamespace::new().unwrap(); + let handshake_address = ipc.handshake_endpoint(); + + let (shutdown_0, task_0) = spawn_mock_engine_task_with_ready( + handshake_address.clone(), + vec![0x00, 0x00], + two_rank_ready(0), + |dealer, push| { + Box::pin(async move { + let utility = recv_engine_message(dealer).await; + let call_id = decode_value(&utility[1]).as_array().unwrap()[1].as_u64().unwrap(); + send_outputs( + push, + UtilityCallOutput { + engine_index: 0, + timestamp: 0.0, + output: UtilityOutput { + call_id: call_id.into(), + failure_message: None, + result: Some(utility_result_value(true)), + }, + } + .into(), + ) + .await; + }) + }, + ); + let (shutdown_1, task_1) = spawn_mock_engine_task_with_ready( + handshake_address.clone(), + vec![0x01, 0x00], + two_rank_ready(1), + |dealer, push| { + Box::pin(async move { + let utility = recv_engine_message(dealer).await; + let call_id = decode_value(&utility[1]).as_array().unwrap()[1].as_u64().unwrap(); + send_outputs( + push, + UtilityCallOutput { + engine_index: 1, + timestamp: 0.0, + output: UtilityOutput { + call_id: call_id.into(), + failure_message: Some("rank failed".to_string()), + result: None, + }, + } + .into(), + ) + .await; + }) + }, + ); + + let client = connect_client_with_ipc( + handshake_test_config( + handshake_address, + 2, + "test-model", + Duration::from_secs(2), + 0, + None, + ), + &ipc, + ) + .await; + + let results = client.call_utility_per_engine::("add_lora", ()).await.unwrap(); + assert_eq!(results.len(), 2); + assert!(matches!(results[0], Ok(true))); + assert!(matches!( + &results[1], + Err(Error::UtilityCallFailed { message, .. }) if message == "rank failed" + )); + + let _ = shutdown_0.send(()); + let _ = shutdown_1.send(()); + task_0.await.unwrap(); + task_1.await.unwrap(); + client.shutdown().await.unwrap(); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn call_utility_on_engine_targets_selected_nonzero_rank() { + init_tracing(); + let ipc = IpcNamespace::new().unwrap(); + let handshake_address = ipc.handshake_endpoint(); + + let ready = |rank| EngineCoreReadyResponse { + data_parallel_size: 6, + data_parallel_rank: rank, + ..crate::mock_engine::default_ready_response() + }; + let (shutdown_0, task_0) = spawn_mock_engine_task_with_ready( + handshake_address.clone(), + vec![0x04, 0x00], + ready(4), + |_dealer, _push| Box::pin(async {}), + ); + let (shutdown_1, task_1) = spawn_mock_engine_task_with_ready( + handshake_address.clone(), + vec![0x05, 0x00], + ready(5), + |dealer, push| { + Box::pin(async move { + let utility = recv_engine_message(dealer).await; + let payload = decode_value(&utility[1]); + let array = payload.as_array().expect("utility payload"); + assert_eq!(array[2], Value::from("remove_lora")); + assert_eq!(array[3], Value::Array(vec![Value::from(17)])); + let call_id = array[1].as_u64().expect("call_id"); + send_outputs( + push, + UtilityCallOutput { + engine_index: 5, + timestamp: 0.0, + output: UtilityOutput { + call_id: call_id.into(), + failure_message: None, + result: Some(utility_result_value(true)), + }, + } + .into(), + ) + .await; + }) + }, + ); + + let client = connect_client_with_ipc( + handshake_test_config( + handshake_address, + 2, + "test-model", + Duration::from_secs(2), + 0, + None, + ), + &ipc, + ) + .await; + + let removed = timeout( + Duration::from_secs(2), + client.call_utility_on_engine::(5, "remove_lora", (17u64,)), + ) + .await + .expect("targeted utility call timed out") + .unwrap(); + assert!(removed); + + let _ = shutdown_0.send(()); + let _ = shutdown_1.send(()); + task_0.await.unwrap(); + task_1.await.unwrap(); + client.shutdown().await.unwrap(); +} diff --git a/rust/src/engine-core-client/src/tests/python_compat.py b/rust/src/engine-core-client/src/tests/python_compat.py index c68e2b6b471e..64112db7b77f 100755 --- a/rust/src/engine-core-client/src/tests/python_compat.py +++ b/rust/src/engine-core-client/src/tests/python_compat.py @@ -50,6 +50,20 @@ class EngineCoreSamplingParams(msgspec.Struct, dict=True, omit_defaults=True): output_kind: RequestOutputKind = RequestOutputKind.DELTA +class LoRARequest( + msgspec.Struct, + array_like=True, + omit_defaults=True, +): + lora_name: str + lora_int_id: int + lora_path: str = "" + base_model_name: str | None = None + tensorizer_config_dict: dict | None = None + load_inplace: bool = False + is_3d_lora_weight: bool = False + + class EngineCoreRequest( msgspec.Struct, array_like=True, @@ -61,7 +75,7 @@ class EngineCoreRequest( sampling_params: EngineCoreSamplingParams | None pooling_params: object | None arrival_time: float - lora_request: object | None = None + lora_request: LoRARequest | None = None cache_salt: str | None = None data_parallel_rank: int | None = None prompt_embeds: object | None = None @@ -136,6 +150,15 @@ class EngineCoreOutputs( ), pooling_params=None, arrival_time=42.5, + lora_request=LoRARequest( + lora_name="adapter-a", + lora_int_id=17, + lora_path="/models/adapter-a", + base_model_name="Qwen/Qwen3-0.6B", + tensorizer_config_dict={"format": "safetensors"}, + load_inplace=True, + is_3d_lora_weight=True, + ), client_index=0, ) @@ -353,18 +376,6 @@ def engine_outputs_wire(output): ) -@dataclass -class KVEventsConfig: - enable_kv_cache_events: bool - publisher: str - endpoint: str - replay_endpoint: str | None - buffer_steps: int - hwm: int - max_queue_size: int - topic: str - - @dataclass class EngineCoreReadyResponse: max_model_len: int @@ -382,9 +393,15 @@ class EngineCoreReadyResponse: max_num_seqs: int max_num_batched_tokens: int instance_id: str + kv_event_block_size: int + supports_lora: bool + max_loras: int kv_cache_size_tokens: int | None = None kv_cache_max_concurrency: float | None = None - kv_events_config: KVEventsConfig | None = None + kv_role: str | None = None + kv_events_publisher: str | None = None + kv_events_endpoint: str | None = None + kv_events_topic: str | None = None ready_response = EngineCoreReadyResponse( @@ -396,23 +413,20 @@ class EngineCoreReadyResponse: vllm_version="0.0.0", data_parallel_size=1, world_size=1, - tensor_parallel_size=1, + tensor_parallel_size=2, pipeline_parallel_size=1, decode_context_parallel_size=1, data_parallel_rank=0, - max_num_seqs=256, - max_num_batched_tokens=8192, - instance_id="test-instance", - kv_events_config=KVEventsConfig( - enable_kv_cache_events=True, - publisher="zmq", - endpoint="tcp://127.0.0.1:5557", - replay_endpoint="tcp://127.0.0.1:5558", - buffer_steps=10_000, - hwm=100_000, - max_queue_size=100_000, - topic="kv", - ), + max_num_seqs=64, + max_num_batched_tokens=4096, + instance_id="engine-0", + kv_role="kv_both", + kv_events_publisher="zmq", + kv_events_endpoint="tcp://127.0.0.1:5557", + kv_events_topic="kv", + kv_event_block_size=256, + supports_lora=True, + max_loras=8, ) print(msgspec.msgpack.encode(request).hex()) diff --git a/rust/src/engine-core-client/src/transport.rs b/rust/src/engine-core-client/src/transport.rs index d3e3ecae2bf2..b55ea783206f 100644 --- a/rust/src/engine-core-client/src/transport.rs +++ b/rust/src/engine-core-client/src/transport.rs @@ -3,6 +3,7 @@ use std::collections::{BTreeMap, BTreeSet}; use std::fmt::Debug; +use std::net::TcpListener; use std::ops::Deref; use std::time::Duration; @@ -22,12 +23,14 @@ use crate::protocol::handshake::{ EngineCoreReadyResponse, HandshakeAddresses, HandshakeInitMessage, ReadyMessage, }; use crate::protocol::output::{EngineCoreOutputs, decode_engine_core_outputs}; -use crate::protocol::{decode_msgpack, encode_msgpack}; +use crate::protocol::{OpaqueValue, decode_msgpack, encode_msgpack}; /// Dedicated single-frame sentinel emitted by Python `EngineCoreProc` when the /// engine dies. pub const ENGINE_CORE_DEAD_SENTINEL: &[u8] = b"ENGINE_CORE_DEAD"; +const DATA_PARALLEL_BOOTSTRAP_PORT_COUNT: usize = 5; + /// Opaque routing identity of one engine on the frontend transport. #[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] pub struct EngineId(Bytes); @@ -121,6 +124,8 @@ pub struct ConnectedTransport { pub output_address: String, /// All engines connected through the startup handshake. pub engines: Vec, + /// Validated contiguous data-parallel rank span managed by this transport. + pub managed_data_parallel_span: ManagedDataParallelSpan, /// Optional engine-facing coordinator transport used for in-process wave /// coordination. pub coordinator: Option, @@ -131,6 +136,13 @@ pub struct ConnectedTransport { pub output_socket: PullSocket, } +/// Data-parallel ranks owned by one validated engine transport. +#[derive(Clone, Copy, Debug)] +pub struct ManagedDataParallelSpan { + pub start_rank: u32, + pub size: u32, +} + #[derive(Clone, Debug, EnumAsInner)] enum EngineStartupState { HelloReceived, @@ -181,9 +193,10 @@ pub async fn connect_handshake( handshake_socket.bind(handshake_address).await?; let mut engines = BTreeMap::new(); + let empty_parallel_config = BTreeMap::new(); - // 3. Receive HELLO from every engine and send a matching INIT. When coordinator mode is - // enabled, the engines will not emit READY until the coordinator barrier below completes. + // 3. Receive HELLO from every engine. Non-coordinated engines receive INIT immediately. A + // coordinated cohort receives one shared DP bootstrap after every expected HELLO arrives. while engines.len() < engine_count { debug!( handshake_address, @@ -207,17 +220,20 @@ pub async fn connect_handshake( } debug!(handshake_address, ?engine_id, "received HELLO from engine"); - send_init_message( - &mut handshake_socket, - &engine_id, - &input_address, - &output_address, - coordinator.as_ref(), - ) - .await?; - debug!(handshake_address, ?engine_id, "sent INIT to engine"); - engines.insert(engine_id.clone(), EngineStartupState::HelloReceived); + + if coordinator.is_none() { + send_init_message( + &mut handshake_socket, + &engine_id, + &input_address, + &output_address, + None, + &empty_parallel_config, + ) + .await?; + debug!(handshake_address, ?engine_id, "sent INIT to engine"); + } } Some("READY") => { if coordinator.is_some() { @@ -247,6 +263,28 @@ pub async fn connect_handshake( } } + // Probe bootstrap ports only after the complete coordinated cohort is waiting for INIT. The + // reservations are released immediately before Python must bind them. + if coordinator.is_some() { + let parallel_config = data_parallel_bootstrap(local_host, engine_count)?; + for engine_id in engines.keys() { + send_init_message( + &mut handshake_socket, + engine_id, + &input_address, + &output_address, + coordinator.as_ref(), + ¶llel_config, + ) + .await?; + debug!( + handshake_address, + ?engine_id, + "sent coordinated INIT to engine" + ); + } + } + // 4. Optional coordinator startup gate. Without coordinator there is nothing to do. if let Some(coordinator) = coordinator.as_mut() { coordinator.wait_for_startup_gate(engine_count, ready_timeout).await?; @@ -298,7 +336,7 @@ pub async fn connect_handshake( } // 6. Wait for every engine to connect to the shared input socket and register itself. - let engines = + let (engines, managed_data_parallel_span) = wait_for_input_registrations(&mut input_socket, engines.into_keys(), ready_timeout).await?; debug!( engine_count = engines.len(), @@ -315,6 +353,7 @@ pub async fn connect_handshake( input_send, output_socket, engines, + managed_data_parallel_span, coordinator, }) } @@ -343,7 +382,7 @@ pub async fn connect_bootstrapped( let mut output_socket = PullSocket::new(); let output_address = output_socket.bind(output_address).await?.to_string(); - let engines = wait_for_input_registrations( + let (engines, managed_data_parallel_span) = wait_for_input_registrations( &mut input_socket, (0..engine_count) .map(|offset| EngineId::from_engine_index(engine_start_index + offset as u32)), @@ -361,6 +400,7 @@ pub async fn connect_bootstrapped( input_address, output_address, engines, + managed_data_parallel_span, coordinator: None, input_send, output_socket, @@ -420,6 +460,7 @@ async fn send_init_message( input_address: &str, output_address: &str, coordinator: Option<&CoordinatorBootstrap>, + parallel_config: &BTreeMap, ) -> Result<()> { let init_message = HandshakeInitMessage { addresses: HandshakeAddresses { @@ -429,7 +470,7 @@ async fn send_init_message( coordinator_output: coordinator.map(|c| c.output_address.clone()), frontend_stats_publish_address: None, }, - parallel_config: Default::default(), + parallel_config: parallel_config.clone(), }; let payload = encode_msgpack(&init_message)?; let message = ZmqMessage::try_from(vec![engine_id.to_frame(), Bytes::from(payload)]) @@ -438,6 +479,49 @@ async fn send_init_message( Ok(()) } +/// Build the Python-compatible bootstrap shared by a coordinated DP cohort. +/// +/// TCP listeners reserve distinct ports while the map is assembled. They must +/// be released before INIT is sent because Python owns the actual listeners, +/// leaving an unavoidable probe-to-bind race after this function returns. +fn data_parallel_bootstrap( + advertised_host: &str, + engine_count: usize, +) -> Result> { + if engine_count <= 1 { + return Ok(BTreeMap::new()); + } + + let listeners = (0..DATA_PARALLEL_BOOTSTRAP_PORT_COUNT) + .map(|_| TcpListener::bind((advertised_host, 0))) + .collect::>>()?; + let mut ports = listeners + .iter() + .map(|listener| listener.local_addr().map(|address| u64::from(address.port()))) + .collect::>>()?; + let master_port = ports.pop().expect("bootstrap port count is nonzero"); + drop(listeners); + + Ok(BTreeMap::from([ + ( + "_data_parallel_master_port_list".to_string(), + OpaqueValue::Array(ports.into_iter().map(OpaqueValue::from).collect()), + ), + ( + "data_parallel_master_ip".to_string(), + OpaqueValue::from(advertised_host), + ), + ( + "data_parallel_master_port".to_string(), + OpaqueValue::from(master_port), + ), + ( + "data_parallel_size".to_string(), + OpaqueValue::from(engine_count as u64), + ), + ])) +} + /// Receive the input registration message from each engine and validate its /// identity. /// @@ -451,7 +535,7 @@ async fn wait_for_input_registrations( input_socket: &mut RouterSocket, expected_engines: impl IntoIterator, ready_timeout: Duration, -) -> Result> { +) -> Result<(Vec, ManagedDataParallelSpan)> { let expected_engines = expected_engines.into_iter().collect::>(); let mut pending = expected_engines.iter().cloned().collect::>(); let mut ready_responses = BTreeMap::new(); @@ -493,7 +577,7 @@ async fn wait_for_input_registrations( ready_responses.insert(actual_id, ready_response); } - Ok(expected_engines + let engines = expected_engines .into_iter() .map(|engine_id| { let ready_response = ready_responses @@ -504,7 +588,110 @@ async fn wait_for_input_registrations( ready_response, } }) - .collect()) + .collect::>(); + let managed_data_parallel_span = validate_ready_responses(&engines)?; + Ok((engines, managed_data_parallel_span)) +} + +/// Validate the private engine/frontend startup contract before publishing a +/// client. Incomplete or internally inconsistent metadata is a startup error. +fn validate_ready_responses(engines: &[ConnectedEngine]) -> Result { + let Some(first) = engines.first() else { + bail_unexpected_handshake_message!("no engine ready responses were received"); + }; + let expected = &first.ready_response; + let mut ranks = BTreeSet::new(); + + for engine in engines { + let response = &engine.ready_response; + if response.world_size == 0 + || response.tensor_parallel_size == 0 + || response.pipeline_parallel_size == 0 + || response.data_parallel_size == 0 + || response.max_num_seqs == 0 + || response.max_num_batched_tokens == 0 + { + bail_unexpected_handshake_message!( + "engine {:?} reported zero topology or scheduler capacity: {:?}", + engine.engine_id, + response + ); + } + if response.data_parallel_rank as u64 >= response.data_parallel_size { + bail_unexpected_handshake_message!( + "engine {:?} reported data-parallel rank {} outside size {}", + engine.engine_id, + response.data_parallel_rank, + response.data_parallel_size + ); + } + if !ranks.insert(response.data_parallel_rank) { + bail_unexpected_handshake_message!( + "duplicate data-parallel rank {} in engine ready responses", + response.data_parallel_rank + ); + } + if response.data_parallel_size > 1 + && let Some(engine_index) = engine.engine_id.engine_index() + && engine_index != response.data_parallel_rank + { + bail_unexpected_handshake_message!( + "engine identity rank {engine_index} does not match reported data-parallel rank {}", + response.data_parallel_rank + ); + } + if response.supports_lora != (response.max_loras > 0) { + bail_unexpected_handshake_message!( + "engine {:?} reported inconsistent LoRA capability (supports_lora={}, max_loras={})", + engine.engine_id, + response.supports_lora, + response.max_loras + ); + } + + let uniform = response.block_size == expected.block_size + && response.dtype == expected.dtype + && response.vllm_version == expected.vllm_version + && response.world_size == expected.world_size + && response.data_parallel_size == expected.data_parallel_size + && response.tensor_parallel_size == expected.tensor_parallel_size + && response.pipeline_parallel_size == expected.pipeline_parallel_size + && response.decode_context_parallel_size == expected.decode_context_parallel_size + && response.max_num_seqs == expected.max_num_seqs + && response.max_num_batched_tokens == expected.max_num_batched_tokens + && response.kv_role == expected.kv_role + && response.kv_events_publisher == expected.kv_events_publisher + && response.kv_events_endpoint == expected.kv_events_endpoint + && response.kv_events_topic == expected.kv_events_topic + && response.kv_event_block_size == expected.kv_event_block_size + && response.supports_lora == expected.supports_lora + && response.max_loras == expected.max_loras; + if !uniform { + bail_unexpected_handshake_message!( + "engine {:?} reported topology or capabilities inconsistent with engine {:?}", + engine.engine_id, + first.engine_id + ); + } + } + + let start_rank = + *ranks.first().expect("non-empty engines produce non-empty data-parallel ranks"); + let Ok(size) = u32::try_from(ranks.len()) else { + bail_unexpected_handshake_message!("managed data-parallel rank count exceeds u32"); + }; + for (offset, rank) in ranks.into_iter().enumerate() { + let Ok(offset) = u32::try_from(offset) else { + bail_unexpected_handshake_message!("managed data-parallel rank count exceeds u32"); + }; + if start_rank.checked_add(offset) != Some(rank) { + bail_unexpected_handshake_message!( + "managed data-parallel ranks must be contiguous, starting at rank {start_rank}" + ); + } + } + + Ok(ManagedDataParallelSpan { start_rank, size }) } /// Send an encoded message to the engine through the input socket. @@ -585,7 +772,9 @@ pub async fn run_output_loop( #[cfg(test)] mod tests { - use super::bind_local_sockets; + use super::{ConnectedEngine, bind_local_sockets, validate_ready_responses}; + use crate::EngineId; + use crate::mock_engine::default_ready_response; #[tokio::test] async fn bind_local_sockets_resolves_zero_port_bindings() { @@ -596,4 +785,102 @@ mod tests { assert!(output_address.starts_with("tcp://127.0.0.1:")); assert_ne!(input_address, output_address); } + + #[test] + fn ready_validation_accepts_uniform_data_parallel_metadata() { + let engines = [0, 1].map(|rank| { + let mut ready_response = default_ready_response(); + ready_response.data_parallel_size = 2; + ready_response.data_parallel_rank = rank; + ConnectedEngine { + engine_id: EngineId::from_engine_index(rank), + ready_response, + } + }); + validate_ready_responses(&engines).expect("valid ready responses"); + } + + #[test] + fn ready_validation_accepts_contiguous_offset_data_parallel_ranks() { + let engines = [2, 3].map(|rank| { + let mut ready_response = default_ready_response(); + ready_response.data_parallel_size = 4; + ready_response.data_parallel_rank = rank; + ConnectedEngine { + engine_id: EngineId::from_engine_index(rank), + ready_response, + } + }); + let span = validate_ready_responses(&engines).expect("contiguous offset ranks are valid"); + assert_eq!(span.start_rank, 2); + assert_eq!(span.size, 2); + } + + #[test] + fn ready_validation_rejects_noncontiguous_data_parallel_ranks() { + let engines = [0, 2].map(|rank| { + let mut ready_response = default_ready_response(); + ready_response.data_parallel_size = 4; + ready_response.data_parallel_rank = rank; + ConnectedEngine { + engine_id: EngineId::from_engine_index(rank), + ready_response, + } + }); + assert!(validate_ready_responses(&engines).is_err()); + } + + #[test] + fn ready_validation_rejects_duplicate_data_parallel_ranks() { + let engines = [0, 1].map(|rank| { + let mut ready_response = default_ready_response(); + ready_response.data_parallel_size = 2; + ConnectedEngine { + engine_id: EngineId::from_engine_index(rank), + ready_response, + } + }); + assert!(validate_ready_responses(&engines).is_err()); + } + + #[test] + fn ready_validation_rejects_duplicate_rank_for_single_rank_topology() { + let engines = [0, 1].map(|engine_index| ConnectedEngine { + engine_id: EngineId::from_engine_index(engine_index), + ready_response: default_ready_response(), + }); + assert!(validate_ready_responses(&engines).is_err()); + } + + #[test] + fn ready_validation_rejects_nonuniform_topology() { + let mut first = default_ready_response(); + first.data_parallel_size = 2; + let mut second = first.clone(); + second.data_parallel_rank = 1; + second.tensor_parallel_size = 2; + let engines = [ + ConnectedEngine { + engine_id: EngineId::from_engine_index(0), + ready_response: first, + }, + ConnectedEngine { + engine_id: EngineId::from_engine_index(1), + ready_response: second, + }, + ]; + assert!(validate_ready_responses(&engines).is_err()); + } + + #[test] + fn ready_validation_rejects_inconsistent_lora_capacity() { + let mut ready_response = default_ready_response(); + ready_response.supports_lora = true; + ready_response.max_loras = 0; + let engines = [ConnectedEngine { + engine_id: EngineId::from_engine_index(0), + ready_response, + }]; + assert!(validate_ready_responses(&engines).is_err()); + } } diff --git a/rust/src/llm/src/output.rs b/rust/src/llm/src/output.rs index c8f01b9fb9a3..fb0c6f73d15b 100644 --- a/rust/src/llm/src/output.rs +++ b/rust/src/llm/src/output.rs @@ -11,7 +11,10 @@ use futures::stream::FusedStream; use futures::{Stream, StreamExt as _, pin_mut}; use serde::{Deserialize, Serialize}; use vllm_engine_core_client::protocol::logprobs::Logprobs; -use vllm_engine_core_client::protocol::output::{EngineCoreFinishReason, StopReason}; +use vllm_engine_core_client::protocol::output::{ + EngineCoreFinishReason, StopReason, concatenate_routed_experts, +}; +use vllm_engine_core_client::protocol::tensor::WireNdArray; use vllm_engine_core_client::{AbortCause, EngineCoreOutputStream}; use crate::error::Result; @@ -44,6 +47,8 @@ pub struct CollectedGenerateOutput { /// Connector-specific encoder cache transfer parameters for disaggregated /// serving. pub ec_transfer_params: Option, + /// Routed-expert IDs concatenated along the token axis. + pub routed_experts: Option, } /// Prompt-scoped metadata emitted only once on the first [`GenerateOutput`] for @@ -155,6 +160,8 @@ pub struct GenerateOutput { /// Connector-specific encoder cache transfer parameters for disaggregated /// serving. pub ec_transfer_params: Option, + /// Routed-expert IDs emitted by engine-core for this output chunk. + pub routed_experts: Option, } impl GenerateOutput { @@ -202,6 +209,7 @@ impl GenerateOutput { cached_token_count: 0, kv_transfer_params: None, ec_transfer_params: None, + routed_experts: None, } } } @@ -291,6 +299,7 @@ impl Stream for GenerateOutputStream { cached_token_count, kv_transfer_params: raw.kv_transfer_params, ec_transfer_params: raw.ec_transfer_params, + routed_experts: raw.routed_experts, }; Poll::Ready(Some(Ok(output))) @@ -337,9 +346,13 @@ impl> + Send> T { let mut prompt_token_ids = None; let mut prompt_logprobs = None; let mut cached_token_count = 0; + let mut routed_experts_chunks = Vec::new(); let mut collected: Option = None; while let Some(output) = stream.next().await.transpose()? { + if let Some(routed_experts) = output.routed_experts { + routed_experts_chunks.push(routed_experts); + } cached_token_count = cached_token_count.max(output.cached_token_count); if let Some(info) = output.prompt_info { if prompt_token_ids.is_none() { @@ -374,6 +387,7 @@ impl> + Send> T { }, kv_transfer_params: None, ec_transfer_params: None, + routed_experts: None, }); } @@ -387,6 +401,7 @@ impl> + Send> T { }; collected.kv_transfer_params = output.kv_transfer_params; collected.ec_transfer_params = output.ec_transfer_params; + collected.routed_experts = concatenate_routed_experts(routed_experts_chunks)?; return Ok(collected); } } diff --git a/rust/src/server/Cargo.toml b/rust/src/server/Cargo.toml index 1ed397d589cc..d512bd7a1673 100644 --- a/rust/src/server/Cargo.toml +++ b/rust/src/server/Cargo.toml @@ -21,6 +21,7 @@ llm-multimodal.workspace = true openssl.workspace = true prost.workspace = true prost-types.workspace = true +rmp-serde.workspace = true rmpv.workspace = true serde.workspace = true serde_json.workspace = true @@ -59,7 +60,6 @@ async-openai = { workspace = true, features = ["full"] } bytes.workspace = true clap.workspace = true expect-test.workspace = true -rmp-serde.workspace = true serial_test.workspace = true tempfile.workspace = true tokio = { workspace = true, features = ["test-util"] } diff --git a/rust/src/server/src/grpc/control.rs b/rust/src/server/src/grpc/control.rs index 897209822252..c6934e55669f 100644 --- a/rust/src/server/src/grpc/control.rs +++ b/rust/src/server/src/grpc/control.rs @@ -1,25 +1,65 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project +use std::path::PathBuf; use std::sync::Arc; use thiserror_ext::AsReport as _; use tonic::{Request, Response, Status}; +use tonic_health::server::HealthReporter; use vllm_engine_core_client::protocol::handshake::EngineCoreReadyResponse; -use super::{ControlServer, pb}; +use super::{AdmissionGuard, AdmissionState, ControlServer, lora_rpc, pb}; +use crate::lora_path::runtime_lora_allowed_path_prefixes; use crate::state::AppState; pub(crate) type ControlGrpcService = ControlServer; +const GRPC_API_VERSION: &str = "vllm"; +const GRPC_CAPABILITIES: &[&str] = &[ + "generate.sampling.v2", + "generate.preprocessed_mm.v1", + "generate.routed_experts.v1", +]; + /// gRPC control service backed by the shared application state. pub struct ControlServiceImpl { state: Arc, + admission: Arc, + health_reporter: Option, + lora_allowed_path_prefixes: Option>, + runtime_lora_updating_enabled: bool, } impl ControlServiceImpl { pub fn new(state: Arc) -> Self { - Self { state } + Self::with_admission(state, Arc::new(AdmissionState::default()), None) + } + + pub(crate) fn with_admission( + state: Arc, + admission: Arc, + health_reporter: Option, + ) -> Self { + Self { + state, + admission, + health_reporter, + lora_allowed_path_prefixes: runtime_lora_allowed_path_prefixes().map(Arc::from), + runtime_lora_updating_enabled: crate::routes::runtime_lora_updating_enabled(), + } + } + + #[cfg(test)] + pub(crate) fn with_lora_allowed_path_prefixes(mut self, prefixes: Vec) -> Self { + self.lora_allowed_path_prefixes = Some(prefixes.into()); + self + } + + #[cfg(test)] + pub(crate) fn with_runtime_lora_updating(mut self, enabled: bool) -> Self { + self.runtime_lora_updating_enabled = enabled; + self } fn ready(&self) -> &EngineCoreReadyResponse { @@ -28,17 +68,29 @@ impl ControlServiceImpl { fn parallelism_info(&self) -> pb::ParallelismInfo { let ready = self.ready(); + let (data_parallel_start_rank, managed_data_parallel_size) = + self.state.engine_core_client().managed_data_parallel_span(); pb::ParallelismInfo { tensor_parallel_size: ready.tensor_parallel_size, pipeline_parallel_size: ready.pipeline_parallel_size, data_parallel_size: ready.data_parallel_size.min(u64::from(u32::MAX)) as u32, data_parallel_rank: ready.data_parallel_rank, decode_context_parallel_size: ready.decode_context_parallel_size, + data_parallel_start_rank, + managed_data_parallel_size, } } -} -const GRPC_API_VERSION: &str = "vllm"; + async fn report_not_serving(&self) { + if let Some(reporter) = &self.health_reporter { + crate::set_generate_not_serving(reporter).await; + } + } + + fn try_admit(&self) -> Option { + self.admission.try_admit() + } +} #[tonic::async_trait] impl pb::control_server::Control for ControlServiceImpl { @@ -57,6 +109,11 @@ impl pb::control_server::Control for ControlServiceImpl { total_kv_blocks: self.state.engine_core_client().total_num_gpu_blocks(), max_running_requests: ready.max_num_seqs, max_batched_tokens: ready.max_num_batched_tokens, + max_loras: ready.max_loras, + capabilities: GRPC_CAPABILITIES + .iter() + .map(|capability| (*capability).to_string()) + .collect(), })) } @@ -64,14 +121,17 @@ impl pb::control_server::Control for ControlServiceImpl { &self, _request: Request, ) -> Result, Status> { + let ready = self.ready(); let served = self.state.served_model_names(); Ok(Response::new(pb::ModelInfo { model_id: self.state.chat.text().model_id().to_string(), served_model_name: self.state.primary_model_name().to_string(), served_model_aliases: served.iter().skip(1).cloned().collect(), + tokenizer_modes: Vec::new(), // GenerateRequest accepts both prompt representations. supports_text_input: true, supports_token_ids_input: true, + supports_lora: ready.supports_lora, supports_multimodal: self.state.chat.supports_multimodal(), reasoning_parser: self .state @@ -104,32 +164,140 @@ impl pb::control_server::Control for ControlServiceImpl { Ok(Response::new(pb::AbortResponse {})) } + async fn drain( + &self, + _request: Request, + ) -> Result, Status> { + self.admission.begin_drain(); + self.report_not_serving().await; + let in_flight = self.admission.in_flight().min(u64::from(u32::MAX)) as u32; + let state = if in_flight == 0 { + pb::DrainState::Complete + } else { + pb::DrainState::InProgress + }; + Ok(Response::new(pb::DrainResponse { + state: state as i32, + in_flight_requests: in_flight, + message: String::new(), + })) + } + + async fn load_lora( + &self, + request: Request, + ) -> Result, Status> { + if !self.runtime_lora_updating_enabled { + return Err(Status::failed_precondition( + "runtime LoRA updating is disabled", + )); + } + let _guard = self + .try_admit() + .ok_or_else(|| Status::unavailable("gRPC service is draining"))?; + lora_rpc::load_lora( + &self.state, + self.lora_allowed_path_prefixes.as_deref(), + request, + ) + .await + } + + async fn unload_lora( + &self, + request: Request, + ) -> Result, Status> { + if !self.runtime_lora_updating_enabled { + return Err(Status::failed_precondition( + "runtime LoRA updating is disabled", + )); + } + let _guard = self + .try_admit() + .ok_or_else(|| Status::unavailable("gRPC service is draining"))?; + lora_rpc::unload_lora(&self.state, request).await + } + + async fn list_loras( + &self, + request: Request, + ) -> Result, Status> { + lora_rpc::list_loras(&self.state, request).await + } + async fn get_kv_event_sources( &self, _request: Request, ) -> Result, Status> { - let client = self.state.engine_core_client(); - let sources = client.ready_responses().into_iter().filter_map(kv_event_source).collect(); + let sources = self + .state + .engine_core_client() + .ready_responses() + .into_iter() + .filter_map(kv_event_source) + .collect(); Ok(Response::new(pb::GetKvEventSourcesResponse { sources })) } } -pub(super) fn kv_event_source(response: &EngineCoreReadyResponse) -> Option { - let config = response.kv_events_config.as_ref()?; - if !config.enable_kv_cache_events || config.publisher != "zmq" { +fn kv_event_source(response: &EngineCoreReadyResponse) -> Option { + if response.kv_events_publisher.as_deref() != Some("zmq") { return None; } - + let endpoint = offset_endpoint_port( + response.kv_events_endpoint.as_ref()?, + response.data_parallel_rank, + ); Some(pb::KvEventSource { transport: "zmq".to_string(), - endpoint: config.endpoint.clone(), - topic: config.topic.clone(), - replay_endpoint: config.replay_endpoint.clone().unwrap_or_default(), + endpoint_addr: Some(kv_endpoint_from_zmq(&endpoint)?), + topic: response.kv_events_topic.clone().unwrap_or_default(), + replay_endpoint: String::new(), data_parallel_rank: Some(response.data_parallel_rank), encoding: "msgpack".to_string(), schema_version: 1, - buffer_steps: config.buffer_steps, - hwm: config.hwm, - max_queue_size: config.max_queue_size, + buffer_steps: 0, + hwm: 0, + max_queue_size: 0, }) } + +fn offset_endpoint_port(endpoint: &str, data_parallel_rank: u32) -> String { + if data_parallel_rank == 0 || endpoint.is_empty() { + return endpoint.to_string(); + } + if endpoint.contains("inproc") { + return format!("{endpoint}_dp{data_parallel_rank}"); + } + if endpoint.contains("tcp") + && let Some((base_addr, port)) = endpoint.rsplit_once(':') + && let Ok(base_port) = port.parse::() + { + return format!("{base_addr}:{}", base_port + data_parallel_rank); + } + endpoint.to_string() +} + +fn kv_endpoint_from_zmq(endpoint: &str) -> Option { + let rest = endpoint.strip_prefix("tcp://").unwrap_or(endpoint); + let (host, port) = rest.rsplit_once(':')?; + let port = port.parse().ok()?; + let host = match host.trim_matches(|character| character == '[' || character == ']') { + "*" | "0.0.0.0" | "::" | "" => advertise_host(), + concrete => concrete.to_string(), + }; + Some(pb::KvEventEndpoint { + host, + port, + protocol: "tcp".to_string(), + }) +} + +fn advertise_host() -> String { + std::net::UdpSocket::bind("0.0.0.0:0") + .and_then(|socket| { + socket.connect("10.255.255.255:1")?; + Ok(socket.local_addr()?.ip().to_string()) + }) + .unwrap_or_else(|_| "127.0.0.1".to_string()) +} diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index 2f13422fecc3..64a47e744fda 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -4,16 +4,75 @@ //! Conversion between gRPC protobuf types and internal `vllm-text` //! request/response types. +mod multimodal; +mod sampling; +mod xargs; + +use multimodal::convert_mm_features; +pub(crate) use multimodal::media_parts_from_request; +#[cfg(test)] +use multimodal::{mm_cache_identifier, preflight_msgpack}; +use sampling::build_sampling_params; +use xargs::parse_vllm_xargs_json; + use tonic::Status; use uuid::Uuid; use vllm_engine_core_client::protocol::output::StopReason; -use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputsParams; +use vllm_engine_core_client::protocol::tensor::{WireArrayData, WireTensor}; use vllm_text::{ - DecodedLogprobs, DecodedPromptLogprobs, FinishReason, Finished, Prompt, SamplingParams, - TextDecodeOptions, TextRequest, + DecodedLogprobs, DecodedPromptLogprobs, FinishReason, Finished, Prompt, TextDecodeOptions, + TextRequest, }; use super::pb; +use super::struct_json::{json_to_prost_struct, prost_struct_to_json}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum KvRole { + Aggregated, + Prefill, + Decode, +} + +pub fn role_from_kv_role(kv_role: Option<&str>) -> KvRole { + match kv_role { + Some("kv_producer") => KvRole::Prefill, + Some("kv_consumer") => KvRole::Decode, + _ => KvRole::Aggregated, + } +} + +pub fn validate_disaggregated_request( + request: &pb::GenerateRequest, + role: KvRole, +) -> Result<(), Status> { + let has_transfer_params = request + .kv + .as_ref() + .and_then(|kv| kv.kv_transfer_params.as_ref()) + .is_some_and(|params| !params.fields.is_empty()); + if role == KvRole::Decode && !has_transfer_params { + return Err(Status::invalid_argument( + "kv.kv_transfer_params is required for decode requests", + )); + } + Ok(()) +} + +pub fn mark_prefill_request(request: &mut TextRequest) { + let params = request + .sampling_params + .vllm_xargs + .get_or_insert_with(Default::default) + .entry("kv_transfer_params".to_string()) + .or_insert_with(|| serde_json::Value::Object(Default::default())); + if let Some(params) = params.as_object_mut() { + params.insert( + "do_remote_decode".to_string(), + serde_json::Value::Bool(true), + ); + } +} // ======================================================================================== // Request conversion @@ -42,11 +101,35 @@ pub fn to_text_request( )); } + if !req.media.is_empty() && !req.mm_features.is_empty() { + return Err(Status::invalid_argument( + "media and mm_features are mutually exclusive", + )); + } + if !req.lora_name.is_empty() && (!req.media.is_empty() || !req.mm_features.is_empty()) { + return Err(Status::invalid_argument( + "native gRPC does not yet advertise tower-LoRA multimodal cache semantics; multimodal requests with LoRA are unsupported", + )); + } + let prompt = match req.prompt { Some(pb::generate_request::Prompt::Text(text)) => Prompt::Text(text), Some(pb::generate_request::Prompt::TokenIds(ids)) => Prompt::TokenIds(ids.ids), None => return Err(Status::invalid_argument("prompt is required")), }; + match &prompt { + Prompt::TokenIds(ids) if req.routed_experts_prompt_start as usize >= ids.len() => { + return Err(Status::invalid_argument( + "routed_experts_prompt_start must be less than the prompt length", + )); + } + Prompt::Text(_) if req.routed_experts_prompt_start != 0 => { + return Err(Status::invalid_argument( + "nonzero routed_experts_prompt_start requires a token-ID prompt", + )); + } + _ => {} + } let request_id = if req.request_id.is_empty() { Uuid::new_v4().to_string() @@ -62,21 +145,27 @@ pub fn to_text_request( let mut sampling_params = build_sampling_params(req.temperature, sampling, decoding, stopping, response)?; + sampling_params.routed_experts_prompt_start = req.routed_experts_prompt_start; + + if let Some(raw) = req.vllm_xargs_json.as_deref() { + let xargs = parse_vllm_xargs_json(raw)?; + sampling_params.vllm_xargs = Some(xargs); + } // Thread KVCacheParameters → SamplingParams fields. if let Some(kv) = kv { // Thread kv_transfer_params through vllm_xargs, matching the HTTP route // convention. if let Some(kv_struct) = kv.kv_transfer_params.as_ref() { - let kv_json = proto_struct_to_json(kv_struct); + let kv_json = prost_struct_to_json(kv_struct); let map = sampling_params.vllm_xargs.get_or_insert_with(Default::default); + if map.contains_key("kv_transfer_params") { + return Err(Status::invalid_argument( + "kv_transfer_params cannot be supplied in both vllm_xargs_json and kv", + )); + } map.insert("kv_transfer_params".to_string(), kv_json); } - if let Some(ec_struct) = kv.ec_transfer_params.as_ref() { - let ec_json = proto_struct_to_json(ec_struct); - let map = sampling_params.vllm_xargs.get_or_insert_with(Default::default); - map.insert("ec_transfer_params".to_string(), ec_json); - } if kv.bypass_prefix_cache { sampling_params.skip_reading_prefix_cache = Some(true); } @@ -89,10 +178,12 @@ pub fn to_text_request( min_tokens: stopping.map_or(0, |s| s.min_new_tokens), }; + let mm_features = convert_mm_features(&req.mm_features, &prompt)?; + Ok(TextRequest { request_id, prompt, - mm_features: None, + mm_features, sampling_params, decode_options, intermediate: stream, @@ -106,156 +197,6 @@ pub fn to_text_request( }) } -fn build_sampling_params( - temperature: Option, - sampling: Option<&pb::RandomSampling>, - decoding: Option<&pb::DecodingParameters>, - stopping: Option<&pb::StoppingCriteria>, - response: Option<&pb::ResponseOptions>, -) -> Result { - // Temperature is a top-level GenerateRequest field. Default to greedy (0.0) for - // the gRPC API when the caller does not specify a value. This differs from - // the HTTP/OpenAI API (which defaults to 1.0) and matches the convention of - // programmatic generation APIs. - let temperature = temperature.or(Some(0.0)); - let mut params = SamplingParams { - temperature, - ..SamplingParams::default() - }; - - // RandomSampling: for every remaining sampling field the protobuf default (`0`) - // is treated as "unset" and leaves the resolved value to the lowering - // stage, which falls back to the model-provided default or a - // neutral/disabled value otherwise. - if let Some(s) = sampling { - // num_sequences (n > 1) is not supported yet by the TextLlm layer; the response - // path also hardcodes SequenceOutput.index = 0, so accepting >1 would silently - // truncate output cardinality. Reject explicitly. - if s.num_sequences > 1 { - return Err(Status::invalid_argument( - "num_sequences > 1 is not supported", - )); - } - if s.top_k != 0 { - params.top_k = Some(s.top_k); - } - if s.top_p != 0.0 { - params.top_p = Some(s.top_p); - } - if s.min_p != 0.0 { - params.min_p = Some(s.min_p); - } - params.seed = s.seed; - } - - // DecodingParameters - if let Some(d) = decoding { - if d.presence_penalty != 0.0 { - params.presence_penalty = Some(d.presence_penalty); - } - if d.frequency_penalty != 0.0 { - params.frequency_penalty = Some(d.frequency_penalty); - } - if d.repetition_penalty != 0.0 { - params.repetition_penalty = Some(d.repetition_penalty); - } - if !d.logit_bias.is_empty() { - params.logit_bias = Some(d.logit_bias.clone()); - } - if !d.allowed_token_ids.is_empty() { - params.allowed_token_ids = Some(d.allowed_token_ids.clone()); - } - params.structured_outputs = convert_structured_output(d)?; - } - - // StoppingCriteria - if let Some(s) = stopping { - if s.max_new_tokens != 0 { - params.max_tokens = Some(s.max_new_tokens); - } - if s.min_new_tokens != 0 { - params.min_tokens = Some(s.min_new_tokens); - } - if !s.stop_token_ids.is_empty() { - params.stop_token_ids = Some(s.stop_token_ids.clone()); - } - params.ignore_eos = s.ignore_eos; - } - - // ResponseOptions → logprobs - if let Some(r) = response { - if r.output_logprobs { - let (count, token_ids) = candidate_logprob_spec(r.output_candidates.as_ref()); - params.logprobs = Some(count); - params.logprob_token_ids = token_ids; - } - if r.prompt_logprobs { - // The engine-core protocol has only one shared `logprob_token_ids` field - // for output and prompt logprobs, so a per-token-id selector for prompt - // candidates can't be honored independently. Reject it instead of silently - // dropping the list. - if matches!( - r.prompt_candidates.as_ref().and_then(|c| c.select.as_ref()), - Some(pb::candidate_tokens::Select::TokenIds(_)) - ) { - return Err(Status::invalid_argument( - "prompt_candidates token_ids selector is not supported", - )); - } - let (count, _) = candidate_logprob_spec(r.prompt_candidates.as_ref()); - params.prompt_logprobs = Some(count); - } - } - - Ok(params) -} - -/// Map the proto `CandidateTokens` selector to a `(logprobs_count, -/// logprob_token_ids)` pair. -/// -/// - `top_n(k)` → `(k, None)` — return top-k candidates by probability -/// - `all` → `(-1, None)` — return the full vocabulary -/// - `token_ids(n)` → `(1, Some(vec of n token ids))` — return logprobs for specific tokens (the -/// count `n` is stored in the proto as the number of token IDs that follow, but the actual IDs -/// are carried via `logprob_token_ids` on `SamplingParams`) -/// - absent → `(1, None)` — just the sampled/scored token -fn candidate_logprob_spec(candidates: Option<&pb::CandidateTokens>) -> (i32, Option>) { - match candidates.and_then(|c| c.select.as_ref()) { - Some(pb::candidate_tokens::Select::TopN(n)) => (*n as i32, None), - Some(pb::candidate_tokens::Select::All(true)) => (-1, None), - Some(pb::candidate_tokens::Select::TokenIds(ids)) => (1, Some(ids.ids.clone())), - _ => (1, None), - } -} - -fn convert_structured_output( - d: &pb::DecodingParameters, -) -> Result, Status> { - let so = match d.structured_output.as_ref() { - None => return Ok(None), - Some(so) => so, - }; - use pb::decoding_parameters::StructuredOutput; - let params = match so { - StructuredOutput::Json(schema) => { - let json: serde_json::Value = serde_json::from_str(schema) - .map_err(|e| Status::invalid_argument(format!("invalid json schema: {e}")))?; - StructuredOutputsParams::json(json) - } - StructuredOutput::Regex(regex) => StructuredOutputsParams::regex(regex.clone()), - StructuredOutput::Choice(choices) => { - StructuredOutputsParams::choice(choices.choices.clone()) - } - StructuredOutput::Grammar(grammar) => StructuredOutputsParams::grammar(grammar.clone()), - StructuredOutput::JsonObject(true) => StructuredOutputsParams::json_object(), - StructuredOutput::JsonObject(false) => return Ok(None), - StructuredOutput::StructuralTag(tag) => { - StructuredOutputsParams::structural_tag(tag.clone()) - } - }; - Ok(Some(params)) -} - // ======================================================================================== // Response conversion // ======================================================================================== @@ -317,6 +258,22 @@ pub fn to_sequence_output( ranks: rank_values, candidate_tokens: candidates, finish_info: finished.map(|f| to_finish_info(f, token_ids)), + routed_experts: finished + .and_then(|finished| finished.routed_experts.as_ref().map(routed_experts_to_proto)), + } +} + +fn routed_experts_to_proto(tensor: &WireTensor) -> pb::RoutedExpertsTensor { + let data = match &tensor.data { + WireArrayData::RawView(data) => data.clone(), + WireArrayData::AuxIndex(_) => { + unreachable!("engine-core output arrays are resolved before response conversion") + } + }; + pb::RoutedExpertsTensor { + dtype: tensor.dtype.clone(), + shape: tensor.shape.iter().map(|&dim| dim as u64).collect(), + data, } } @@ -350,8 +307,7 @@ fn to_finish_info(finished: &Finished, token_ids: &[u32]) -> pb::FinishInfo { num_output_tokens: finished.usage.output_token_count as u32, finish_reason, stop_reason, - kv_transfer_params: finished.kv_transfer_params.as_ref().and_then(json_to_proto_struct), - ec_transfer_params: finished.ec_transfer_params.as_ref().and_then(json_to_proto_struct), + kv_transfer_params: finished.kv_transfer_params.as_ref().and_then(json_to_prost_struct), } } @@ -397,72 +353,30 @@ fn positions_to_proto( if let Some(first) = pos.entries.first() { logprobs.push(first.logprob); ranks.push(first.rank); - } - - // Extra candidates beyond the first. - let entries = pos.entries.iter().skip(1); - candidates.push(pb::CandidateTokenInfo { - tokens: entries - .map(|e| pb::candidate_token_info::TokenInfo { - id: e.token_id, - logprob: e.logprob, - rank: e.rank, - }) - .collect(), - }); - } - (logprobs, ranks, candidates) -} - -// ======================================================================================== -// KV transfer params conversion (serde_json::Value ↔ prost_types::Struct) -// ======================================================================================== - -fn proto_struct_to_json(s: &prost_types::Struct) -> serde_json::Value { - serde_json::Value::Object( - s.fields.iter().map(|(k, v)| (k.clone(), proto_value_to_json(v))).collect(), - ) -} - -fn proto_value_to_json(v: &prost_types::Value) -> serde_json::Value { - use prost_types::value::Kind; - match v.kind.as_ref() { - None | Some(Kind::NullValue(_)) => serde_json::Value::Null, - Some(Kind::BoolValue(b)) => serde_json::Value::Bool(*b), - Some(Kind::NumberValue(n)) => serde_json::json!(*n), - Some(Kind::StringValue(s)) => serde_json::Value::String(s.clone()), - Some(Kind::ListValue(list)) => { - serde_json::Value::Array(list.values.iter().map(proto_value_to_json).collect()) + // Engine-core can include the sampled token again in its top-k + // alternatives. The gRPC schema carries that token separately in + // the parallel token_ids/logprobs/ranks fields, so do not repeat it + // in CandidateTokenInfo (whose contract is alternatives only). + candidates.push(pb::CandidateTokenInfo { + tokens: pos + .entries + .iter() + .skip(1) + .filter(|entry| entry.token_id != first.token_id) + .map(|entry| pb::candidate_token_info::TokenInfo { + id: entry.token_id, + logprob: entry.logprob, + rank: entry.rank, + }) + .collect(), + }); + } else { + candidates.push(pb::CandidateTokenInfo { tokens: vec![] }); } - Some(Kind::StructValue(s)) => proto_struct_to_json(s), } -} -fn json_to_proto_struct(value: &serde_json::Value) -> Option { - match value { - serde_json::Value::Object(map) => Some(prost_types::Struct { - fields: map.iter().map(|(k, v)| (k.clone(), json_to_proto_value(v))).collect(), - }), - _ => None, - } -} - -fn json_to_proto_value(v: &serde_json::Value) -> prost_types::Value { - use prost_types::value::Kind; - let kind = match v { - serde_json::Value::Null => Kind::NullValue(0), - serde_json::Value::Bool(b) => Kind::BoolValue(*b), - serde_json::Value::Number(n) => Kind::NumberValue(n.as_f64().unwrap_or(0.0)), - serde_json::Value::String(s) => Kind::StringValue(s.clone()), - serde_json::Value::Array(arr) => Kind::ListValue(prost_types::ListValue { - values: arr.iter().map(json_to_proto_value).collect(), - }), - serde_json::Value::Object(map) => Kind::StructValue(prost_types::Struct { - fields: map.iter().map(|(k, v)| (k.clone(), json_to_proto_value(v))).collect(), - }), - }; - prost_types::Value { kind: Some(kind) } + (logprobs, ranks, candidates) } // ======================================================================================== @@ -499,11 +413,25 @@ impl ResponseOpts { #[cfg(test)] mod tests { + use std::collections::BTreeMap; + + use vllm_engine_core_client::protocol::multimodal::{ + MmBatchedField, MmField, MmFieldElem, MmKwargValue, + }; use vllm_engine_core_client::protocol::output::StopReason; - use vllm_text::{FinishReason, Finished, Prompt}; + use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputConstraint; + use vllm_engine_core_client::protocol::tensor::WireTensor; + use vllm_text::{ + DecodedLogprobs, DecodedPositionLogprobs, DecodedTokenLogprob, FinishReason, Finished, + Prompt, + }; use super::pb::finish_info::{FinishReason as PbFinishReason, StopReason as PbStopReason}; - use super::{ResponseOpts, pb, to_finish_info, to_sequence_output, to_text_request}; + use super::{ + KvRole, ResponseOpts, mark_prefill_request, media_parts_from_request, mm_cache_identifier, + pb, preflight_msgpack, to_finish_info, to_sequence_output, to_text_request, + validate_disaggregated_request, + }; fn base_request() -> pb::GenerateRequest { pb::GenerateRequest { @@ -514,6 +442,45 @@ mod tests { } } + #[test] + fn media_requires_a_source() { + let error = media_parts_from_request(&[pb::MediaItem { + modality: pb::Modality::Image as i32, + ..Default::default() + }]) + .unwrap_err(); + assert_eq!(error.code(), tonic::Code::InvalidArgument); + } + + #[test] + fn media_rejects_unsupported_modalities() { + let error = media_parts_from_request(&[pb::MediaItem { + modality: pb::Modality::Audio as i32, + source: Some(pb::media_item::Source::Url( + "https://example.test/a.wav".into(), + )), + ..Default::default() + }]) + .unwrap_err(); + assert_eq!(error.code(), tonic::Code::Unimplemented); + } + + #[test] + fn decode_requires_kv_transfer_params() { + let error = validate_disaggregated_request(&base_request(), KvRole::Decode).unwrap_err(); + assert_eq!(error.code(), tonic::Code::InvalidArgument); + } + + #[test] + fn prefill_marks_remote_decode() { + let mut text = to_text_request(base_request(), false, &["test-model".to_string()]).unwrap(); + mark_prefill_request(&mut text); + assert_eq!( + text.sampling_params.vllm_xargs.as_ref().unwrap()["kv_transfer_params"]["do_remote_decode"], + serde_json::Value::Bool(true) + ); + } + #[test] fn temperature_propagates_from_top_level_request_field() { let req = pb::GenerateRequest { @@ -545,6 +512,39 @@ mod tests { assert_eq!(text.sampling_params.seed, None); } + #[test] + fn routed_experts_prompt_start_reaches_text_sampling_params() { + let req = pb::GenerateRequest { + prompt: Some(pb::generate_request::Prompt::TokenIds(pb::TokenIds { + ids: vec![1, 2], + })), + routed_experts_prompt_start: 1, + ..base_request() + }; + + let text = to_text_request(req, false, &["test-model".to_string()]).unwrap(); + + assert_eq!(text.sampling_params.routed_experts_prompt_start, 1); + } + + #[test] + fn routed_experts_prompt_start_rejects_unverifiable_or_out_of_range_values() { + let text_prompt = pb::GenerateRequest { + routed_experts_prompt_start: 1, + ..base_request() + }; + assert!(to_text_request(text_prompt, false, &["test-model".to_string()]).is_err()); + + let token_prompt = pb::GenerateRequest { + prompt: Some(pb::generate_request::Prompt::TokenIds(pb::TokenIds { + ids: vec![1, 2], + })), + routed_experts_prompt_start: 2, + ..base_request() + }; + assert!(to_text_request(token_prompt, false, &["test-model".to_string()]).is_err()); + } + #[test] fn zero_seed_is_valid() { let req = pb::GenerateRequest { @@ -586,6 +586,293 @@ mod tests { assert!(matches!(text.prompt, Prompt::Text(s) if s == "hi")); } + #[test] + fn extended_sampling_fields_reach_text_request_losslessly() { + let req = pb::GenerateRequest { + prompt: Some(pb::generate_request::Prompt::TokenIds(pb::TokenIds { + ids: vec![1, 2, 3], + })), + sampling: Some(pb::RandomSampling { + top_k: Some(-1), + top_p: Some(0.0), + min_p: Some(0.0), + ..Default::default() + }), + decoding: Some(pb::DecodingParameters { + presence_penalty: Some(0.0), + frequency_penalty: Some(0.0), + repetition_penalty: Some(0.0), + logit_bias: [(7, -1.25)].into_iter().collect(), + allowed_token_ids: vec![7, 8], + bad_words: vec!["blocked".to_string()], + structured_output: Some(pb::decoding_parameters::StructuredOutput::Regex( + "[a-z]+".to_string(), + )), + structured_output_disable_any_whitespace: true, + structured_output_disable_additional_properties: true, + structured_output_whitespace_pattern: Some("\\s*".to_string()), + }), + stopping: Some(pb::StoppingCriteria { + thinking_token_budget: Some(64), + ..Default::default() + }), + response: Some(pb::ResponseOptions { + output_logprobs: true, + output_candidates: Some(pb::CandidateTokens { + select: Some(pb::candidate_tokens::Select::TokenIds(pb::TokenIds { + ids: vec![7, 8], + })), + }), + ..Default::default() + }), + vllm_xargs_json: Some(br#"{"custom_integer":9007199254740993}"#.to_vec()), + ..base_request() + }; + + let text = to_text_request(req, false, &["test-model".to_string()]).unwrap(); + + assert_eq!(text.sampling_params.top_k, Some(0)); + assert_eq!(text.sampling_params.top_p, Some(0.0)); + assert_eq!(text.sampling_params.min_p, Some(0.0)); + assert_eq!(text.sampling_params.presence_penalty, Some(0.0)); + assert_eq!(text.sampling_params.frequency_penalty, Some(0.0)); + assert_eq!(text.sampling_params.repetition_penalty, Some(0.0)); + assert_eq!(text.sampling_params.thinking_token_budget, Some(64)); + assert_eq!(text.sampling_params.logit_bias.as_ref().unwrap()[&7], -1.25); + assert_eq!(text.sampling_params.allowed_token_ids, Some(vec![7, 8])); + assert_eq!( + text.sampling_params.bad_words, + Some(vec!["blocked".to_string()]) + ); + assert_eq!(text.sampling_params.logprobs, Some(2)); + assert_eq!(text.sampling_params.logprob_token_ids, Some(vec![7, 8])); + let structured = text.sampling_params.structured_outputs.unwrap(); + assert_eq!( + structured.constraint, + StructuredOutputConstraint::Regex("[a-z]+".to_string()) + ); + assert!(structured.options.disable_any_whitespace); + assert!(structured.options.disable_additional_properties); + assert_eq!( + structured.options.whitespace_pattern.as_deref(), + Some("\\s*") + ); + assert_eq!( + text.sampling_params.vllm_xargs.as_ref().unwrap()["custom_integer"], + serde_json::json!(9_007_199_254_740_993_u64) + ); + } + + #[test] + fn preprocessed_multimodal_features_reach_text_request() { + let kwargs = BTreeMap::from([( + "num_tiles".to_string(), + MmFieldElem { + data: Some(MmKwargValue::Int(2)), + field: MmField::Batched(MmBatchedField { keep_on_cpu: true }), + }, + )]); + let kwargs_msgpack = rmp_serde::to_vec_named(&kwargs).unwrap(); + let identifier = mm_cache_identifier("image", &kwargs_msgpack); + let req = pb::GenerateRequest { + prompt: Some(pb::generate_request::Prompt::TokenIds(pb::TokenIds { + ids: vec![10, 99, 99, 20], + })), + mm_features: vec![pb::PreprocessedMultimodalFeature { + modality: "image".to_string(), + mm_hash: identifier.clone(), + position: Some(pb::MultimodalPlaceholder { + offset: 1, + length: 2, + is_embed: vec![true, false], + }), + cache_identifier: identifier.clone(), + kwargs_msgpack: Some(kwargs_msgpack), + }], + ..base_request() + }; + + let text = to_text_request(req, false, &["test-model".to_string()]).unwrap(); + let features = text.mm_features.unwrap(); + + assert_eq!(features.len(), 1); + assert_eq!(features[0].modality, "image"); + assert!(features[0].identifier.starts_with("grpc-mm:")); + assert_eq!(features[0].mm_hash.as_deref(), Some(identifier.as_str())); + assert_eq!(features[0].data.as_ref(), Some(&kwargs)); + assert_eq!(features[0].mm_position.offset, 1); + assert_eq!(features[0].mm_position.length, 2); + assert!(features[0].mm_position.is_embed.is_some()); + } + + #[test] + fn multimodal_lora_is_rejected_until_tower_cache_semantics_are_advertised() { + for request in [ + pb::GenerateRequest { + lora_name: "adapter-a".to_string(), + media: vec![pb::MediaItem::default()], + ..base_request() + }, + pb::GenerateRequest { + lora_name: "adapter-a".to_string(), + mm_features: vec![pb::PreprocessedMultimodalFeature::default()], + ..base_request() + }, + ] { + let error = to_text_request(request, false, &["test-model".to_string()]).unwrap_err(); + assert_eq!(error.code(), tonic::Code::InvalidArgument); + assert!(error.message().contains("tower-LoRA")); + } + } + + #[test] + fn preprocessed_multimodal_cache_hit_without_data_is_rejected() { + let req = pb::GenerateRequest { + prompt: Some(pb::generate_request::Prompt::TokenIds(pb::TokenIds { + ids: vec![10, 99, 20], + })), + mm_features: vec![pb::PreprocessedMultimodalFeature { + modality: "image".to_string(), + mm_hash: "image-hash".to_string(), + position: Some(pb::MultimodalPlaceholder { + offset: 1, + length: 1, + is_embed: Vec::new(), + }), + kwargs_msgpack: None, + cache_identifier: String::new(), + }], + ..base_request() + }; + + let error = to_text_request(req, false, &["test-model".to_string()]).unwrap_err(); + assert_eq!(error.code(), tonic::Code::InvalidArgument); + } + + #[test] + fn multimodal_cache_identity_is_data_bound_without_language_lora_scoping() { + fn request(value: i64) -> pb::GenerateRequest { + let kwargs = BTreeMap::from([( + "value".to_string(), + MmFieldElem { + data: Some(MmKwargValue::Int(value)), + field: MmField::Batched(MmBatchedField { keep_on_cpu: true }), + }, + )]); + let kwargs_msgpack = rmp_serde::to_vec_named(&kwargs).unwrap(); + let identifier = mm_cache_identifier("image", &kwargs_msgpack); + pb::GenerateRequest { + prompt: Some(pb::generate_request::Prompt::TokenIds(pb::TokenIds { + ids: vec![10, 99, 20], + })), + mm_features: vec![pb::PreprocessedMultimodalFeature { + modality: "image".to_string(), + mm_hash: identifier.clone(), + position: Some(pb::MultimodalPlaceholder { + offset: 1, + length: 1, + is_embed: Vec::new(), + }), + cache_identifier: identifier, + kwargs_msgpack: Some(kwargs_msgpack), + }], + ..base_request() + } + } + + let first = to_text_request(request(1), false, &["test-model".to_string()]).unwrap(); + let second = to_text_request(request(2), false, &["test-model".to_string()]).unwrap(); + assert_ne!( + first.mm_features.as_ref().unwrap()[0].identifier, + second.mm_features.as_ref().unwrap()[0].identifier + ); + assert!(first.mm_features.as_ref().unwrap()[0].identifier.starts_with("grpc-mm:")); + } + + #[test] + fn multimodal_cache_identity_has_a_cross_language_fixed_vector() { + assert_eq!( + mm_cache_identifier("image", b"abc"), + "grpc-mm:c2f2df4bb94911d850921fa6d577ee0713ad5884c276f85330a8e50137f6a59d" + ); + } + + #[test] + fn multimodal_cache_identity_mismatch_is_rejected() { + let kwargs = BTreeMap::from([( + "value".to_string(), + MmFieldElem { + data: Some(MmKwargValue::Int(1)), + field: MmField::Batched(MmBatchedField { keep_on_cpu: true }), + }, + )]); + let kwargs_msgpack = rmp_serde::to_vec_named(&kwargs).unwrap(); + let req = pb::GenerateRequest { + prompt: Some(pb::generate_request::Prompt::TokenIds(pb::TokenIds { + ids: vec![10, 99, 20], + })), + mm_features: vec![pb::PreprocessedMultimodalFeature { + modality: "image".to_string(), + mm_hash: "image-hash".to_string(), + position: Some(pb::MultimodalPlaceholder { + offset: 1, + length: 1, + is_embed: Vec::new(), + }), + kwargs_msgpack: Some(kwargs_msgpack.clone()), + cache_identifier: mm_cache_identifier("image", &kwargs_msgpack), + }], + ..base_request() + }; + + let error = to_text_request(req, false, &["test-model".to_string()]).unwrap_err(); + assert_eq!(error.code(), tonic::Code::InvalidArgument); + } + + #[test] + fn vllm_xargs_reject_reserved_kv_control() { + let req = pb::GenerateRequest { + vllm_xargs_json: Some(br#"{"kv_transfer_params":{}}"#.to_vec()), + ..base_request() + }; + let error = to_text_request(req, false, &["test-model".to_string()]).unwrap_err(); + assert_eq!(error.code(), tonic::Code::InvalidArgument); + } + + #[test] + fn msgpack_preflight_rejects_trailing_and_amplified_values() { + let mut nodes = 0; + assert!(preflight_msgpack(&[0xc0, 0xc0], &mut nodes).is_err()); + + let mut nodes = 0; + let oversized_array = [0xdd, 0x00, 0x01, 0x00, 0x01]; + assert!(preflight_msgpack(&oversized_array, &mut nodes).is_err()); + } + + #[test] + fn malformed_preprocessed_multimodal_feature_is_rejected() { + let req = pb::GenerateRequest { + prompt: Some(pb::generate_request::Prompt::TokenIds(pb::TokenIds { + ids: vec![1, 2], + })), + mm_features: vec![pb::PreprocessedMultimodalFeature { + modality: "image".to_string(), + mm_hash: "image-hash".to_string(), + position: Some(pb::MultimodalPlaceholder { + offset: 1, + length: 2, + is_embed: Vec::new(), + }), + kwargs_msgpack: Some(vec![0xc1]), + cache_identifier: "invalid".to_string(), + }], + ..base_request() + }; + + let error = to_text_request(req, false, &["test-model".to_string()]).unwrap_err(); + assert_eq!(error.code(), tonic::Code::InvalidArgument); + } + fn finished(reason: FinishReason) -> Finished { Finished { usage: vllm_llm::TokenUsage { @@ -596,6 +883,7 @@ mod tests { finish_reason: reason, kv_transfer_params: None, ec_transfer_params: None, + routed_experts: None, } } @@ -679,4 +967,63 @@ mod tests { assert_eq!(finish.finish_reason, PbFinishReason::Stop as i32); assert_eq!(finish.stop_reason, Some(PbStopReason::EosTokenId(30))); } + + #[test] + fn to_sequence_output_emits_typed_routed_experts() { + let mut fin = finished(FinishReason::Length); + fin.routed_experts = Some(WireTensor::from_raw("|u1", vec![1, 2, 2], vec![1, 2, 3, 4])); + + let output = to_sequence_output("", &[], None, Some(&fin), &ResponseOpts::default()); + let routed = output.routed_experts.unwrap(); + assert_eq!(routed.dtype, "|u1"); + assert_eq!(routed.shape, vec![1, 2, 2]); + assert_eq!(routed.data, vec![1, 2, 3, 4]); + } + + #[test] + fn output_logprobs_do_not_repeat_the_sampled_token_as_a_candidate() { + let logprobs = DecodedLogprobs { + positions: vec![DecodedPositionLogprobs { + entries: vec![ + DecodedTokenLogprob { + token_id: 42, + token: "sampled".into(), + logprob: -0.1, + rank: 1, + }, + DecodedTokenLogprob { + token_id: 42, + token: "sampled".into(), + logprob: -0.1, + rank: 1, + }, + DecodedTokenLogprob { + token_id: 7, + token: "alternate".into(), + logprob: -1.2, + rank: 2, + }, + ], + }], + }; + let opts = ResponseOpts { + output_token_ids: true, + output_logprobs: true, + ..Default::default() + }; + + let output = to_sequence_output("sampled", &[42], Some(&logprobs), None, &opts); + + assert_eq!(output.logprobs, vec![-0.1]); + assert_eq!(output.ranks, vec![1]); + assert_eq!(output.candidate_tokens.len(), 1); + assert_eq!( + output.candidate_tokens[0] + .tokens + .iter() + .map(|candidate| candidate.id) + .collect::>(), + vec![7] + ); + } } diff --git a/rust/src/server/src/grpc/convert/multimodal.rs b/rust/src/server/src/grpc/convert/multimodal.rs new file mode 100644 index 000000000000..8b80239b35e2 --- /dev/null +++ b/rust/src/server/src/grpc/convert/multimodal.rs @@ -0,0 +1,540 @@ +use std::collections::HashMap; +use std::io::Cursor; + +use serde::Deserialize as _; +use sha2::{Digest, Sha256}; +use tonic::Status; +use vllm_chat::MediaContentPart; +use vllm_engine_core_client::protocol::multimodal::{ + MmFeatureSpec, MmField, MmKwargValue, MmKwargsItem, MmSlice, PlaceholderRange, +}; +use vllm_engine_core_client::protocol::tensor::{WireArrayData, WireTensor}; +use vllm_text::Prompt; + +use super::pb; + +const MAX_MM_FEATURE_BYTES: usize = 16 * 1024 * 1024; +const MAX_MM_FEATURES: usize = 64; +const MAX_MM_DEPTH: usize = 32; +const MAX_MM_NODES: usize = 65_536; +const MAX_MM_FIELDS_PER_ITEM: usize = 256; +const MAX_MM_KEY_BYTES: usize = 256; +const MAX_MM_HASH_BYTES: usize = 256; +const MAX_MM_TENSOR_RANK: usize = 32; + +pub(crate) fn media_parts_from_request( + media: &[pb::MediaItem], +) -> Result, Status> { + let mut parts = Vec::with_capacity(media.len()); + for item in media { + let modality = pb::Modality::try_from(item.modality).map_err(|_| { + Status::invalid_argument(format!("unknown media modality {}", item.modality)) + })?; + match modality { + pb::Modality::Image | pb::Modality::Unspecified => {} + other => { + return Err(Status::unimplemented(format!( + "media modality {other:?} is not supported by the gRPC service" + ))); + } + } + let uuid = (!item.uuid.is_empty()).then(|| item.uuid.clone()); + let part = match item.source.as_ref() { + Some(pb::media_item::Source::Url(url)) | Some(pb::media_item::Source::DataUri(url)) => { + MediaContentPart::ImageUrl { + url: url.clone(), + detail: None, + uuid, + } + } + Some(pb::media_item::Source::RawBytes(bytes)) => MediaContentPart::ImageData { + data: bytes.clone(), + mime_type: (!item.mime_type.is_empty()).then(|| item.mime_type.clone()), + uuid, + detail: None, + }, + None => return Err(Status::invalid_argument("media item has no source")), + }; + parts.push(part); + } + Ok(parts) +} + +pub(super) fn convert_mm_features( + features: &[pb::PreprocessedMultimodalFeature], + prompt: &Prompt, +) -> Result>, Status> { + if features.is_empty() { + return Ok(None); + } + let Prompt::TokenIds(token_ids) = prompt else { + return Err(Status::invalid_argument( + "preprocessed multimodal features require token_ids input", + )); + }; + if features.len() > MAX_MM_FEATURES || features.len() > token_ids.len() { + return Err(Status::resource_exhausted( + "too many preprocessed multimodal features", + )); + } + + let mut encoded_bytes = 0usize; + let mut wire_nodes = 0usize; + let mut converted = Vec::with_capacity(features.len()); + for feature in features { + if feature.modality.is_empty() || feature.modality.len() > 64 { + return Err(Status::invalid_argument( + "multimodal feature modality must contain between 1 and 64 bytes", + )); + } + if feature.mm_hash.is_empty() || feature.mm_hash.len() > MAX_MM_HASH_BYTES { + return Err(Status::invalid_argument( + "multimodal feature mm_hash must contain between 1 and 256 bytes", + )); + } + let position = feature + .position + .as_ref() + .ok_or_else(|| Status::invalid_argument("multimodal feature position is required"))?; + let offset = usize::try_from(position.offset) + .map_err(|_| Status::invalid_argument("multimodal feature offset is too large"))?; + let length = usize::try_from(position.length) + .map_err(|_| Status::invalid_argument("multimodal feature length is too large"))?; + if length == 0 { + return Err(Status::invalid_argument( + "multimodal feature length must be positive", + )); + } + let end = offset + .checked_add(length) + .ok_or_else(|| Status::invalid_argument("multimodal feature range overflows"))?; + if end > token_ids.len() { + return Err(Status::invalid_argument( + "multimodal feature range exceeds token_ids", + )); + } + if !position.is_embed.is_empty() && position.is_embed.len() != length { + return Err(Status::invalid_argument( + "multimodal feature is_embed length must match position length", + )); + } + let is_embed = if position.is_embed.is_empty() { + None + } else { + Some( + WireTensor::from_bool(vec![length], position.is_embed.clone()) + .map_err(Status::invalid_argument)?, + ) + }; + + let raw = feature.kwargs_msgpack.as_deref().ok_or_else(|| { + Status::invalid_argument( + "multimodal feature kwargs_msgpack is required; unverified cache hits are unsupported", + ) + })?; + encoded_bytes = encoded_bytes + .checked_add(raw.len()) + .ok_or_else(|| Status::resource_exhausted("multimodal feature payload is too large"))?; + if encoded_bytes > MAX_MM_FEATURE_BYTES { + return Err(Status::resource_exhausted( + "multimodal feature payload exceeds 16 MiB", + )); + } + let item = decode_mm_kwargs(raw, &mut wire_nodes)?; + validate_mm_kwargs_item(&item)?; + let identifier = mm_cache_identifier(&feature.modality, raw); + if feature.cache_identifier != identifier { + return Err(Status::invalid_argument( + "multimodal feature cache_identifier does not match its canonical payload identity", + )); + } + if feature.mm_hash != identifier { + return Err(Status::invalid_argument( + "multimodal feature mm_hash does not match its canonical payload identity", + )); + } + + converted.push(MmFeatureSpec { + data: Some(item), + modality: feature.modality.clone(), + identifier: identifier.clone(), + mm_position: PlaceholderRange { + offset, + length, + is_embed, + }, + mm_hash: Some(identifier), + }); + } + converted.sort_by_key(|feature| feature.mm_position.offset); + for pair in converted.windows(2) { + let previous_end = pair[0] + .mm_position + .offset + .checked_add(pair[0].mm_position.length) + .ok_or_else(|| Status::invalid_argument("multimodal feature range overflows"))?; + if previous_end > pair[1].mm_position.offset { + return Err(Status::invalid_argument( + "multimodal feature ranges must not overlap", + )); + } + } + validate_mm_field_metadata(&converted)?; + Ok(Some(converted)) +} + +fn validate_mm_kwargs_item(item: &MmKwargsItem) -> Result<(), Status> { + if item.is_empty() || item.len() > MAX_MM_FIELDS_PER_ITEM { + return Err(Status::invalid_argument( + "multimodal kwargs item must contain between 1 and 256 fields", + )); + } + for (key, element) in item { + if key.is_empty() || key.len() > MAX_MM_KEY_BYTES { + return Err(Status::invalid_argument( + "multimodal kwargs keys must contain between 1 and 256 bytes", + )); + } + let value = element.data.as_ref().ok_or_else(|| { + Status::invalid_argument("multimodal kwargs fields must carry inline data") + })?; + validate_mm_kwarg_value(value, 0)?; + } + Ok(()) +} + +fn validate_mm_kwarg_value(value: &MmKwargValue, depth: usize) -> Result<(), Status> { + if depth > 32 { + return Err(Status::invalid_argument( + "multimodal kwargs nesting exceeds 32 levels", + )); + } + match value { + MmKwargValue::Tensor(tensor) => validate_wire_tensor(tensor), + MmKwargValue::List(values) => { + for value in values { + validate_mm_kwarg_value(value, depth + 1)?; + } + Ok(()) + } + MmKwargValue::Int(_) | MmKwargValue::Float(_) => Ok(()), + } +} + +fn validate_wire_tensor(tensor: &WireTensor) -> Result<(), Status> { + if tensor.shape.len() > MAX_MM_TENSOR_RANK { + return Err(Status::invalid_argument( + "multimodal tensor rank exceeds 32", + )); + } + let width = match tensor.dtype.as_str() { + "bool" | "uint8" | "int8" => 1, + "float16" | "bfloat16" | "uint16" | "int16" => 2, + "float32" | "uint32" | "int32" => 4, + "float64" | "uint64" | "int64" => 8, + dtype => { + return Err(Status::invalid_argument(format!( + "unsupported multimodal tensor dtype {dtype:?}" + ))); + } + }; + let numel = tensor + .shape + .iter() + .try_fold(1usize, |count, dim| count.checked_mul(*dim)) + .ok_or_else(|| Status::invalid_argument("multimodal tensor shape overflows"))?; + let expected = numel + .checked_mul(width) + .ok_or_else(|| Status::invalid_argument("multimodal tensor byte length overflows"))?; + match &tensor.data { + WireArrayData::RawView(bytes) if bytes.len() == expected => Ok(()), + WireArrayData::RawView(bytes) => Err(Status::invalid_argument(format!( + "multimodal tensor byte length {} does not match expected {expected}", + bytes.len() + ))), + WireArrayData::AuxIndex(_) => Err(Status::invalid_argument( + "multimodal kwargs must encode tensors inline", + )), + } +} + +pub(super) fn mm_cache_identifier(modality: &str, raw: &[u8]) -> String { + let mut hasher = Sha256::new(); + hasher.update(b"vllm.grpc.preprocessed-mm.v1"); + hasher.update((modality.len() as u64).to_be_bytes()); + hasher.update(modality.as_bytes()); + hasher.update((raw.len() as u64).to_be_bytes()); + hasher.update(raw); + format!("grpc-mm:{:x}", hasher.finalize()) +} + +fn decode_mm_kwargs(raw: &[u8], nodes: &mut usize) -> Result { + preflight_msgpack(raw, nodes)?; + let mut deserializer = rmp_serde::Deserializer::new(Cursor::new(raw)); + deserializer.set_max_depth(MAX_MM_DEPTH); + let item = MmKwargsItem::deserialize(&mut deserializer).map_err(|error| { + Status::invalid_argument(format!("invalid multimodal kwargs msgpack: {error}")) + })?; + if deserializer.position() != raw.len() as u64 { + return Err(Status::invalid_argument( + "multimodal kwargs msgpack contains trailing data", + )); + } + Ok(item) +} + +pub(super) fn preflight_msgpack(raw: &[u8], nodes: &mut usize) -> Result<(), Status> { + let mut cursor = 0usize; + scan_msgpack_value(raw, &mut cursor, 0, nodes)?; + if cursor != raw.len() { + return Err(Status::invalid_argument( + "multimodal kwargs msgpack contains trailing data", + )); + } + Ok(()) +} + +fn scan_msgpack_value( + raw: &[u8], + cursor: &mut usize, + depth: usize, + nodes: &mut usize, +) -> Result<(), Status> { + if depth > MAX_MM_DEPTH { + return Err(Status::resource_exhausted( + "multimodal kwargs nesting exceeds 32 levels", + )); + } + *nodes = nodes + .checked_add(1) + .ok_or_else(|| Status::resource_exhausted("multimodal kwargs are too complex"))?; + if *nodes > MAX_MM_NODES { + return Err(Status::resource_exhausted( + "multimodal kwargs contain too many values", + )); + } + let marker = take_msgpack(raw, cursor, 1)?[0]; + match marker { + 0x00..=0x7f | 0xc0 | 0xc2 | 0xc3 | 0xe0..=0xff => Ok(()), + 0x80..=0x8f => { + scan_msgpack_children(raw, cursor, depth, nodes, (marker & 0x0f) as usize * 2) + } + 0x90..=0x9f => scan_msgpack_children(raw, cursor, depth, nodes, (marker & 0x0f) as usize), + 0xa0..=0xbf => skip_msgpack(raw, cursor, (marker & 0x1f) as usize), + 0xc1 => Err(Status::invalid_argument("reserved MessagePack marker")), + 0xc4 => { + let len = read_msgpack_u8(raw, cursor)? as usize; + skip_msgpack(raw, cursor, len) + } + 0xc5 => { + let len = read_msgpack_u16(raw, cursor)? as usize; + skip_msgpack(raw, cursor, len) + } + 0xc6 => { + let len = read_msgpack_u32(raw, cursor)? as usize; + skip_msgpack(raw, cursor, len) + } + 0xc7 => { + let len = read_msgpack_u8(raw, cursor)? as usize; + skip_msgpack(raw, cursor, len + 1) + } + 0xc8 => { + let len = read_msgpack_u16(raw, cursor)? as usize; + skip_msgpack(raw, cursor, len + 1) + } + 0xc9 => { + let len = read_msgpack_u32(raw, cursor)? as usize; + skip_msgpack(raw, cursor, len + 1) + } + 0xca => skip_msgpack(raw, cursor, 4), + 0xcb => skip_msgpack(raw, cursor, 8), + 0xcc | 0xd0 => skip_msgpack(raw, cursor, 1), + 0xcd | 0xd1 => skip_msgpack(raw, cursor, 2), + 0xce | 0xd2 => skip_msgpack(raw, cursor, 4), + 0xcf | 0xd3 => skip_msgpack(raw, cursor, 8), + 0xd4 => skip_msgpack(raw, cursor, 2), + 0xd5 => skip_msgpack(raw, cursor, 3), + 0xd6 => skip_msgpack(raw, cursor, 5), + 0xd7 => skip_msgpack(raw, cursor, 9), + 0xd8 => skip_msgpack(raw, cursor, 17), + 0xd9 => { + let len = read_msgpack_u8(raw, cursor)? as usize; + skip_msgpack(raw, cursor, len) + } + 0xda => { + let len = read_msgpack_u16(raw, cursor)? as usize; + skip_msgpack(raw, cursor, len) + } + 0xdb => { + let len = read_msgpack_u32(raw, cursor)? as usize; + skip_msgpack(raw, cursor, len) + } + 0xdc => { + let count = read_msgpack_u16(raw, cursor)? as usize; + scan_msgpack_children(raw, cursor, depth, nodes, count) + } + 0xdd => { + let count = read_msgpack_u32(raw, cursor)? as usize; + scan_msgpack_children(raw, cursor, depth, nodes, count) + } + 0xde => { + let count = (read_msgpack_u16(raw, cursor)? as usize) + .checked_mul(2) + .ok_or_else(|| Status::resource_exhausted("MessagePack map is too large"))?; + scan_msgpack_children(raw, cursor, depth, nodes, count) + } + 0xdf => { + let count = (read_msgpack_u32(raw, cursor)? as usize) + .checked_mul(2) + .ok_or_else(|| Status::resource_exhausted("MessagePack map is too large"))?; + scan_msgpack_children(raw, cursor, depth, nodes, count) + } + } +} + +fn scan_msgpack_children( + raw: &[u8], + cursor: &mut usize, + depth: usize, + nodes: &mut usize, + count: usize, +) -> Result<(), Status> { + if count > MAX_MM_NODES.saturating_sub(*nodes) { + return Err(Status::resource_exhausted( + "multimodal kwargs contain too many values", + )); + } + for _ in 0..count { + scan_msgpack_value(raw, cursor, depth + 1, nodes)?; + } + Ok(()) +} + +fn take_msgpack<'a>(raw: &'a [u8], cursor: &mut usize, len: usize) -> Result<&'a [u8], Status> { + let end = cursor + .checked_add(len) + .filter(|end| *end <= raw.len()) + .ok_or_else(|| Status::invalid_argument("truncated multimodal kwargs msgpack"))?; + let bytes = &raw[*cursor..end]; + *cursor = end; + Ok(bytes) +} + +fn skip_msgpack(raw: &[u8], cursor: &mut usize, len: usize) -> Result<(), Status> { + take_msgpack(raw, cursor, len).map(|_| ()) +} + +fn read_msgpack_u8(raw: &[u8], cursor: &mut usize) -> Result { + Ok(take_msgpack(raw, cursor, 1)?[0]) +} + +fn read_msgpack_u16(raw: &[u8], cursor: &mut usize) -> Result { + Ok(u16::from_be_bytes( + take_msgpack(raw, cursor, 2)?.try_into().expect("fixed two-byte slice"), + )) +} + +fn read_msgpack_u32(raw: &[u8], cursor: &mut usize) -> Result { + Ok(u32::from_be_bytes( + take_msgpack(raw, cursor, 4)?.try_into().expect("fixed four-byte slice"), + )) +} + +fn validate_mm_field_metadata(features: &[MmFeatureSpec]) -> Result<(), Status> { + let mut occurrences: HashMap<(String, String), usize> = HashMap::new(); + let mut fields: HashMap<(String, String), MmField> = HashMap::new(); + for feature in features { + let item = feature.data.as_ref().expect("inline multimodal data is required above"); + for (key, element) in item { + let identity = (feature.modality.clone(), key.clone()); + *occurrences.entry(identity.clone()).or_default() += 1; + if let Some(previous) = fields.get(&identity) { + if previous != &element.field { + return Err(Status::invalid_argument( + "multimodal field configuration differs across items", + )); + } + } else { + fields.insert(identity, element.field.clone()); + } + } + } + for feature in features { + let item = feature.data.as_ref().expect("inline multimodal data is required above"); + for (key, element) in item { + let count = occurrences[&(feature.modality.clone(), key.clone())]; + validate_mm_field( + &element.field, + element.data.as_ref().expect("inline multimodal field data is required above"), + count, + )?; + } + } + Ok(()) +} + +fn validate_mm_field( + field: &MmField, + data: &MmKwargValue, + occurrences: usize, +) -> Result<(), Status> { + match field { + MmField::Batched(_) => Ok(()), + MmField::Shared(shared) => { + if shared.batch_size == 0 || shared.batch_size != occurrences { + return Err(Status::invalid_argument( + "multimodal shared-field batch_size must match item count", + )); + } + Ok(()) + } + MmField::Flat(flat) => { + if flat.slices.is_empty() || flat.slices.len() != occurrences { + return Err(Status::invalid_argument( + "multimodal flat-field slices must match item count", + )); + } + for slice in &flat.slices { + match slice { + MmSlice::Slice(slice) => validate_slice_step(slice.step)?, + MmSlice::Slices(slices) => { + if slices.is_empty() || slices.len() > MAX_MM_TENSOR_RANK { + return Err(Status::invalid_argument( + "multimodal flat-field slice tuple must contain 1 to 32 slices", + )); + } + for slice in slices { + validate_slice_step(slice.step)?; + } + } + } + } + match data { + MmKwargValue::Tensor(tensor) => { + let rank = i32::try_from(tensor.shape.len()).unwrap_or(i32::MAX); + if rank == 0 || flat.dim < -rank || flat.dim >= rank { + return Err(Status::invalid_argument( + "multimodal flat-field dim is outside the tensor rank", + )); + } + } + _ if flat.dim != 0 => { + return Err(Status::invalid_argument( + "multimodal non-tensor flat fields require dim=0", + )); + } + _ => {} + } + Ok(()) + } + } +} + +fn validate_slice_step(step: Option) -> Result<(), Status> { + if step == Some(0) { + return Err(Status::invalid_argument( + "multimodal slice step must not be zero", + )); + } + Ok(()) +} diff --git a/rust/src/server/src/grpc/convert/sampling.rs b/rust/src/server/src/grpc/convert/sampling.rs new file mode 100644 index 000000000000..94a29d9e835a --- /dev/null +++ b/rust/src/server/src/grpc/convert/sampling.rs @@ -0,0 +1,163 @@ +use tonic::Status; +use vllm_engine_core_client::protocol::structured_outputs::{ + StructuredOutputOptions, StructuredOutputsParams, +}; +use vllm_text::SamplingParams; + +use super::pb; + +pub(super) fn build_sampling_params( + temperature: Option, + sampling: Option<&pb::RandomSampling>, + decoding: Option<&pb::DecodingParameters>, + stopping: Option<&pb::StoppingCriteria>, + response: Option<&pb::ResponseOptions>, +) -> Result { + // Temperature is a top-level GenerateRequest field. Default to greedy (0.0) for + // the gRPC API when the caller does not specify a value. This differs from + // the HTTP/OpenAI API (which defaults to 1.0) and matches the convention of + // programmatic generation APIs. + let temperature = temperature.or(Some(0.0)); + let mut params = SamplingParams { + temperature, + ..SamplingParams::default() + }; + + // Optional scalar presence preserves explicit zero and sentinel values; + // omitted fields remain available for model defaults during lowering. + if let Some(s) = sampling { + // num_sequences (n > 1) is not supported yet by the TextLlm layer; the response + // path also hardcodes SequenceOutput.index = 0, so accepting >1 would silently + // truncate output cardinality. Reject explicitly. + if s.num_sequences > 1 { + return Err(Status::invalid_argument( + "num_sequences > 1 is not supported", + )); + } + params.top_k = s + .top_k + .map(|value| match value { + -1 => Ok(0), + 0.. => u32::try_from(value) + .map_err(|_| Status::invalid_argument("top_k exceeds uint32")), + _ => Err(Status::invalid_argument("top_k must be at least -1")), + }) + .transpose()?; + params.top_p = s.top_p; + params.min_p = s.min_p; + params.seed = s.seed; + } + + // DecodingParameters + if let Some(d) = decoding { + params.presence_penalty = d.presence_penalty; + params.frequency_penalty = d.frequency_penalty; + params.repetition_penalty = d.repetition_penalty; + if !d.logit_bias.is_empty() { + params.logit_bias = Some(d.logit_bias.clone()); + } + if !d.allowed_token_ids.is_empty() { + params.allowed_token_ids = Some(d.allowed_token_ids.clone()); + } + if !d.bad_words.is_empty() { + params.bad_words = Some(d.bad_words.clone()); + } + params.structured_outputs = convert_structured_output(d)?; + } + + // StoppingCriteria + if let Some(s) = stopping { + if s.max_new_tokens != 0 { + params.max_tokens = Some(s.max_new_tokens); + } + if s.min_new_tokens != 0 { + params.min_tokens = Some(s.min_new_tokens); + } + if !s.stop_token_ids.is_empty() { + params.stop_token_ids = Some(s.stop_token_ids.clone()); + } + params.ignore_eos = s.ignore_eos; + params.thinking_token_budget = s.thinking_token_budget; + } + + // ResponseOptions → logprobs + if let Some(r) = response { + if r.output_logprobs { + let (count, token_ids) = candidate_logprob_spec(r.output_candidates.as_ref()); + params.logprobs = Some(count); + params.logprob_token_ids = token_ids; + } + if r.prompt_logprobs { + // The engine-core protocol has only one shared `logprob_token_ids` field + // for output and prompt logprobs, so a per-token-id selector for prompt + // candidates can't be honored independently. Reject it instead of silently + // dropping the list. + if matches!( + r.prompt_candidates.as_ref().and_then(|c| c.select.as_ref()), + Some(pb::candidate_tokens::Select::TokenIds(_)) + ) { + return Err(Status::invalid_argument( + "prompt_candidates token_ids selector is not supported", + )); + } + let (count, _) = candidate_logprob_spec(r.prompt_candidates.as_ref()); + params.prompt_logprobs = Some(count); + } + } + + Ok(params) +} + +/// Map the proto `CandidateTokens` selector to a `(logprobs_count, +/// logprob_token_ids)` pair. +/// +/// - `top_n(k)` → `(k, None)` — return top-k candidates by probability +/// - `all` → `(-1, None)` — return the full vocabulary +/// - `token_ids(n)` → `(1, Some(vec of n token ids))` — return logprobs for specific tokens (the +/// count `n` is stored in the proto as the number of token IDs that follow, but the actual IDs +/// are carried via `logprob_token_ids` on `SamplingParams`) +/// - absent → `(1, None)` — just the sampled/scored token +fn candidate_logprob_spec(candidates: Option<&pb::CandidateTokens>) -> (i32, Option>) { + match candidates.and_then(|c| c.select.as_ref()) { + Some(pb::candidate_tokens::Select::TopN(n)) => (*n as i32, None), + Some(pb::candidate_tokens::Select::All(true)) => (-1, None), + Some(pb::candidate_tokens::Select::TokenIds(ids)) => ( + ids.ids.len().try_into().unwrap_or(i32::MAX), + Some(ids.ids.clone()), + ), + _ => (1, None), + } +} + +fn convert_structured_output( + d: &pb::DecodingParameters, +) -> Result, Status> { + let so = match d.structured_output.as_ref() { + None => return Ok(None), + Some(so) => so, + }; + use pb::decoding_parameters::StructuredOutput; + let mut params = match so { + StructuredOutput::Json(schema) => { + let json: serde_json::Value = serde_json::from_str(schema) + .map_err(|e| Status::invalid_argument(format!("invalid json schema: {e}")))?; + StructuredOutputsParams::json(json) + } + StructuredOutput::Regex(regex) => StructuredOutputsParams::regex(regex.clone()), + StructuredOutput::Choice(choices) => { + StructuredOutputsParams::choice(choices.choices.clone()) + } + StructuredOutput::Grammar(grammar) => StructuredOutputsParams::grammar(grammar.clone()), + StructuredOutput::JsonObject(true) => StructuredOutputsParams::json_object(), + StructuredOutput::JsonObject(false) => return Ok(None), + StructuredOutput::StructuralTag(tag) => { + StructuredOutputsParams::structural_tag(tag.clone()) + } + }; + params.options = StructuredOutputOptions { + disable_any_whitespace: d.structured_output_disable_any_whitespace, + disable_additional_properties: d.structured_output_disable_additional_properties, + whitespace_pattern: d.structured_output_whitespace_pattern.clone(), + }; + Ok(Some(params)) +} diff --git a/rust/src/server/src/grpc/convert/xargs.rs b/rust/src/server/src/grpc/convert/xargs.rs new file mode 100644 index 000000000000..7c1f126c2fa7 --- /dev/null +++ b/rust/src/server/src/grpc/convert/xargs.rs @@ -0,0 +1,77 @@ +use std::collections::HashMap; + +use tonic::Status; + +const MAX_VLLM_XARGS_BYTES: usize = 64 * 1024; +const MAX_VLLM_XARGS_KEYS: usize = 64; +const MAX_VLLM_XARGS_DEPTH: usize = 16; +const MAX_VLLM_XARGS_NODES: usize = 1024; +const MAX_KEY_BYTES: usize = 256; + +pub(super) fn parse_vllm_xargs_json( + raw: &[u8], +) -> Result, Status> { + if raw.len() > MAX_VLLM_XARGS_BYTES { + return Err(Status::resource_exhausted("vllm_xargs_json exceeds 64 KiB")); + } + let value: serde_json::Value = serde_json::from_slice(raw).map_err(|error| { + Status::invalid_argument(format!("vllm_xargs_json must be a JSON object: {error}")) + })?; + let object = value + .as_object() + .ok_or_else(|| Status::invalid_argument("vllm_xargs_json must be a JSON object"))?; + if object.len() > MAX_VLLM_XARGS_KEYS { + return Err(Status::resource_exhausted( + "vllm_xargs_json contains too many keys", + )); + } + if object.contains_key("kv_transfer_params") { + return Err(Status::invalid_argument( + "vllm_xargs_json key kv_transfer_params is reserved for typed KV parameters", + )); + } + let mut nodes = 0usize; + validate_json_budget(&value, 0, &mut nodes)?; + serde_json::from_value(value).map_err(|error| { + Status::invalid_argument(format!("vllm_xargs_json must be a JSON object: {error}")) + }) +} + +fn validate_json_budget( + value: &serde_json::Value, + depth: usize, + nodes: &mut usize, +) -> Result<(), Status> { + if depth > MAX_VLLM_XARGS_DEPTH { + return Err(Status::resource_exhausted( + "vllm_xargs_json nesting exceeds 16 levels", + )); + } + *nodes = nodes + .checked_add(1) + .ok_or_else(|| Status::resource_exhausted("vllm_xargs_json is too complex"))?; + if *nodes > MAX_VLLM_XARGS_NODES { + return Err(Status::resource_exhausted( + "vllm_xargs_json contains too many values", + )); + } + match value { + serde_json::Value::Array(values) => { + for value in values { + validate_json_budget(value, depth + 1, nodes)?; + } + } + serde_json::Value::Object(values) => { + for (key, value) in values { + if key.len() > MAX_KEY_BYTES { + return Err(Status::resource_exhausted( + "vllm_xargs_json key exceeds 256 bytes", + )); + } + validate_json_budget(value, depth + 1, nodes)?; + } + } + _ => {} + } + Ok(()) +} diff --git a/rust/src/server/src/grpc/inference.rs b/rust/src/server/src/grpc/inference.rs index 56fa40924cf1..3dca3b7fae7b 100644 --- a/rust/src/server/src/grpc/inference.rs +++ b/rust/src/server/src/grpc/inference.rs @@ -10,10 +10,10 @@ use tokio::sync::mpsc; use tokio_stream::wrappers::ReceiverStream; use tonic::{Request, Response, Status}; use tracing::info; -use vllm_text::{DecodedTextEvent, TextOutputStreamExt as _}; +use vllm_text::{DecodedTextEvent, Prompt, TextOutputStreamExt as _, TextRequest}; use super::convert::{self, ResponseOpts}; -use super::{InferenceServer, pb}; +use super::{AdmissionState, InferenceServer, pb}; use crate::state::AppState; pub(crate) type InferenceGrpcService = InferenceServer; @@ -21,11 +21,73 @@ pub(crate) type InferenceGrpcService = InferenceServer; /// gRPC inference service backed by the shared application state. pub struct InferenceServiceImpl { state: Arc, + admission: Arc, } impl InferenceServiceImpl { pub fn new(state: Arc) -> Self { - Self { state } + Self::with_admission(state, Arc::new(AdmissionState::default())) + } + + pub(crate) fn with_admission(state: Arc, admission: Arc) -> Self { + Self { state, admission } + } + + async fn prepare_request( + &self, + proto_request: pb::GenerateRequest, + stream: bool, + ) -> Result<(TextRequest, crate::lora::LoraLease), Status> { + let ready = self.state.engine_core_client().ready_response(); + let role = convert::role_from_kv_role(ready.kv_role.as_deref()); + convert::validate_disaggregated_request(&proto_request, role)?; + let supports_lora = ready.supports_lora; + let media = convert::media_parts_from_request(&proto_request.media)?; + let lora_name = proto_request.lora_name.clone(); + let mut text_request = + convert::to_text_request(proto_request, stream, self.state.served_model_names())?; + + let mut lora_lease = None; + if !lora_name.is_empty() { + if !supports_lora { + return Err(Status::failed_precondition( + "engine was not started with LoRA enabled", + )); + } + let mut resolution = self.state.resolve_model_with_loras(Some(&lora_name)).await; + lora_lease = resolution.lease.take(); + if !self.state.lora_state_is_consistent() { + return Err(Status::failed_precondition( + "LoRA state differs across engine ranks; restart the engine", + )); + } + text_request.lora_request = Some(resolution.lora_request.ok_or_else(|| { + Status::not_found(format!("LoRA adapter `{lora_name}` is not loaded")) + })?); + } + // Language-only LoRA does not change preprocessed multimodal features. + // Tower or connector adapters need an explicit cache-identity contract. + if !media.is_empty() { + let Prompt::TokenIds(mut token_ids) = text_request.prompt else { + return Err(Status::invalid_argument( + "multimodal gRPC requests must provide token_ids input", + )); + }; + let mm_features = self + .state + .chat + .prepare_media(media, &mut token_ids) + .await + .map_err(|error| Status::internal(error.to_report_string()))?; + text_request.prompt = Prompt::TokenIds(token_ids); + text_request.mm_features = mm_features; + } + + if role == convert::KvRole::Prefill { + convert::mark_prefill_request(&mut text_request); + } + + Ok((text_request, lora_lease)) } } @@ -39,16 +101,20 @@ impl pb::inference_server::Inference for InferenceServiceImpl { &self, request: Request, ) -> Result, Status> { + let _guard = self + .admission + .try_admit() + .ok_or_else(|| Status::unavailable("gRPC service is draining"))?; let proto_req = request.into_inner(); let response_opts = ResponseOpts::from_proto(proto_req.response.as_ref()); - let text_request = - convert::to_text_request(proto_req, false, self.state.served_model_names())?; + let (text_request, lora_lease) = self.prepare_request(proto_req, false).await?; let request_id = text_request.request_id.clone(); info!(%request_id, "grpc generate (unary)"); let stream = self.state.chat.text().generate(text_request).await; let stream = stream.map_err(text_error_to_status)?; + let stream = crate::lora::hold_lora_lease(stream, lora_lease); let collected = stream.collect_output().await.map_err(text_error_to_status)?; @@ -64,6 +130,7 @@ impl pb::inference_server::Inference for InferenceServiceImpl { finish_reason: collected.finish_reason, kv_transfer_params: collected.kv_transfer_params, ec_transfer_params: collected.ec_transfer_params, + routed_experts: collected.routed_experts, }; let outputs = convert::to_sequence_output( @@ -85,20 +152,25 @@ impl pb::inference_server::Inference for InferenceServiceImpl { &self, request: Request, ) -> Result, Status> { + let guard = self + .admission + .try_admit() + .ok_or_else(|| Status::unavailable("gRPC service is draining"))?; let proto_req = request.into_inner(); let response_opts = ResponseOpts::from_proto(proto_req.response.as_ref()); - let text_request = - convert::to_text_request(proto_req, true, self.state.served_model_names())?; + let (text_request, lora_lease) = self.prepare_request(proto_req, true).await?; let request_id = text_request.request_id.clone(); info!(%request_id, "grpc generate (stream)"); let stream = self.state.chat.text().generate(text_request).await; let stream = stream.map_err(text_error_to_status)?; + let stream = crate::lora::hold_lora_lease(stream, lora_lease); let (tx, rx) = mpsc::channel(32); tokio::spawn(async move { + let _guard = guard; futures::pin_mut!(stream); while let Some(event) = stream.next().await { let response = match event { diff --git a/rust/src/server/src/grpc/lora_rpc.rs b/rust/src/server/src/grpc/lora_rpc.rs new file mode 100644 index 000000000000..2a5424fe3cb2 --- /dev/null +++ b/rust/src/server/src/grpc/lora_rpc.rs @@ -0,0 +1,172 @@ +use std::path::Path; + +use thiserror_ext::AsReport as _; +use tonic::{Request, Response, Status}; +use vllm_engine_core_client::protocol::lora::LoraRequest; + +use super::pb; +use crate::lora::{LoadExactLoraError, UnloadLoraError}; +use crate::lora_path::validate_lora_path_access; +use crate::state::AppState; + +fn ensure_enabled(state: &AppState) -> Result<(), Status> { + state + .engine_core_client() + .ready_response() + .supports_lora + .then_some(()) + .ok_or_else(|| Status::failed_precondition("engine was not started with LoRA enabled")) +} + +fn ensure_consistent(state: &AppState) -> Result<(), Status> { + state.lora_state_is_consistent().then_some(()).ok_or_else(|| { + Status::failed_precondition("LoRA state differs across engine ranks; restart the engine") + }) +} + +pub(super) async fn load_lora( + state: &std::sync::Arc, + allowed_path_prefixes: Option<&[std::path::PathBuf]>, + request: Request, +) -> Result, Status> { + ensure_enabled(state)?; + ensure_consistent(state)?; + let request = request.into_inner(); + let load_inplace = request.load_inplace; + let adapter = normalize_adapter( + request.adapter.ok_or_else(|| Status::invalid_argument("adapter is required"))?, + allowed_path_prefixes, + ) + .await?; + let (adapter, already_loaded) = + state + .load_lora_exact(adapter, load_inplace) + .await + .map_err(|error| match error { + LoadExactLoraError::Inconsistent => Status::failed_precondition( + "LoRA state differs across engine ranks; restart the engine", + ), + LoadExactLoraError::BaseModelName { lora_name } => Status::already_exists(format!( + "LoRA adapter `{lora_name}` conflicts with a served base model" + )), + LoadExactLoraError::Conflict { existing } => conflict(&existing), + LoadExactLoraError::Engine(error) => Status::internal(error.to_report_string()), + LoadExactLoraError::NotLoaded { lora_name } => Status::internal(format!( + "one or more engine ranks rejected LoRA adapter `{lora_name}`" + )), + })?; + Ok(Response::new(pb::LoadLoraResponse { + adapter: Some(to_proto(&adapter)), + already_loaded, + })) +} + +pub(super) async fn unload_lora( + state: &std::sync::Arc, + request: Request, +) -> Result, Status> { + ensure_enabled(state)?; + ensure_consistent(state)?; + let name = request.into_inner().lora_name; + if name.trim().is_empty() { + return Err(Status::invalid_argument("lora_name is required")); + } + + let adapter = state + .served_lora_requests() + .await + .into_iter() + .find(|adapter| adapter.lora_name == name) + .ok_or_else(|| Status::not_found(format!("LoRA adapter `{name}` is not loaded")))?; + let adapter = + state + .unload_lora(&name, Some(adapter.lora_int_id)) + .await + .map_err(|error| match error { + UnloadLoraError::Inconsistent => Status::failed_precondition( + "LoRA state differs across engine ranks; restart the engine", + ), + UnloadLoraError::NotFound { lora_name } => { + Status::not_found(format!("LoRA adapter `{lora_name}` is not loaded")) + } + UnloadLoraError::IntIdMismatch { .. } => { + Status::internal("LoRA registry changed during unload") + } + UnloadLoraError::Engine(error) => Status::internal(error.to_report_string()), + })?; + Ok(Response::new(pb::UnloadLoraResponse { + adapter: Some(to_proto(&adapter)), + })) +} + +pub(super) async fn list_loras( + state: &std::sync::Arc, + _request: Request, +) -> Result, Status> { + ensure_enabled(state)?; + ensure_consistent(state)?; + let mut adapters = state.served_lora_requests().await; + adapters.sort_by(|left, right| left.lora_name.cmp(&right.lora_name)); + let adapters = adapters.iter().map(to_proto).collect(); + Ok(Response::new(pb::ListLorasResponse { adapters })) +} + +async fn normalize_adapter( + adapter: pb::LoraAdapter, + allowed_path_prefixes: Option<&[std::path::PathBuf]>, +) -> Result { + if adapter.lora_id <= 0 { + return Err(Status::invalid_argument("lora_id must be positive")); + } + if adapter.lora_name.trim().is_empty() { + return Err(Status::invalid_argument("lora_name is required")); + } + let path = Path::new(&adapter.source_path); + if !path.is_absolute() { + return Err(Status::invalid_argument("source_path must be absolute")); + } + let canonical = validate_lora_path_access(&adapter.source_path, allowed_path_prefixes) + .await + .map_err(|error| { + if error.is_client_error() { + Status::invalid_argument(error.public_message()) + } else { + tracing::error!(error = %error, "runtime LoRA path policy validation failed"); + Status::internal(error.public_message()) + } + })? + .ok_or_else(|| Status::invalid_argument("source_path must be a local path"))?; + let metadata = tokio::fs::metadata(&canonical) + .await + .map_err(|error| Status::invalid_argument(format!("invalid source_path: {error}")))?; + if !metadata.is_dir() { + return Err(Status::invalid_argument("source_path must be a directory")); + } + + let request = LoraRequest { + lora_name: adapter.lora_name, + lora_int_id: u64::try_from(adapter.lora_id) + .map_err(|_| Status::invalid_argument("lora_id must be positive"))?, + lora_path: canonical.to_string_lossy().into_owned(), + base_model_name: None, + tensorizer_config_dict: None, + load_inplace: false, + is_3d_lora_weight: false, + }; + Ok(request) +} + +fn to_proto(adapter: &LoraRequest) -> pb::LoraAdapter { + pb::LoraAdapter { + lora_id: adapter.lora_int_id.min(i64::MAX as u64) as i64, + lora_name: adapter.lora_name.clone(), + source_path: adapter.lora_path.clone(), + } +} + +fn conflict(existing: &LoraRequest) -> Status { + Status::already_exists(format!( + "conflicts with loaded LoRA `{}` (id {})", + existing.lora_name, existing.lora_int_id + )) +} diff --git a/rust/src/server/src/grpc/mod.rs b/rust/src/server/src/grpc/mod.rs index 06c7c9259eb8..2f0eea71afc8 100644 --- a/rust/src/server/src/grpc/mod.rs +++ b/rust/src/server/src/grpc/mod.rs @@ -7,6 +7,11 @@ mod control; mod convert; mod health; mod inference; +mod lora_rpc; +mod struct_json; + +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; /// Generated protobuf/gRPC types for the `vllm` package. pub mod pb { @@ -21,5 +26,51 @@ pub use inference::InferenceServiceImpl; pub use pb::control_server::ControlServer; pub use pb::inference_server::InferenceServer; +/// Drain/admission state shared by the inference and control services. +/// +/// `Drain` arrives on the control service but must stop admitting work on the +/// inference service, so both hold the same `Arc`. +#[derive(Default)] +pub(crate) struct AdmissionState { + draining: AtomicBool, + in_flight: AtomicU64, +} + +impl AdmissionState { + fn is_draining(&self) -> bool { + self.draining.load(Ordering::SeqCst) + } + + pub(crate) fn begin_drain(&self) { + self.draining.store(true, Ordering::SeqCst); + } + + pub(crate) fn in_flight(&self) -> u64 { + self.in_flight.load(Ordering::SeqCst) + } + + /// Reserve an in-flight slot, or return `None` once draining has begun. + pub(crate) fn try_admit(self: &Arc) -> Option { + if self.is_draining() { + return None; + } + self.in_flight.fetch_add(1, Ordering::SeqCst); + // Re-check: `begin_drain` may land between the check and the increment. + if self.is_draining() { + self.in_flight.fetch_sub(1, Ordering::SeqCst); + return None; + } + Some(AdmissionGuard(self.clone())) + } +} + +pub(crate) struct AdmissionGuard(Arc); + +impl Drop for AdmissionGuard { + fn drop(&mut self) { + self.0.in_flight.fetch_sub(1, Ordering::SeqCst); + } +} + #[cfg(test)] mod tests; diff --git a/rust/src/server/src/grpc/struct_json.rs b/rust/src/server/src/grpc/struct_json.rs new file mode 100644 index 000000000000..779347203df3 --- /dev/null +++ b/rust/src/server/src/grpc/struct_json.rs @@ -0,0 +1,84 @@ +//! Lossless-enough conversion between JSON values and protobuf `Struct`. + +pub(crate) fn json_to_prost_struct(value: &serde_json::Value) -> Option { + match value { + serde_json::Value::Object(map) => Some(prost_types::Struct { + fields: map + .iter() + .map(|(key, value)| (key.clone(), json_to_prost_value(value))) + .collect(), + }), + _ => None, + } +} + +fn json_to_prost_value(value: &serde_json::Value) -> prost_types::Value { + use prost_types::value::Kind; + let kind = match value { + serde_json::Value::Null => Kind::NullValue(prost_types::NullValue::NullValue as i32), + serde_json::Value::Bool(value) => Kind::BoolValue(*value), + serde_json::Value::Number(value) => Kind::NumberValue(value.as_f64().unwrap_or(0.0)), + serde_json::Value::String(value) => Kind::StringValue(value.clone()), + serde_json::Value::Array(values) => Kind::ListValue(prost_types::ListValue { + values: values.iter().map(json_to_prost_value).collect(), + }), + serde_json::Value::Object(values) => Kind::StructValue(prost_types::Struct { + fields: values + .iter() + .map(|(key, value)| (key.clone(), json_to_prost_value(value))) + .collect(), + }), + }; + prost_types::Value { kind: Some(kind) } +} + +pub(crate) fn prost_struct_to_json(value: &prost_types::Struct) -> serde_json::Value { + serde_json::Value::Object( + value + .fields + .iter() + .map(|(key, value)| (key.clone(), prost_value_to_json(value))) + .collect(), + ) +} + +fn prost_value_to_json(value: &prost_types::Value) -> serde_json::Value { + use prost_types::value::Kind; + match &value.kind { + None | Some(Kind::NullValue(_)) => serde_json::Value::Null, + Some(Kind::BoolValue(value)) => serde_json::Value::Bool(*value), + Some(Kind::NumberValue(value)) => number_to_json(*value), + Some(Kind::StringValue(value)) => serde_json::Value::String(value.clone()), + Some(Kind::ListValue(values)) => { + serde_json::Value::Array(values.values.iter().map(prost_value_to_json).collect()) + } + Some(Kind::StructValue(value)) => prost_struct_to_json(value), + } +} + +fn number_to_json(value: f64) -> serde_json::Value { + if value.is_finite() + && value.fract() == 0.0 + && value >= i64::MIN as f64 + && value <= i64::MAX as f64 + { + serde_json::Value::Number((value as i64).into()) + } else { + serde_json::Number::from_f64(value) + .map(serde_json::Value::Number) + .unwrap_or(serde_json::Value::Null) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn integral_struct_numbers_return_as_json_integers() { + let input = serde_json::json!({"port": 8000, "ratio": 1.5}); + let round_trip = prost_struct_to_json(&json_to_prost_struct(&input).unwrap()); + assert_eq!(round_trip, input); + assert!(round_trip["port"].is_i64()); + } +} diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index 8265fefa050e..34f966b7fcef 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -253,9 +253,15 @@ async fn setup_grpc_service( Arc::new(FakeTextBackend) as Arc, ); let state = Arc::new(AppState::new(vec!["test-model".to_string()], chat)); + // Inference and control must share one AdmissionState so `Drain` on the + // control service stops the inference service admitting new work. + let admission = std::sync::Arc::new(super::AdmissionState::default()); ( - InferenceServer::new(InferenceServiceImpl::new(state.clone())), - ControlServer::new(ControlServiceImpl::new(state)), + InferenceServer::new(InferenceServiceImpl::with_admission( + state.clone(), + admission.clone(), + )), + ControlServer::new(ControlServiceImpl::with_admission(state, admission, None)), engine_health, engine_task, ) @@ -811,8 +817,8 @@ async fn unary_generate_with_sampling_params() { prompt: Some(pb::generate_request::Prompt::Text("test".to_string())), temperature: Some(0.7), sampling: Some(pb::RandomSampling { - top_k: 50, - top_p: 0.9, + top_k: Some(50), + top_p: Some(0.9), seed: Some(42), ..Default::default() }), diff --git a/rust/src/server/src/grpc/tests/lora.rs b/rust/src/server/src/grpc/tests/lora.rs new file mode 100644 index 000000000000..8c1b0f1a2fa4 --- /dev/null +++ b/rust/src/server/src/grpc/tests/lora.rs @@ -0,0 +1,192 @@ +use std::path::PathBuf; + +use vllm_engine_core_client::mock_engine::default_ready_response; +use vllm_engine_core_client::protocol::decode_value; +use vllm_engine_core_client::protocol::output::{EngineCoreOutputs, UtilityCallOutput}; +use vllm_engine_core_client::protocol::utility::{UtilityOutput, UtilityResultEnvelope}; + +use super::*; + +async fn grpc_lora_test_server( + engine_id: impl Into, + allowed_path_prefixes: Vec, + runtime_updates_enabled: bool, + run: F, +) -> ( + ControlClient, + tokio::task::JoinHandle<()>, + MockEngineTask, +) +where + F: for<'a> FnOnce(&'a mut DealerSocket, &'a mut PushSocket) -> TestFuture<'a> + Send + 'static, +{ + let mut ready = default_ready_response(); + ready.supports_lora = true; + ready.max_loras = 1; + let (service, engine_health, engine_task) = + setup_grpc_service_with_ready_and_engine(engine_id, ready, run).await; + let service = service + .with_lora_allowed_path_prefixes(allowed_path_prefixes) + .with_runtime_lora_updating(runtime_updates_enabled); + let (_generate, control, _health, server_task, engine_task) = + start_grpc_test_server(service, engine_health, engine_task).await; + (control, server_task, engine_task) +} + +async fn send_utility_result(push: &mut PushSocket, call_id: u64, result: bool) { + let output: EngineCoreOutputs = UtilityCallOutput { + engine_index: 0, + timestamp: 0.0, + output: UtilityOutput { + call_id: call_id.into(), + failure_message: None, + result: Some(UtilityResultEnvelope::without_type_info(rmpv::Value::from( + result, + ))), + }, + } + .into(); + send_outputs(push, output).await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn load_accepts_directory_under_injected_prefix() { + let temp = tempfile::tempdir().expect("create temporary LoRA root"); + let allowed = temp.path().join("allowed"); + let adapter = allowed.join("adapter-a"); + tokio::fs::create_dir_all(&adapter).await.expect("create adapter directory"); + let canonical_adapter = tokio::fs::canonicalize(&adapter).await.expect("canonical adapter"); + let expected_path = canonical_adapter.to_string_lossy().into_owned(); + let engine_expected_path = expected_path.clone(); + + let (mut client, server_task, engine_task) = grpc_lora_test_server( + b"engine-grpc-lora-allowed", + vec![allowed], + true, + move |dealer, push| { + boxed_test_future(async move { + let utility = recv_engine_message(dealer).await; + let payload = decode_value(&utility[1]).expect("decode utility payload"); + let array = payload.as_array().expect("utility payload array"); + assert_eq!(array[2], rmpv::Value::from("add_lora")); + let lora = array[3].as_array().expect("utility args")[0] + .as_array() + .expect("LoRA request tuple"); + assert_eq!(lora[0], rmpv::Value::from("adapter-a")); + assert_eq!(lora[1], rmpv::Value::from(1)); + assert_eq!(lora[2], rmpv::Value::from(engine_expected_path)); + send_utility_result(push, array[1].as_u64().expect("call id"), true).await; + }) + }, + ) + .await; + + let response = client + .load_lora(pb::LoadLoraRequest { + adapter: Some(pb::LoraAdapter { + lora_id: 1, + lora_name: "adapter-a".to_string(), + source_path: adapter.to_string_lossy().into_owned(), + }), + load_inplace: false, + }) + .await + .expect("load adapter under configured prefix") + .into_inner(); + assert_eq!( + response.adapter.expect("loaded adapter").source_path, + expected_path + ); + + engine_task.await.expect("mock engine task"); + server_task.abort(); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn load_rejects_directory_outside_injected_prefix() { + let temp = tempfile::tempdir().expect("create temporary LoRA root"); + let allowed = temp.path().join("allowed"); + let outside = temp.path().join("outside").join("adapter-a"); + tokio::fs::create_dir_all(&allowed).await.expect("create allowed directory"); + tokio::fs::create_dir_all(&outside).await.expect("create outside adapter"); + + let (mut client, server_task, engine_task) = grpc_lora_test_server( + b"engine-grpc-lora-outside", + vec![allowed], + true, + |_dealer, _push| boxed_test_future(async {}), + ) + .await; + let error = client + .load_lora(pb::LoadLoraRequest { + adapter: Some(pb::LoraAdapter { + lora_id: 1, + lora_name: "adapter-a".to_string(), + source_path: outside.to_string_lossy().into_owned(), + }), + load_inplace: false, + }) + .await + .expect_err("reject adapter outside configured prefix"); + assert_eq!(error.code(), tonic::Code::InvalidArgument); + + engine_task.await.expect("mock engine task"); + server_task.abort(); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn load_does_not_expose_unavailable_configured_prefix() { + let temp = tempfile::tempdir().expect("create temporary LoRA root"); + let adapter = temp.path().join("adapter-a"); + let missing_prefix = temp.path().join("secret-configured-prefix"); + tokio::fs::create_dir_all(&adapter).await.expect("create adapter directory"); + + let (mut client, server_task, engine_task) = grpc_lora_test_server( + b"engine-grpc-lora-missing-prefix", + vec![missing_prefix.clone()], + true, + |_dealer, _push| boxed_test_future(async {}), + ) + .await; + let error = client + .load_lora(pb::LoadLoraRequest { + adapter: Some(pb::LoraAdapter { + lora_id: 1, + lora_name: "adapter-a".to_string(), + source_path: adapter.to_string_lossy().into_owned(), + }), + load_inplace: false, + }) + .await + .expect_err("unavailable configured prefix is a server error"); + + assert_eq!(error.code(), tonic::Code::Internal); + assert!(!error.message().contains(&missing_prefix.to_string_lossy().into_owned())); + assert_eq!( + error.message(), + "Runtime LoRA path policy is unavailable; check the server configuration." + ); + + engine_task.await.expect("mock engine task"); + server_task.abort(); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn load_requires_runtime_lora_updating() { + let (mut client, server_task, engine_task) = grpc_lora_test_server( + b"engine-grpc-lora-disabled", + Vec::new(), + false, + |_dealer, _push| boxed_test_future(async {}), + ) + .await; + + let error = client + .load_lora(pb::LoadLoraRequest::default()) + .await + .expect_err("runtime LoRA updating is disabled"); + assert_eq!(error.code(), tonic::Code::FailedPrecondition); + + engine_task.await.expect("mock engine task"); + server_task.abort(); +} diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs index 224ec672aec7..dc43ff59ce1c 100644 --- a/rust/src/server/src/lib.rs +++ b/rust/src/server/src/lib.rs @@ -1,6 +1,3 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project - //! Minimal OpenAI-compatible HTTP server above [`vllm_chat`]. mod config; @@ -8,6 +5,7 @@ mod error; mod grpc; mod listener; mod lora; +mod lora_path; mod middleware; mod routes; mod runtime; @@ -36,10 +34,12 @@ use hyper_util::rt::{TokioIo, TokioTimer}; use hyper_util::server::graceful::GracefulShutdown; use hyper_util::service::TowerToHyperService; use tokio::net::TcpListener; +use tokio::sync::watch; use tokio::time::{Instant, sleep_until}; use tokio_util::sync::CancellationToken; use tonic::transport::Server as TonicServer; -use tonic_health::server::health_reporter; +use tonic_health::ServingStatus; +use tonic_health::server::{HealthReporter, health_reporter}; use tower::ServiceExt as _; use tracing::{info, trace, warn}; use vllm_chat::{ChatLlm, LoadModelBackendsOptions, load_model_backends}; @@ -51,6 +51,11 @@ use vllm_text::TextLlm; use crate::listener::{Listener, MaybeTlsListener}; use crate::routes::build_router; use crate::server_info::ServerInfoSnapshot; + +// Preprocessed multimodal tensors are capped at 16 MiB; leave bounded room for +// protobuf framing and ordinary request metadata without opening a 64 MiB +// allocation-amplification surface on every Generate RPC. +const GRPC_MAX_REQUEST_SIZE: usize = 20 * 1024 * 1024; use crate::state::AppState; /// How often the server PINGs an idle gRPC connection to reap a dead peer; @@ -59,6 +64,65 @@ const GRPC_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(7200); /// How long the server waits for a keepalive PING reply before dropping the gRPC /// connection. 20s matches the gRPC-core default. const GRPC_KEEPALIVE_TIMEOUT: Duration = Duration::from_secs(20); +const GRPC_LORA_HEALTH_POLL_INTERVAL: Duration = Duration::from_secs(1); + +async fn set_generate_not_serving(health_reporter: &HealthReporter) { + health_reporter + .set_not_serving::() + .await; + health_reporter.set_service_status("", ServingStatus::NotServing).await; +} + +async fn set_grpc_not_serving(health_reporter: &HealthReporter) { + set_generate_not_serving(health_reporter).await; + health_reporter + .set_not_serving::>() + .await; +} + +async fn wait_until_engine_unhealthy(mut engine_health: watch::Receiver) { + loop { + if !*engine_health.borrow_and_update() { + return; + } + if engine_health.changed().await.is_err() { + return; + } + } +} + +async fn monitor_grpc_health( + health_reporter: HealthReporter, + engine_health: watch::Receiver, + shutdown: CancellationToken, +) { + tokio::select! { + _ = wait_until_engine_unhealthy(engine_health) => {} + _ = shutdown.cancelled() => {} + } + + set_grpc_not_serving(&health_reporter).await; +} + +async fn monitor_lora_health( + state: Arc, + health_reporter: HealthReporter, + shutdown: CancellationToken, +) { + let mut interval = tokio::time::interval(GRPC_LORA_HEALTH_POLL_INTERVAL); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + tokio::select! { + _ = shutdown.cancelled() => return, + _ = interval.tick() => { + if !state.lora_state_is_consistent() { + set_grpc_not_serving(&health_reporter).await; + return; + } + } + } + } +} /// Resolve the public model names accepted by the frontend. fn effective_served_model_names(model: &str, served_model_name: &[String]) -> Vec { @@ -101,7 +165,6 @@ async fn build_state(config: &Config) -> Result> { .default_chat_template_kwargs .clone() .unwrap_or_default(), - limit_mm_per_prompt: config.limit_mm_per_prompt.clone(), }, ) .await @@ -181,15 +244,12 @@ where result = build_state(&config) => result?, _ = shutdown.cancelled() => return Ok(()), }; - let model = state.primary_model_name().to_owned(); - let app = extend_router(build_router(state.clone())); - - info!(model, "starting vLLM server"); - let listener = Listener::bind(&config.listener_mode) .await .context("failed to bind listener for OpenAI server")?; let bind_address = listener.local_addr_display()?; + let model = state.primary_model_name().to_owned(); + let app = extend_router(build_router(state.clone())); // Optionally bind the gRPC Inference server on a separate port. Bind // synchronously here so bind errors (port in use, permission denied, ...) @@ -212,10 +272,16 @@ where let engine_health = state.engine_core_client().subscribe_health(); health_reporter.set_serving::().await; health_reporter.set_serving::().await; - let control_service = - grpc::ControlGrpcService::new(grpc::ControlServiceImpl::new(state.clone())); - let inference_service = - grpc::InferenceGrpcService::new(grpc::InferenceServiceImpl::new(state.clone())); + let admission = std::sync::Arc::new(grpc::AdmissionState::default()); + let control_service = grpc::ControlGrpcService::new(grpc::ControlServiceImpl::with_admission( + state.clone(), + admission.clone(), + Some(health_reporter.clone()), + )); + let inference_service = grpc::InferenceGrpcService::new( + grpc::InferenceServiceImpl::with_admission(state.clone(), admission), + ) + .max_decoding_message_size(GRPC_MAX_REQUEST_SIZE); let svc = TonicServer::builder() .http2_keepalive_interval(Some(GRPC_KEEPALIVE_INTERVAL)) .http2_keepalive_timeout(Some(GRPC_KEEPALIVE_TIMEOUT)) @@ -223,14 +289,8 @@ where .add_service(health_service) .add_service(control_service) .add_service(inference_service); - Some(( - addr, - grpc_listener, - svc, - grpc_tls, - health_reporter, - engine_health, - )) + info!(%addr, tls = grpc_tls.is_some(), "starting gRPC server"); + Some((grpc_listener, svc, grpc_tls, health_reporter, engine_health)) } else { None }; @@ -240,7 +300,7 @@ where } else { "http" }; - let model = model.as_str(); + info!(%bind_address, %scheme, %model, "starting OpenAI server"); // Run HTTP and gRPC concurrently under a child token of the caller's shutdown // token. Caller cancellation propagates into both protocols; if either @@ -250,6 +310,14 @@ where let force_shutdown = CancellationToken::new(); let shutdown_deadline = Arc::new(OnceLock::new()); + let (grpc_server_setup, grpc_health_setup) = match grpc_setup { + Some((listener, service, tls, reporter, engine_health)) => ( + Some((listener, service, tls)), + Some((reporter, engine_health)), + ), + None => (None, None), + }; + // Spawn a task to trigger `force_shutdown` after shutdown deadline elapses. tokio::spawn({ let shutdown = server_shutdown.clone(); @@ -294,11 +362,6 @@ where }; let server = serve_connections(listener, app, shutdown.cancelled_owned(), timeouts); - info!( - bind_address, - scheme, model, "OpenAI server is ready to accept requests" - ); - let result = tokio::select! { result = server => { result.context("HTTP server failed") @@ -319,46 +382,48 @@ where let server_shutdown = server_shutdown.clone(); let force_shutdown = force_shutdown.clone(); async move { - let Some((addr, grpc_listener, svc, grpc_tls, health_reporter, engine_health)) = - grpc_setup - else { + let Some((grpc_listener, svc, grpc_tls)) = grpc_server_setup else { // No gRPC configured: just wait for shutdown so we do not race the // join! by resolving early and tripping the cancellation token. shutdown.cancelled().await; return Ok(()); }; - let tls = grpc_tls.is_some(); let incoming = match grpc_tls { Some(context) => MaybeTlsListener::tls(grpc_listener, context), None => MaybeTlsListener::plain(grpc_listener), }; - let server = - svc.serve_with_incoming_shutdown(incoming, shutdown.clone().cancelled_owned()); - let health_monitor = grpc::monitor_health(health_reporter, engine_health, shutdown); - - info!(%addr, tls, model, "gRPC server is ready to accept requests"); - - let server = async move { - let result = tokio::select! { - result = server => { - result.context("gRPC server failed") - } - _ = force_shutdown.cancelled() => { - warn!("gRPC graceful shutdown deadline elapsed; aborting server"); - Ok(()) - } - }; - - server_shutdown.cancel(); - result + let server = svc.serve_with_incoming_shutdown(incoming, shutdown.cancelled_owned()); + + let result = tokio::select! { + result = server => { + result.context("gRPC server failed") + } + _ = force_shutdown.cancelled() => { + warn!("gRPC graceful shutdown deadline elapsed; aborting server"); + Ok(()) + } }; - let (result, ()) = tokio::join!(server, health_monitor); + server_shutdown.cancel(); result } }; - let (http_res, grpc_res) = tokio::join!(http_fut, grpc_fut); + let grpc_health_fut = { + let state = state.clone(); + let shutdown = server_shutdown.child_token(); + async move { + let Some((health_reporter, engine_health)) = grpc_health_setup else { + return; + }; + let lora_health = + monitor_lora_health(state, health_reporter.clone(), shutdown.child_token()); + let engine_health = monitor_grpc_health(health_reporter, engine_health, shutdown); + tokio::join!(engine_health, lora_health); + } + }; + + let (http_res, grpc_res, ()) = tokio::join!(http_fut, grpc_fut, grpc_health_fut); http_res.and(grpc_res)?; let shutdown_deadline = shutdown_deadline diff --git a/rust/src/server/src/lora.rs b/rust/src/server/src/lora.rs index 896ba32aac45..cceca74b312c 100644 --- a/rust/src/server/src/lora.rs +++ b/rust/src/server/src/lora.rs @@ -1,33 +1,81 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project -use std::sync::atomic::{AtomicU64, Ordering}; +use std::pin::Pin; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::task::{Context, Poll}; +use futures::Stream; +use futures::future::join_all; use indexmap::IndexMap; -use tokio::sync::{Mutex, RwLock}; +use tokio::sync::{Mutex, OwnedRwLockReadGuard, RwLock}; use vllm_engine_core_client::EngineCoreClient; +use vllm_engine_core_client::Error as EngineCoreError; use vllm_engine_core_client::protocol::lora::LoraRequest; +#[cfg(not(test))] +const LORA_MUTATION_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); +#[cfg(test)] +const LORA_MUTATION_TIMEOUT: std::time::Duration = std::time::Duration::from_millis(250); + +struct MutationGuard<'a> { + consistent: &'a AtomicBool, + final_state_proven: bool, +} + +impl<'a> MutationGuard<'a> { + fn new(consistent: &'a AtomicBool) -> Self { + Self { + consistent, + final_state_proven: false, + } + } + + fn prove_final_state(&mut self) { + self.final_state_proven = true; + } +} + +impl Drop for MutationGuard<'_> { + fn drop(&mut self) { + if !self.final_state_proven { + self.consistent.store(false, Ordering::Release); + } + } +} + /// Snapshot of the currently served model names plus the requested LoRA, if /// the model name resolves to a dynamic adapter. -#[derive(Debug, Clone)] +pub(crate) type LoraLease = Option>; + +#[derive(Debug)] pub(crate) struct LoraModelResolution { pub model_names: Vec, pub lora_request: Option, + pub lease: LoraLease, +} + +#[derive(Clone)] +struct LoadedLora { + request: LoraRequest, + lease: std::sync::Arc>, } /// Runtime registry for dynamically loaded LoRA adapters. pub(crate) struct LoraManager { - /// Dynamically loaded LoRA adapters keyed by public model name, in load order. - requests: RwLock>, + /// Loaded adapters and their generation leases, keyed by public model name. + registry: RwLock>, /// Monotonic adapter id allocator. LoRA ids are one-indexed. id_counter: AtomicU64, /// Serialize dynamic LoRA registry updates around engine utility calls. update_lock: Mutex<()>, + /// False after a failed compensation leaves per-engine state indeterminate. + consistent: AtomicBool, } #[derive(Debug)] pub(crate) enum LoadLoraError { + Inconsistent, AlreadyLoaded { lora_name: String }, BaseModelName { lora_name: String }, Engine(vllm_engine_core_client::Error), @@ -36,6 +84,7 @@ pub(crate) enum LoadLoraError { #[derive(Debug)] pub(crate) enum UnloadLoraError { + Inconsistent, NotFound { lora_name: String, }, @@ -45,24 +94,44 @@ pub(crate) enum UnloadLoraError { actual: u64, }, Engine(vllm_engine_core_client::Error), - NotRemoved { - lora_name: String, - lora_int_id: u64, - }, +} + +#[derive(Debug)] +pub(crate) enum LoadExactLoraError { + Inconsistent, + BaseModelName { lora_name: String }, + Conflict { existing: LoraRequest }, + Engine(vllm_engine_core_client::Error), + NotLoaded { lora_name: String }, +} + +enum ApplyLoadError { + Engine(EngineCoreError), + Rejected, } impl LoraManager { pub fn new() -> Self { Self { - requests: RwLock::new(IndexMap::new()), + registry: RwLock::new(IndexMap::new()), id_counter: AtomicU64::new(0), update_lock: Mutex::new(()), + consistent: AtomicBool::new(true), } } + pub fn is_consistent(&self) -> bool { + self.consistent.load(Ordering::Acquire) + } + /// Snapshot loaded LoRA adapters in load order. pub async fn served_lora_requests(&self) -> Vec { - self.requests.read().await.values().cloned().collect() + self.registry + .read() + .await + .values() + .map(|loaded| loaded.request.clone()) + .collect() } /// Resolve the requested model against one consistent LoRA registry @@ -72,14 +141,34 @@ impl LoraManager { base_model_names: &[String], model_name: Option<&str>, ) -> LoraModelResolution { - let requests = self.requests.read().await; - let mut model_names = base_model_names.to_vec(); - model_names.extend(requests.keys().cloned()); - let lora_request = model_name.and_then(|name| requests.get(name).cloned()); + loop { + let candidate = match model_name { + Some(name) => { + self.registry.read().await.get(name).map(|loaded| loaded.lease.clone()) + } + None => None, + }; + let lease = match &candidate { + Some(lease) => Some(lease.clone().read_owned().await), + None => None, + }; + let registry = self.registry.read().await; + let current = model_name.and_then(|name| registry.get(name)); + if !same_lease(current, candidate.as_ref()) { + drop(registry); + drop(lease); + tokio::task::yield_now().await; + continue; + } - LoraModelResolution { - model_names, - lora_request, + let mut model_names = base_model_names.to_vec(); + model_names.extend(registry.keys().cloned()); + let lora_request = current.map(|loaded| loaded.request.clone()); + return LoraModelResolution { + model_names, + lora_request, + lease, + }; } } @@ -94,19 +183,29 @@ impl LoraManager { is_3d_lora_weight: bool, ) -> Result { let _guard = self.update_lock.lock().await; + if !self.is_consistent() { + return Err(LoadLoraError::Inconsistent); + } if base_model_names.iter().any(|name| name == &lora_name) { return Err(LoadLoraError::BaseModelName { lora_name }); } - if !load_inplace && self.requests.read().await.contains_key(&lora_name) { - return Err(LoadLoraError::AlreadyLoaded { lora_name }); - } - - let lora_int_id = self - .requests + let previous = self + .registry .read() .await .get(&lora_name) - .map(|request| request.lora_int_id) + .map(|loaded| (loaded.request.clone(), loaded.lease.clone())); + if previous.is_some() && !load_inplace { + return Err(LoadLoraError::AlreadyLoaded { lora_name }); + } + let _generation_guard = match previous.as_ref().map(|(_, lease)| lease.clone()) { + Some(lease) => Some(lease.write_owned().await), + None => None, + }; + + let lora_int_id = previous + .as_ref() + .map(|(request, _)| request.lora_int_id) .unwrap_or_else(|| self.id_counter.fetch_add(1, Ordering::Relaxed) + 1); let lora_request = LoraRequest::new( lora_name.clone(), @@ -116,15 +215,162 @@ impl LoraManager { is_3d_lora_weight, ); - let loaded = engine_core_client - .add_lora(&lora_request) + let mut mutation = self + .apply_load( + engine_core_client, + &lora_request, + previous.as_ref().map(|(request, _)| request), + ) + .await + .map_err(|error| match error { + ApplyLoadError::Engine(error) => LoadLoraError::Engine(error), + ApplyLoadError::Rejected => LoadLoraError::NotLoaded { + lora_name: lora_name.clone(), + }, + })?; + let lease = previous + .map(|(_, lease)| lease) + .unwrap_or_else(|| std::sync::Arc::new(RwLock::new(()))); + let mut stored_request = lora_request; + stored_request.load_inplace = false; + self.registry.write().await.insert( + lora_name, + LoadedLora { + request: stored_request.clone(), + lease, + }, + ); + mutation.prove_final_state(); + Ok(stored_request) + } + + /// Load an adapter with a caller-supplied ID. + pub async fn load_lora_exact( + &self, + engine_core_client: &EngineCoreClient, + base_model_names: &[String], + mut lora_request: LoraRequest, + load_inplace: bool, + ) -> Result<(LoraRequest, bool), LoadExactLoraError> { + let _guard = self.update_lock.lock().await; + if !self.is_consistent() { + return Err(LoadExactLoraError::Inconsistent); + } + if base_model_names.iter().any(|name| name == &lora_request.lora_name) { + return Err(LoadExactLoraError::BaseModelName { + lora_name: lora_request.lora_name, + }); + } + + let registry = self.registry.read().await; + let same_name = registry.get(&lora_request.lora_name); + if let Some(existing) = registry.values().find(|loaded| { + loaded.request.lora_name != lora_request.lora_name + && (loaded.request.lora_int_id == lora_request.lora_int_id + || loaded.request.lora_path == lora_request.lora_path) + }) { + return Err(LoadExactLoraError::Conflict { + existing: existing.request.clone(), + }); + } + let previous = if let Some(existing) = same_name { + if same_wire_identity(&existing.request, &lora_request) && !load_inplace { + return Ok((existing.request.clone(), true)); + } + if !load_inplace + || existing.request.lora_name != lora_request.lora_name + || existing.request.lora_int_id != lora_request.lora_int_id + { + return Err(LoadExactLoraError::Conflict { + existing: existing.request.clone(), + }); + } + Some((existing.request.clone(), existing.lease.clone())) + } else { + None + }; + drop(registry); + + let _generation_guard = match previous.as_ref().map(|(_, lease)| lease.clone()) { + Some(lease) => Some(lease.write_owned().await), + None => None, + }; + lora_request.load_inplace = load_inplace; + + let mut mutation = self + .apply_load( + engine_core_client, + &lora_request, + previous.as_ref().map(|(request, _)| request), + ) .await - .map_err(LoadLoraError::Engine)?; - if !loaded { - return Err(LoadLoraError::NotLoaded { lora_name }); + .map_err(|error| match error { + ApplyLoadError::Engine(error) => LoadExactLoraError::Engine(error), + ApplyLoadError::Rejected => LoadExactLoraError::NotLoaded { + lora_name: lora_request.lora_name.clone(), + }, + })?; + + self.id_counter.fetch_max(lora_request.lora_int_id, Ordering::Relaxed); + lora_request.load_inplace = false; + self.registry.write().await.insert( + lora_request.lora_name.clone(), + LoadedLora { + request: lora_request.clone(), + lease: previous + .map(|(_, lease)| lease) + .unwrap_or_else(|| std::sync::Arc::new(RwLock::new(()))), + }, + ); + mutation.prove_final_state(); + Ok((lora_request, false)) + } + + /// Apply one load or replacement across every engine rank. The returned + /// guard remains uncommitted until the caller updates the frontend + /// registry, so cancellation between engine success and registry commit + /// still fails closed. + async fn apply_load<'a>( + &'a self, + engine_core_client: &EngineCoreClient, + lora_request: &LoraRequest, + previous: Option<&LoraRequest>, + ) -> Result, ApplyLoadError> { + let mut mutation = MutationGuard::new(&self.consistent); + let results = match call_all_bounded::( + engine_core_client, + "add_lora", + (lora_request,), + ) + .await + { + Ok(results) => results, + Err(error) => { + if matches!(&error, EngineCoreError::UtilityCallTimeout { .. }) { + if restore_load_state(engine_core_client, lora_request, previous).await { + mutation.prove_final_state(); + } + } else { + // Outer errors occur before dispatch, so engine state is + // unchanged. + mutation.prove_final_state(); + } + return Err(ApplyLoadError::Engine(error)); + } + }; + + let Some(failure_index) = results.iter().position(|result| !matches!(result, Ok(true))) + else { + return Ok(mutation); + }; + if restore_load_state(engine_core_client, lora_request, previous).await { + mutation.prove_final_state(); + } + match results.into_iter().nth(failure_index).unwrap() { + Ok(false) => Err(ApplyLoadError::Rejected), + Err(error) => Err(ApplyLoadError::Engine(error)), + Ok(true) => unreachable!(), } - self.requests.write().await.insert(lora_name, lora_request.clone()); - Ok(lora_request) } /// Remove one dynamic LoRA adapter from the engine and public model @@ -136,11 +382,19 @@ impl LoraManager { requested_lora_int_id: Option, ) -> Result { let _guard = self.update_lock.lock().await; - let lora_request = self.requests.read().await.get(lora_name).cloned().ok_or_else(|| { - UnloadLoraError::NotFound { + if !self.is_consistent() { + return Err(UnloadLoraError::Inconsistent); + } + let (lora_request, lease) = self + .registry + .read() + .await + .get(lora_name) + .map(|loaded| (loaded.request.clone(), loaded.lease.clone())) + .ok_or_else(|| UnloadLoraError::NotFound { lora_name: lora_name.to_string(), - } - })?; + })?; + let _generation_guard = lease.write_owned().await; if let Some(actual) = requested_lora_int_id && actual != lora_request.lora_int_id @@ -152,17 +406,332 @@ impl LoraManager { }); } - let removed = engine_core_client - .remove_lora(lora_request.lora_int_id) + let mut mutation = MutationGuard::new(&self.consistent); + let results = match call_all_bounded::( + engine_core_client, + "remove_lora", + (lora_request.lora_int_id,), + ) + .await + { + Ok(results) => results, + Err(error) => { + if matches!(&error, EngineCoreError::UtilityCallTimeout { .. }) { + if restore_removed(engine_core_client, &lora_request).await { + mutation.prove_final_state(); + } + } else { + mutation.prove_final_state(); + } + return Err(UnloadLoraError::Engine(error)); + } + }; + if let Some(failure_index) = results.iter().position(Result::is_err) { + if restore_removed(engine_core_client, &lora_request).await { + mutation.prove_final_state(); + } + let failure = results.into_iter().nth(failure_index).unwrap(); + return match failure { + Err(error) => Err(UnloadLoraError::Engine(error)), + Ok(_) => unreachable!(), + }; + } + + let removed = self + .registry + .write() .await - .map_err(UnloadLoraError::Engine)?; - if !removed { - return Err(UnloadLoraError::NotRemoved { - lora_name: lora_request.lora_name, - lora_int_id: lora_request.lora_int_id, - }); + .shift_remove(lora_name) + .map(|loaded| loaded.request) + .unwrap_or(lora_request); + mutation.prove_final_state(); + Ok(removed) + } +} + +fn same_wire_identity(left: &LoraRequest, right: &LoraRequest) -> bool { + left.lora_name == right.lora_name + && left.lora_int_id == right.lora_int_id + && left.lora_path == right.lora_path +} + +fn same_lease( + current: Option<&LoadedLora>, + candidate: Option<&std::sync::Arc>>, +) -> bool { + match (current, candidate) { + (Some(current), Some(candidate)) => std::sync::Arc::ptr_eq(¤t.lease, candidate), + (None, None) => true, + _ => false, + } +} + +/// Hold an adapter's shared generation lease until the wrapped stream ends or +/// is dropped. +pub(crate) struct LoraLeaseStream { + stream: Pin>, + _lease: LoraLease, +} + +impl Unpin for LoraLeaseStream {} + +impl Stream for LoraLeaseStream { + type Item = S::Item; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.stream.as_mut().poll_next(cx) + } +} + +pub(crate) fn hold_lora_lease(stream: S, lease: LoraLease) -> LoraLeaseStream +where + S: Stream, +{ + LoraLeaseStream { + stream: Box::pin(stream), + _lease: lease, + } +} + +async fn call_all_bounded( + engine_core_client: &EngineCoreClient, + method: &str, + args: A, +) -> Result>, EngineCoreError> +where + T: serde::de::DeserializeOwned, + A: serde::Serialize + std::fmt::Debug, +{ + tokio::time::timeout( + LORA_MUTATION_TIMEOUT, + engine_core_client.call_utility_per_engine(method, args), + ) + .await + .map_err(|_| EngineCoreError::UtilityCallTimeout { + method: method.to_string(), + timeout: LORA_MUTATION_TIMEOUT, + })? +} + +async fn restore_load_state( + engine_core_client: &EngineCoreClient, + attempted: &LoraRequest, + previous: Option<&LoraRequest>, +) -> bool { + match previous { + // Reloading the same path cannot restore old bytes after an in-place + // filesystem update. Leave the manager inconsistent so serving fails + // closed until restart. + Some(previous) if previous.lora_path == attempted.lora_path => false, + Some(previous) => restore_previous_on_all(engine_core_client, previous).await, + None => remove_from_all(engine_core_client, attempted.lora_int_id).await, + } +} + +async fn remove_from_all(engine_core_client: &EngineCoreClient, lora_int_id: u64) -> bool { + let calls = engine_core_client.ready_responses().into_iter().map(|ready| async move { + engine_core_client + .call_utility_on_engine::( + ready.data_parallel_rank, + "remove_lora", + (lora_int_id,), + ) + .await + }); + // `false` also proves the desired absent state. + tokio::time::timeout(LORA_MUTATION_TIMEOUT, join_all(calls)) + .await + .is_ok_and(|outcomes| outcomes.iter().all(Result::is_ok)) +} + +async fn restore_previous_on_all( + engine_core_client: &EngineCoreClient, + previous: &LoraRequest, +) -> bool { + let mut previous = previous.clone(); + previous.load_inplace = true; + let calls = engine_core_client.ready_responses().into_iter().map(|ready| { + let previous = previous.clone(); + async move { + engine_core_client + .call_utility_on_engine::( + ready.data_parallel_rank, + "add_lora", + (&previous,), + ) + .await + } + }); + tokio::time::timeout(LORA_MUTATION_TIMEOUT, join_all(calls)) + .await + .is_ok_and(|outcomes| outcomes.iter().all(|outcome| matches!(outcome, Ok(true)))) +} + +async fn restore_removed( + engine_core_client: &EngineCoreClient, + lora_request: &LoraRequest, +) -> bool { + let calls = engine_core_client.ready_responses().into_iter().map(|ready| async move { + engine_core_client + .call_utility_on_engine::( + ready.data_parallel_rank, + "add_lora", + (lora_request,), + ) + .await + }); + tokio::time::timeout(LORA_MUTATION_TIMEOUT, join_all(calls)) + .await + .is_ok_and(|outcomes| outcomes.iter().all(|outcome| matches!(outcome, Ok(true)))) +} + +#[cfg(test)] +mod tests { + use super::*; + use vllm_engine_core_client::protocol::decode_value; + use vllm_engine_core_client::protocol::output::{EngineCoreOutputs, UtilityCallOutput}; + use vllm_engine_core_client::protocol::utility::{UtilityOutput, UtilityResultEnvelope}; + use vllm_engine_core_client::test_utils::{IpcNamespace, spawn_mock_engine_task_with_ready}; + use vllm_engine_core_client::{EngineCoreClientConfig, TransportMode}; + use zeromq::ZmqMessage; + use zeromq::prelude::{SocketRecv as _, SocketSend as _}; + + #[test] + fn wire_identity_ignores_internal_load_options() { + let public = LoraRequest::new( + "adapter-a".to_string(), + 17, + "/adapters/a".to_string(), + false, + false, + ); + let mut internal = public.clone(); + internal.load_inplace = true; + internal.is_3d_lora_weight = true; + internal.base_model_name = Some("base".to_string()); + internal.tensorizer_config_dict = Some(rmpv::Value::Map(vec![( + rmpv::Value::from("format"), + rmpv::Value::from("safetensors"), + )])); + + assert!(same_wire_identity(&public, &internal)); + } + + async fn reply_utility(push: &mut zeromq::PushSocket, call_id: u64, result: bool) { + let output: EngineCoreOutputs = UtilityCallOutput { + engine_index: 0, + timestamp: 0.0, + output: UtilityOutput { + call_id: call_id.into(), + failure_message: None, + result: Some(UtilityResultEnvelope::without_type_info(rmpv::Value::from( + result, + ))), + }, } + .into(); + push.send(ZmqMessage::from(rmp_serde::to_vec_named(&output).unwrap())) + .await + .unwrap(); + } + + async fn recv_utility_call_id(dealer: &mut zeromq::DealerSocket, method: &str) -> u64 { + let message = dealer.recv().await.unwrap().into_vec(); + let payload = decode_value(&message[1]).unwrap(); + let array = payload.as_array().unwrap(); + assert_eq!(array[2], rmpv::Value::from(method)); + array[1].as_u64().unwrap() + } + + #[test] + fn lease_identity_rejects_reloaded_adapter() { + let old_lease = std::sync::Arc::new(RwLock::new(())); + let reloaded = LoadedLora { + request: LoraRequest::new( + "adapter-a".to_string(), + 17, + "/adapters/a".to_string(), + false, + false, + ), + lease: std::sync::Arc::new(RwLock::new(())), + }; + + assert!(!same_lease(Some(&reloaded), Some(&old_lease))); + assert!(same_lease(Some(&reloaded), Some(&reloaded.lease))); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn partial_rank_load_failure_rolls_back_every_rank() { + let ipc = IpcNamespace::new().unwrap(); + let handshake = ipc.handshake_endpoint(); + let ready = |rank| { + let mut response = vllm_engine_core_client::mock_engine::default_ready_response(); + response.data_parallel_size = 2; + response.data_parallel_rank = rank; + response + }; + + let (shutdown_zero, task_zero) = spawn_mock_engine_task_with_ready( + handshake.clone(), + vec![0x00, 0x00], + ready(0), + |dealer, push| { + Box::pin(async move { + let load = recv_utility_call_id(dealer, "add_lora").await; + reply_utility(push, load, true).await; + let rollback = recv_utility_call_id(dealer, "remove_lora").await; + reply_utility(push, rollback, true).await; + }) + }, + ); + let (shutdown_one, task_one) = spawn_mock_engine_task_with_ready( + handshake.clone(), + vec![0x01, 0x00], + ready(1), + |dealer, push| { + Box::pin(async move { + let load = recv_utility_call_id(dealer, "add_lora").await; + reply_utility(push, load, false).await; + let rollback = recv_utility_call_id(dealer, "remove_lora").await; + reply_utility(push, rollback, true).await; + }) + }, + ); + + let mut config = EngineCoreClientConfig::new_single(handshake) + .with_model_name("test-model") + .with_local_input_output_addresses( + Some(ipc.input_endpoint()), + Some(ipc.output_endpoint()), + ); + let TransportMode::HandshakeOwner { engine_count, .. } = &mut config.transport_mode else { + unreachable!() + }; + *engine_count = 2; + let client = EngineCoreClient::connect(config).await.unwrap(); + let manager = LoraManager::new(); + let error = manager + .load_lora( + &client, + &["test-model".to_string()], + "adapter-a".to_string(), + "/adapters/a".to_string(), + false, + false, + ) + .await + .unwrap_err(); + assert!(matches!(error, LoadLoraError::NotLoaded { .. })); + assert!(manager.is_consistent()); + assert!(manager.served_lora_requests().await.is_empty()); - Ok(self.requests.write().await.shift_remove(lora_name).unwrap_or(lora_request)) + let _ = shutdown_zero.send(()); + let _ = shutdown_one.send(()); + task_zero.await.unwrap(); + task_one.await.unwrap(); + client.shutdown().await.unwrap(); } + #[path = "inplace.rs"] + mod inplace; } diff --git a/rust/src/server/src/lora/tests/inplace.rs b/rust/src/server/src/lora/tests/inplace.rs new file mode 100644 index 000000000000..144b9edad515 --- /dev/null +++ b/rust/src/server/src/lora/tests/inplace.rs @@ -0,0 +1,129 @@ +use super::*; + +struct ExactLoadFixture { + client: EngineCoreClient, + manager: LoraManager, + shutdown: tokio::sync::oneshot::Sender<()>, + task: tokio::task::JoinHandle<()>, +} + +impl ExactLoadFixture { + async fn start(results: impl IntoIterator) -> Self { + let results = results.into_iter().collect::>(); + let ipc = IpcNamespace::new().unwrap(); + let handshake = ipc.handshake_endpoint(); + let (shutdown, task) = spawn_mock_engine_task_with_ready( + handshake.clone(), + vec![0x00, 0x00], + vllm_engine_core_client::mock_engine::default_ready_response(), + move |dealer, push| { + Box::pin(async move { + for result in results { + let load = recv_utility_call_id(dealer, "add_lora").await; + reply_utility(push, load, result).await; + } + }) + }, + ); + let config = EngineCoreClientConfig::new_single(handshake) + .with_model_name("test-model") + .with_local_input_output_addresses( + Some(ipc.input_endpoint()), + Some(ipc.output_endpoint()), + ); + Self { + client: EngineCoreClient::connect(config).await.unwrap(), + manager: LoraManager::new(), + shutdown, + task, + } + } + + fn adapter(path: &str) -> LoraRequest { + LoraRequest::new("adapter-a".to_string(), 17, path.to_string(), false, false) + } + + async fn load_exact(&self, path: &str, load_inplace: bool) -> Result { + self.manager + .load_lora_exact( + &self.client, + &["test-model".to_string()], + Self::adapter(path), + load_inplace, + ) + .await + .map(|(_, already_loaded)| already_loaded) + } + + async fn finish(self) { + let _ = self.shutdown.send(()); + self.task.await.unwrap(); + self.client.shutdown().await.unwrap(); + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn exact_load_inplace_reloads_same_or_different_path_without_persisting_flag() { + for (first_path, replacement_path) in [ + ("/adapters/step-1", "/adapters/step-2"), + ("/adapters/stable", "/adapters/stable"), + ] { + let fixture = ExactLoadFixture::start([true, true]).await; + fixture.load_exact(first_path, false).await.unwrap(); + + let already_loaded = fixture.load_exact(replacement_path, true).await.unwrap(); + let loaded = fixture.manager.served_lora_requests().await; + + assert!(!already_loaded, "in-place load must reach the engine"); + assert_eq!(loaded.len(), 1); + assert_eq!(loaded[0].lora_name, "adapter-a"); + assert_eq!(loaded[0].lora_int_id, 17); + assert_eq!(loaded[0].lora_path, replacement_path); + assert!(!loaded[0].load_inplace); + fixture.finish().await; + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn failed_same_path_inplace_load_fails_closed() { + let fixture = ExactLoadFixture::start([true, false]).await; + fixture.load_exact("/adapters/stable", false).await.unwrap(); + + let error = fixture.load_exact("/adapters/stable", true).await.unwrap_err(); + + assert!(matches!(error, LoadExactLoraError::NotLoaded { .. })); + assert!(!fixture.manager.is_consistent()); + fixture.finish().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn named_load_inplace_does_not_persist_the_mutation_flag() { + let fixture = ExactLoadFixture::start([true, true]).await; + fixture + .manager + .load_lora( + &fixture.client, + &["test-model".to_string()], + "adapter-a".to_string(), + "/adapters/step-1".to_string(), + false, + false, + ) + .await + .unwrap(); + fixture + .manager + .load_lora( + &fixture.client, + &["test-model".to_string()], + "adapter-a".to_string(), + "/adapters/step-2".to_string(), + true, + false, + ) + .await + .unwrap(); + + assert!(!fixture.manager.served_lora_requests().await[0].load_inplace); + fixture.finish().await; +} diff --git a/rust/src/server/src/lora_path.rs b/rust/src/server/src/lora_path.rs new file mode 100644 index 000000000000..00c0830ccbc4 --- /dev/null +++ b/rust/src/server/src/lora_path.rs @@ -0,0 +1,248 @@ +//! Transport-neutral policy for resolving runtime LoRA adapter sources. + +use std::fmt; +use std::path::{Component, Path, PathBuf}; + +pub(crate) const RUNTIME_LORA_ALLOWED_PATH_PREFIXES_ENV: &str = + "VLLM_RUNTIME_LORA_ALLOWED_PATH_PREFIXES"; + +#[derive(Debug)] +pub(crate) enum LoraPathError { + LocalPathsDisabled, + RelativeLocalPath, + PathUnavailable, + AllowedPrefixUnavailable(PathBuf), + OutsideAllowedPrefixes, +} + +impl LoraPathError { + /// Misconfigured allowed prefixes are server errors; all other failures + /// describe an invalid adapter source supplied by the caller. + pub(crate) fn is_client_error(&self) -> bool { + !matches!(self, Self::AllowedPrefixUnavailable(_)) + } + + /// Stable caller-facing error text. Configuration paths remain available + /// through `Display` for server logs, but never cross the API boundary. + pub(crate) fn public_message(&self) -> &'static str { + match self { + Self::LocalPathsDisabled => { + "Local LoRA adapter paths require VLLM_RUNTIME_LORA_ALLOWED_PATH_PREFIXES to be configured." + } + Self::RelativeLocalPath => { + "Local LoRA adapter paths must be absolute and under one of the prefixes configured by VLLM_RUNTIME_LORA_ALLOWED_PATH_PREFIXES." + } + Self::PathUnavailable => "Local LoRA adapter path must exist and be accessible.", + Self::AllowedPrefixUnavailable(_) => { + "Runtime LoRA path policy is unavailable; check the server configuration." + } + Self::OutsideAllowedPrefixes => { + "Local LoRA adapter path is outside the configured allowed prefixes." + } + } + } +} + +impl fmt::Display for LoraPathError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::LocalPathsDisabled => write!( + formatter, + "Local LoRA adapter paths require {RUNTIME_LORA_ALLOWED_PATH_PREFIXES_ENV} to be configured." + ), + Self::RelativeLocalPath => write!( + formatter, + "Local LoRA adapter paths must be absolute and under one of the prefixes configured by {RUNTIME_LORA_ALLOWED_PATH_PREFIXES_ENV}." + ), + Self::PathUnavailable => { + formatter.write_str("Local LoRA adapter path must exist and be accessible.") + } + Self::AllowedPrefixUnavailable(prefix) => write!( + formatter, + "configured {RUNTIME_LORA_ALLOWED_PATH_PREFIXES_ENV} path prefix `{}` must exist and be accessible", + prefix.display() + ), + Self::OutsideAllowedPrefixes => formatter + .write_str("Local LoRA adapter path is outside the configured allowed prefixes."), + } + } +} + +pub(crate) fn runtime_lora_allowed_path_prefixes() -> Option> { + let prefixes = std::env::var_os(RUNTIME_LORA_ALLOWED_PATH_PREFIXES_ENV)?; + let prefixes: Vec<_> = std::env::split_paths(&prefixes) + .filter(|path| !path.as_os_str().is_empty()) + .collect(); + (!prefixes.is_empty()).then_some(prefixes) +} + +fn looks_like_local_lora_path(lora_path: &str) -> bool { + let path = Path::new(lora_path); + path.is_absolute() + || lora_path.starts_with('~') + || lora_path.starts_with('.') + || path.components().any(|component| matches!(component, Component::ParentDir)) +} + +/// Resolve a local adapter path under the configured allowlist. +/// +/// A `None` result means `lora_path` is a non-local identifier such as a +/// Hugging Face repository id and therefore does not need filesystem policy. +pub(crate) async fn validate_lora_path_access( + lora_path: &str, + allowed_prefixes: Option<&[PathBuf]>, +) -> Result, LoraPathError> { + let path = Path::new(lora_path); + if !looks_like_local_lora_path(lora_path) { + match tokio::fs::try_exists(path).await { + Ok(false) => return Ok(None), + Ok(true) => {} + Err(_) => return Err(LoraPathError::PathUnavailable), + } + } + + let Some(allowed_prefixes) = allowed_prefixes else { + return Err(LoraPathError::LocalPathsDisabled); + }; + + if !path.is_absolute() { + return Err(LoraPathError::RelativeLocalPath); + } + + let canonical_path = tokio::fs::canonicalize(path) + .await + .map_err(|_| LoraPathError::PathUnavailable)?; + let mut canonical_prefixes = Vec::with_capacity(allowed_prefixes.len()); + for prefix in allowed_prefixes { + canonical_prefixes.push( + tokio::fs::canonicalize(prefix) + .await + .map_err(|_| LoraPathError::AllowedPrefixUnavailable(prefix.clone()))?, + ); + } + + canonical_prefixes + .iter() + .any(|prefix| canonical_path.starts_with(prefix)) + .then_some(Some(canonical_path)) + .ok_or(LoraPathError::OutsideAllowedPrefixes) +} + +#[cfg(test)] +mod tests { + use std::fs; + use std::path::PathBuf; + use std::time::{SystemTime, UNIX_EPOCH}; + + use super::{LoraPathError, validate_lora_path_access}; + + fn temp_lora_dir(test_name: &str) -> PathBuf { + let suffix = SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("clock should be after unix epoch") + .as_nanos(); + let path = std::env::temp_dir().join(format!( + "vllm-lora-{test_name}-{}-{suffix}", + std::process::id() + )); + fs::create_dir_all(&path).expect("create temp lora dir"); + path + } + + #[tokio::test] + async fn allows_hf_repo_ids_without_prefixes() { + assert_eq!( + validate_lora_path_access("org/adapter-a", None) + .await + .expect("hf repo id should be allowed"), + None + ); + } + + #[tokio::test] + async fn rejects_local_paths_without_prefixes() { + for path in [ + "/tmp/adapter-a", + "./adapter-a", + "~/adapter-a", + "subdir/../../../etc/sensitive", + ] { + assert!(matches!( + validate_lora_path_access(path, None).await, + Err(LoraPathError::LocalPathsDisabled) + )); + } + } + + #[tokio::test] + async fn rejects_existing_bare_relative_paths_without_prefixes() { + let root = + PathBuf::from("target").join(format!("vllm-lora-relative-{}", std::process::id())); + let adapter = root.join("adapter-a"); + fs::create_dir_all(&adapter).expect("create relative adapter dir"); + + assert!(matches!( + validate_lora_path_access(adapter.to_str().expect("utf-8 temp path"), None).await, + Err(LoraPathError::LocalPathsDisabled) + )); + + fs::remove_dir_all(root).ok(); + } + + #[tokio::test] + async fn allows_absolute_paths_under_configured_prefixes() { + let root = temp_lora_dir("allowed-prefix"); + let allowed = root.join("allowed"); + let adapter = allowed.join("adapter-a"); + fs::create_dir_all(&adapter).expect("create adapter dir"); + + let prefixes = [allowed]; + let resolved = + validate_lora_path_access(adapter.to_str().expect("utf-8 temp path"), Some(&prefixes)) + .await + .expect("path under configured prefix should be allowed"); + assert_eq!( + resolved, + Some(adapter.canonicalize().expect("canonical adapter")) + ); + + fs::remove_dir_all(root).ok(); + } + + #[tokio::test] + async fn rejects_parent_escape_from_configured_prefixes() { + let root = temp_lora_dir("parent-escape"); + let allowed = root.join("allowed"); + let private_adapter = root.join("private").join("adapter-a"); + fs::create_dir_all(&allowed).expect("create allowed dir"); + fs::create_dir_all(&private_adapter).expect("create private adapter dir"); + + let escaped = allowed.join("../private/adapter-a"); + let prefixes = [allowed]; + assert!(matches!( + validate_lora_path_access(escaped.to_str().expect("utf-8 temp path"), Some(&prefixes)) + .await, + Err(LoraPathError::OutsideAllowedPrefixes) + )); + + fs::remove_dir_all(root).ok(); + } + + #[tokio::test] + async fn unavailable_allowed_prefix_is_a_server_error() { + let root = temp_lora_dir("missing-prefix"); + let adapter = root.join("adapter-a"); + fs::create_dir_all(&adapter).expect("create adapter dir"); + let prefixes = [root.join("missing")]; + + let error = + validate_lora_path_access(adapter.to_str().expect("utf-8 temp path"), Some(&prefixes)) + .await + .expect_err("missing configured prefix should fail"); + assert!(!error.is_client_error()); + assert!(error.to_string().contains(prefixes[0].to_str().expect("utf-8 prefix"))); + assert!(!error.public_message().contains(prefixes[0].to_str().expect("utf-8 prefix"))); + + fs::remove_dir_all(root).ok(); + } +} diff --git a/rust/src/server/src/routes.rs b/rust/src/server/src/routes.rs index b9157f6ec92c..3f0d16f56c32 100644 --- a/rust/src/server/src/routes.rs +++ b/rust/src/server/src/routes.rs @@ -39,7 +39,7 @@ fn server_dev_mode_enabled() -> bool { .is_some_and(|value| value != 0) } -fn runtime_lora_updating_enabled() -> bool { +pub(crate) fn runtime_lora_updating_enabled() -> bool { std::env::var("VLLM_ALLOW_RUNTIME_LORA_UPDATING") .ok() .is_some_and(|value| matches!(value.trim().to_lowercase().as_str(), "1" | "true")) diff --git a/rust/src/server/src/routes/inference/generate.rs b/rust/src/server/src/routes/inference/generate.rs index e176938db400..b13194df3af2 100644 --- a/rust/src/server/src/routes/inference/generate.rs +++ b/rust/src/server/src/routes/inference/generate.rs @@ -46,7 +46,13 @@ pub async fn generate( ValidatedJson(body): ValidatedJson, ) -> Response { let request_context = resolve_request_context(&headers, body.request_id.as_deref()); - let lora_resolution = state.resolve_model_with_loras(body.model.as_deref()).await; + let mut lora_resolution = state.resolve_model_with_loras(body.model.as_deref()).await; + if lora_resolution.lora_request.is_some() && !state.lora_state_is_consistent() { + return ApiError::server_error( + "LoRA state differs across engine ranks; restart the engine".to_string(), + ) + .into_response(); + } let prepared = match prepare_generate_request(body, &lora_resolution, request_context) { Ok(prepared) => prepared, Err(error) => return error.into_response(), @@ -73,6 +79,8 @@ pub async fn generate( } }; + let raw_stream = crate::lora::hold_lora_lease(raw_stream, lora_resolution.lease.take()); + if stream { let chunk_stream = generate_chunk_stream( raw_stream, @@ -406,6 +414,7 @@ mod tests { async fn generate_chunk_stream_captures_late_prompt_info() { let stream = stream::iter(vec![ Ok(GenerateOutput { + routed_experts: None, request_id: String::new(), prompt_info: None, token_ids: Vec::new(), @@ -416,6 +425,7 @@ mod tests { ec_transfer_params: None, }), Ok(GenerateOutput { + routed_experts: None, request_id: String::new(), prompt_info: Some(GeneratePromptInfo { prompt_token_ids: Arc::from([11_u32, 22_u32]), @@ -486,6 +496,7 @@ mod tests { token_ids: vec![3], logprobs: None, finish_reason: FinishReason::stop_eos(), + routed_experts: None, usage: vllm_llm::TokenUsage { prompt_token_count: prompt_token_ids.len(), output_token_count: 1, diff --git a/rust/src/server/src/routes/inference/generate/convert.rs b/rust/src/server/src/routes/inference/generate/convert.rs index 02a52b04146b..622a86028856 100644 --- a/rust/src/server/src/routes/inference/generate/convert.rs +++ b/rust/src/server/src/routes/inference/generate/convert.rs @@ -107,6 +107,7 @@ mod tests { LoraModelResolution { model_names: names.iter().map(|s| s.to_string()).collect(), lora_request: None, + lease: None, } } diff --git a/rust/src/server/src/routes/lora.rs b/rust/src/server/src/routes/lora.rs index dad817c69a28..c879e1a0ccd3 100644 --- a/rust/src/server/src/routes/lora.rs +++ b/rust/src/server/src/routes/lora.rs @@ -11,12 +11,11 @@ use validator::Validate; use crate::error::ApiError; use crate::lora::{LoadLoraError, UnloadLoraError}; +use crate::lora_path::{runtime_lora_allowed_path_prefixes, validate_lora_path_access}; use crate::routes::openai::utils::types::Normalizable; use crate::routes::openai::utils::validated_json::ValidatedJson; use crate::state::AppState; -const RUNTIME_LORA_ALLOWED_PATH_PREFIXES_ENV: &str = "VLLM_RUNTIME_LORA_ALLOWED_PATH_PREFIXES"; - #[derive(Debug, Deserialize, Validate)] pub(crate) struct LoadLoraAdapterRequest { lora_name: String, @@ -38,76 +37,6 @@ pub(crate) struct UnloadLoraAdapterRequest { impl Normalizable for UnloadLoraAdapterRequest {} -fn runtime_lora_allowed_path_prefixes() -> Option> { - let prefixes = std::env::var_os(RUNTIME_LORA_ALLOWED_PATH_PREFIXES_ENV)?; - let prefixes: Vec<_> = std::env::split_paths(&prefixes) - .filter(|path| !path.as_os_str().is_empty()) - .collect(); - (!prefixes.is_empty()).then_some(prefixes) -} - -fn looks_like_local_lora_path(lora_path: &str) -> bool { - let path = Path::new(lora_path); - path.is_absolute() - || lora_path.starts_with('~') - || lora_path.starts_with('.') - || path.components().any(|component| matches!(component, Component::ParentDir)) -} - -fn validate_lora_path_access( - lora_path: &str, - allowed_prefixes: Option<&[PathBuf]>, -) -> Result, ApiError> { - let path = Path::new(lora_path); - if !looks_like_local_lora_path(lora_path) && !path.exists() { - return Ok(None); - } - - let Some(allowed_prefixes) = allowed_prefixes else { - return Err(ApiError::invalid_request( - format!( - "Local LoRA adapter paths require {RUNTIME_LORA_ALLOWED_PATH_PREFIXES_ENV} to be configured." - ), - Some("lora_path"), - )); - }; - - if !path.is_absolute() { - return Err(ApiError::invalid_request( - format!( - "Local LoRA adapter paths must be absolute and under one of the prefixes configured by {RUNTIME_LORA_ALLOWED_PATH_PREFIXES_ENV}." - ), - Some("lora_path"), - )); - } - - let canonical_path = path.canonicalize().map_err(|_| { - ApiError::invalid_request( - "Local LoRA adapter path must exist and be accessible.".to_string(), - Some("lora_path"), - ) - })?; - let canonical_prefixes = allowed_prefixes - .iter() - .map(|prefix| { - prefix.canonicalize().map_err(|_| { - ApiError::server_error(format!( - "configured {RUNTIME_LORA_ALLOWED_PATH_PREFIXES_ENV} path prefix must exist and be accessible" - )) - }) - }) - .collect::, _>>()?; - - if !canonical_prefixes.iter().any(|prefix| canonical_path.starts_with(prefix)) { - return Err(ApiError::invalid_request( - "Local LoRA adapter path is outside the configured allowed prefixes.".to_string(), - Some("lora_path"), - )); - } - - Ok(Some(canonical_path.to_string_lossy().into_owned())) -} - /// Dynamically load one LoRA adapter and expose it as an OpenAI model id. pub async fn load_lora_adapter( State(state): State>, @@ -120,7 +49,17 @@ pub async fn load_lora_adapter( )); } let allowed_prefixes = runtime_lora_allowed_path_prefixes(); - let lora_path = validate_lora_path_access(&request.lora_path, allowed_prefixes.as_deref())? + let lora_path = validate_lora_path_access(&request.lora_path, allowed_prefixes.as_deref()) + .await + .map_err(|error| { + if error.is_client_error() { + ApiError::invalid_request(error.public_message().to_string(), Some("lora_path")) + } else { + tracing::error!(error = %error, "runtime LoRA path policy validation failed"); + ApiError::server_error(error.public_message().to_string()) + } + })? + .map(|path| path.to_string_lossy().into_owned()) .unwrap_or(request.lora_path); let lora_name = request.lora_name; @@ -133,6 +72,9 @@ pub async fn load_lora_adapter( ) .await .map_err(|error| match error { + LoadLoraError::Inconsistent => ApiError::server_error( + "LoRA state differs across engine ranks; restart the engine".to_string(), + ), LoadLoraError::AlreadyLoaded { lora_name } => ApiError::invalid_request( format!( "The lora adapter '{lora_name}' has already been loaded. If you want to load the adapter in place, set 'load_inplace' to true." @@ -173,6 +115,9 @@ pub async fn unload_lora_adapter( .unload_lora(&request.lora_name, request.lora_int_id) .await .map_err(|error| match error { + UnloadLoraError::Inconsistent => ApiError::server_error( + "LoRA state differs across engine ranks; restart the engine".to_string(), + ), UnloadLoraError::NotFound { lora_name } => ApiError::model_not_found(lora_name), UnloadLoraError::IntIdMismatch { lora_name, @@ -189,12 +134,6 @@ pub async fn unload_lora_adapter( request.lora_name, error.to_report_string() )), - UnloadLoraError::NotRemoved { - lora_name, - lora_int_id, - } => ApiError::server_error(format!( - "failed to unload LoRA adapter '{lora_name}' with id {lora_int_id}" - )), })?; Ok(format!( @@ -202,98 +141,3 @@ pub async fn unload_lora_adapter( lora_request.lora_name )) } - -#[cfg(test)] -mod tests { - use std::fs; - use std::path::PathBuf; - use std::time::{SystemTime, UNIX_EPOCH}; - - use super::validate_lora_path_access; - - fn temp_lora_dir(test_name: &str) -> PathBuf { - let suffix = SystemTime::now() - .duration_since(UNIX_EPOCH) - .expect("clock should be after unix epoch") - .as_nanos(); - let path = std::env::temp_dir().join(format!( - "vllm-lora-{test_name}-{}-{suffix}", - std::process::id() - )); - fs::create_dir_all(&path).expect("create temp lora dir"); - path - } - - #[test] - fn lora_path_allows_hf_repo_ids_without_prefixes() { - assert_eq!( - validate_lora_path_access("org/adapter-a", None).expect("hf repo id should be allowed"), - None - ); - } - - #[test] - fn lora_path_rejects_local_paths_without_prefixes() { - assert!(validate_lora_path_access("/tmp/adapter-a", None).is_err()); - assert!(validate_lora_path_access("./adapter-a", None).is_err()); - assert!(validate_lora_path_access("~/adapter-a", None).is_err()); - assert!(validate_lora_path_access("subdir/../../../etc/sensitive", None).is_err()); - } - - #[test] - fn lora_path_rejects_existing_bare_relative_paths_without_prefixes() { - let root = - PathBuf::from("target").join(format!("vllm-lora-relative-{}", std::process::id())); - let adapter = root.join("adapter-a"); - fs::create_dir_all(&adapter).expect("create relative adapter dir"); - - assert!( - validate_lora_path_access(adapter.to_str().expect("utf-8 temp path"), None).is_err() - ); - - fs::remove_dir_all(root).ok(); - } - - #[test] - fn lora_path_allows_absolute_paths_under_configured_prefixes() { - let root = temp_lora_dir("allowed-prefix"); - let allowed = root.join("allowed"); - let adapter = allowed.join("adapter-a"); - fs::create_dir_all(&adapter).expect("create adapter dir"); - - let prefixes = [allowed]; - let resolved = - validate_lora_path_access(adapter.to_str().expect("utf-8 temp path"), Some(&prefixes)) - .expect("path under configured prefix should be allowed"); - assert_eq!( - resolved.as_deref(), - Some( - adapter - .canonicalize() - .expect("canonical adapter") - .to_str() - .expect("utf-8 temp path") - ) - ); - - fs::remove_dir_all(root).ok(); - } - - #[test] - fn lora_path_rejects_parent_escape_from_configured_prefixes() { - let root = temp_lora_dir("parent-escape"); - let allowed = root.join("allowed"); - let private_adapter = root.join("private").join("adapter-a"); - fs::create_dir_all(&allowed).expect("create allowed dir"); - fs::create_dir_all(&private_adapter).expect("create private adapter dir"); - - let escaped = allowed.join("../private/adapter-a"); - let prefixes = [allowed]; - assert!( - validate_lora_path_access(escaped.to_str().expect("utf-8 temp path"), Some(&prefixes)) - .is_err() - ); - - fs::remove_dir_all(root).ok(); - } -} diff --git a/rust/src/server/src/routes/openai/chat_completions.rs b/rust/src/server/src/routes/openai/chat_completions.rs index 8a1e9c4383a1..fa562f04c1cb 100644 --- a/rust/src/server/src/routes/openai/chat_completions.rs +++ b/rust/src/server/src/routes/openai/chat_completions.rs @@ -54,7 +54,13 @@ pub async fn chat_completions( ) -> Response { let stream = body.stream; let request_context = resolve_request_context(&headers, body.request_id.as_deref()); - let lora_resolution = state.resolve_model_with_loras(Some(&body.model)).await; + let mut lora_resolution = state.resolve_model_with_loras(Some(&body.model)).await; + if lora_resolution.lora_request.is_some() && !state.lora_state_is_consistent() { + return ApiError::server_error( + "LoRA state differs across engine ranks; restart the engine".to_string(), + ) + .into_response(); + } let prepared = match prepare_chat_request(body, &lora_resolution, request_context) { Ok(prepared) => prepared, @@ -78,6 +84,7 @@ pub async fn chat_completions( }; if stream { + let chat_stream = crate::lora::hold_lora_lease(chat_stream, lora_resolution.lease.take()); let chunk_stream = chat_completion_chunk_stream( chat_stream, prepared.request_id, @@ -90,6 +97,7 @@ pub async fn chat_completions( Sse::new(sse_stream).into_response() } else { + let _lora_lease = lora_resolution.lease.take(); let response = match collect_chat_completion( chat_stream, prepared.request_id, diff --git a/rust/src/server/src/routes/openai/chat_completions/convert.rs b/rust/src/server/src/routes/openai/chat_completions/convert.rs index 2c4050f8f1e3..2f980087461a 100644 --- a/rust/src/server/src/routes/openai/chat_completions/convert.rs +++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs @@ -128,6 +128,7 @@ pub(super) fn prepare_chat_request( request_id: request_id.clone(), messages, sampling_params: SamplingParams { + routed_experts_prompt_start: 0, temperature: request.temperature, top_p: request.top_p, top_k: request.top_k, @@ -444,6 +445,7 @@ mod tests { LoraModelResolution { model_names: names.iter().map(|s| s.to_string()).collect(), lora_request: None, + lease: None, } } diff --git a/rust/src/server/src/routes/openai/completions.rs b/rust/src/server/src/routes/openai/completions.rs index b6f2138f1278..137775e12367 100644 --- a/rust/src/server/src/routes/openai/completions.rs +++ b/rust/src/server/src/routes/openai/completions.rs @@ -53,7 +53,13 @@ pub async fn completions( ) -> Response { let stream = body.stream; let request_context = resolve_request_context(&headers, body.request_id.as_deref()); - let lora_resolution = state.resolve_model_with_loras(Some(&body.model)).await; + let mut lora_resolution = state.resolve_model_with_loras(Some(&body.model)).await; + if lora_resolution.lora_request.is_some() && !state.lora_state_is_consistent() { + return ApiError::server_error( + "LoRA state differs across engine ranks; restart the engine".to_string(), + ) + .into_response(); + } let tokenizer = state.chat.text().tokenizer(); let prepared = match prepare_completion_request( @@ -86,6 +92,8 @@ pub async fn completions( } }; + let text_stream = crate::lora::hold_lora_lease(text_stream, lora_resolution.lease.take()); + if stream { let chunk_stream = completion_chunk_stream( text_stream, @@ -665,6 +673,7 @@ mod tests { ))), kv_transfer_params: None, ec_transfer_params: None, + routed_experts: None, }), }), ]); @@ -770,6 +779,7 @@ mod tests { finish_reason: FinishReason::Length, kv_transfer_params: None, ec_transfer_params: None, + routed_experts: None, }), }), ]); @@ -822,6 +832,7 @@ mod tests { finish_reason: FinishReason::Length, kv_transfer_params: None, ec_transfer_params: None, + routed_experts: None, }), }), ]); @@ -877,6 +888,7 @@ mod tests { finish_reason: FinishReason::Length, kv_transfer_params: None, ec_transfer_params: None, + routed_experts: None, }), }), ]); @@ -949,6 +961,7 @@ mod tests { finish_reason: FinishReason::Length, kv_transfer_params: None, ec_transfer_params: None, + routed_experts: None, }), }), ]); @@ -1023,6 +1036,7 @@ mod tests { finish_reason: FinishReason::Length, kv_transfer_params: None, ec_transfer_params: None, + routed_experts: None, }), }), ]); diff --git a/rust/src/server/src/routes/openai/completions/convert.rs b/rust/src/server/src/routes/openai/completions/convert.rs index 2030a4ad4084..1b66f12458ca 100644 --- a/rust/src/server/src/routes/openai/completions/convert.rs +++ b/rust/src/server/src/routes/openai/completions/convert.rs @@ -109,6 +109,7 @@ pub(super) fn prepare_completion_request( prompt: request.prompt, mm_features: None, sampling_params: SamplingParams { + routed_experts_prompt_start: 0, temperature: request.temperature, top_p: request.top_p, top_k: request.top_k, @@ -217,6 +218,7 @@ mod tests { LoraModelResolution { model_names: names.iter().map(|s| s.to_string()).collect(), lora_request: None, + lease: None, } } diff --git a/rust/src/server/src/routes/openai/models.rs b/rust/src/server/src/routes/openai/models.rs index 2a0f092be5cd..66f4e25b654b 100644 --- a/rust/src/server/src/routes/openai/models.rs +++ b/rust/src/server/src/routes/openai/models.rs @@ -31,7 +31,12 @@ pub async fn list_models(State(state): State>) -> Json = json["data"] + .as_array() + .expect("model data") + .iter() + .map(|model| model["id"].as_str().expect("model id")) + .collect(); + assert_eq!(model_ids, ["Qwen/Qwen1.5-0.5B-Chat"]); + + drop(app); + engine_task.finish().await; +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn load_lora_adapter_rejects_base_model_name_collision() { diff --git a/rust/src/server/src/routes/tests/lora_concurrency.rs b/rust/src/server/src/routes/tests/lora_concurrency.rs new file mode 100644 index 000000000000..24a895809c57 --- /dev/null +++ b/rust/src/server/src/routes/tests/lora_concurrency.rs @@ -0,0 +1,306 @@ +use super::*; + +async fn post_load(app: &mut axum::Router, path: &str, load_inplace: bool) -> StatusCode { + app.call( + Request::builder() + .method("POST") + .uri("/v1/load_lora_adapter") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "lora_name": "adapter-a", + "lora_path": path, + "load_inplace": load_inplace + }) + .to_string(), + )) + .expect("build load request"), + ) + .await + .expect("call load route") + .status() +} + +async fn post_completion(app: &mut axum::Router) -> StatusCode { + app.call( + Request::builder() + .method("POST") + .uri("/v1/completions") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "adapter-a", + "prompt": "hello", + "max_tokens": 1, + "stream": false + }) + .to_string(), + )) + .expect("build completion request"), + ) + .await + .expect("call completion route") + .status() +} + +async fn post_unload(app: &mut axum::Router) -> StatusCode { + app.call( + Request::builder() + .method("POST") + .uri("/v1/unload_lora_adapter") + .header("content-type", "application/json") + .body(Body::from(json!({"lora_name": "adapter-a"}).to_string())) + .expect("build unload request"), + ) + .await + .expect("call unload route") + .status() +} + +async fn receive_lora_utility( + dealer: &mut DealerSocket, + expected_path: &str, + expected_inplace: bool, +) -> u64 { + 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"); + assert_eq!(array[2], Value::from("add_lora")); + let lora = array[3].as_array().unwrap()[0].as_array().unwrap(); + assert_eq!(lora[0], Value::from("adapter-a")); + assert_eq!(lora[1], Value::from(1)); + assert_eq!(lora[2], Value::from(expected_path)); + assert_eq!(lora[5], Value::from(expected_inplace)); + array[1].as_u64().expect("call id") +} + +async fn reply_utility(push: &mut PushSocket, call_id: u64, result: bool) { + send_outputs(push, utility_outputs(call_id, utility_result_value(result))).await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn successful_replacement_waits_for_active_generation() { + let (generation_started_tx, generation_started_rx) = tokio::sync::oneshot::channel(); + let (release_generation_tx, release_generation_rx) = tokio::sync::oneshot::channel(); + let (mut app, engine_task) = test_admin_app_with_engine_script(move |dealer, push| { + boxed_test_future(async move { + let initial = receive_lora_utility(dealer, "org/adapter-a", false).await; + reply_utility(push, initial, true).await; + + let add = recv_engine_message(dealer).await; + assert_eq!(add[0].as_ref(), &[0x00]); + let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap(); + assert_adapter_a_lora_request(&request); + let _ = generation_started_tx.send(()); + + tokio::select! { + result = release_generation_rx => result.expect("release generation"), + _early = recv_engine_message(dealer) => { + panic!("replacement reached the engine before generation completed") + } + } + send_outputs( + push, + engine_outputs_for_request(&request.request_id, default_stream_output_specs()), + ) + .await; + + let replacement = receive_lora_utility(dealer, "org/adapter-b", true).await; + reply_utility(push, replacement, true).await; + }) + }) + .await; + assert_eq!( + post_load(&mut app, "org/adapter-a", false).await, + StatusCode::OK + ); + + let mut generation_app = app.clone(); + let generation = tokio::spawn(async move { post_completion(&mut generation_app).await }); + generation_started_rx.await.unwrap(); + + let mut replacement_app = app.clone(); + let mut replacement = + tokio::spawn(async move { post_load(&mut replacement_app, "org/adapter-b", true).await }); + assert!(tokio::time::timeout(Duration::from_millis(50), &mut replacement).await.is_err()); + release_generation_tx.send(()).unwrap(); + + assert_eq!(generation.await.unwrap(), StatusCode::OK); + assert_eq!(replacement.await.unwrap(), StatusCode::OK); + drop(app); + engine_task.finish().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn rolled_back_replacement_hides_intermediate_adapter_from_generation() { + let (replacement_started_tx, replacement_started_rx) = tokio::sync::oneshot::channel(); + let (fail_replacement_tx, fail_replacement_rx) = tokio::sync::oneshot::channel(); + let (mut app, engine_task) = test_admin_app_with_engine_script(move |dealer, push| { + boxed_test_future(async move { + let initial = receive_lora_utility(dealer, "org/adapter-a", false).await; + reply_utility(push, initial, true).await; + + let replacement = receive_lora_utility(dealer, "org/adapter-b", true).await; + let _ = replacement_started_tx.send(()); + tokio::select! { + result = fail_replacement_rx => result.expect("fail replacement"), + _early = recv_engine_message(dealer) => { + panic!("generation reached the engine during replacement") + } + } + reply_utility(push, replacement, false).await; + + let restore = receive_lora_utility(dealer, "org/adapter-a", true).await; + reply_utility(push, restore, true).await; + + let add = recv_engine_message(dealer).await; + assert_eq!(add[0].as_ref(), &[0x00]); + let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap(); + assert_adapter_a_lora_request(&request); + send_outputs( + push, + engine_outputs_for_request(&request.request_id, default_stream_output_specs()), + ) + .await; + }) + }) + .await; + assert_eq!( + post_load(&mut app, "org/adapter-a", false).await, + StatusCode::OK + ); + + let mut replacement_app = app.clone(); + let replacement = + tokio::spawn(async move { post_load(&mut replacement_app, "org/adapter-b", true).await }); + replacement_started_rx.await.unwrap(); + + let mut generation_app = app.clone(); + let mut generation = tokio::spawn(async move { post_completion(&mut generation_app).await }); + assert!(tokio::time::timeout(Duration::from_millis(50), &mut generation).await.is_err()); + fail_replacement_tx.send(()).unwrap(); + + assert_eq!( + replacement.await.unwrap(), + StatusCode::INTERNAL_SERVER_ERROR + ); + assert_eq!(generation.await.unwrap(), StatusCode::OK); + drop(app); + engine_task.finish().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn timed_out_load_runs_bounded_compensation() { + let (mut app, engine_task) = test_admin_app_with_engine_script(|dealer, push| { + boxed_test_future(async move { + let _load = recv_engine_message(dealer).await; + let remove = recv_engine_message(dealer).await; + let payload = decode_value(&remove[1]).expect("decode compensation"); + let array = payload.as_array().expect("utility payload"); + assert_eq!(array[2], Value::from("remove_lora")); + reply_utility(push, array[1].as_u64().unwrap(), true).await; + }) + }) + .await; + + assert_eq!( + post_load(&mut app, "org/adapter-a", false).await, + StatusCode::INTERNAL_SERVER_ERROR + ); + let models = app + .call(Request::builder().uri("/v1/models").body(Body::empty()).unwrap()) + .await + .unwrap(); + let body = to_bytes(models.into_body(), usize::MAX).await.unwrap(); + let json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(json["data"].as_array().unwrap().len(), 1); + drop(app); + engine_task.finish().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn cancelled_load_marks_registry_inconsistent() { + let dispatched = std::sync::Arc::new(tokio::sync::Notify::new()); + let engine_notice = dispatched.clone(); + let (app, engine_task) = test_admin_app_with_engine_script(move |dealer, _push| { + let engine_notice = engine_notice.clone(); + boxed_test_future(async move { + let _load = recv_engine_message(dealer).await; + engine_notice.notify_one(); + std::future::pending::<()>().await; + }) + }) + .await; + + let mut loading_app = app.clone(); + let load = + tokio::spawn(async move { post_load(&mut loading_app, "org/adapter-a", false).await }); + dispatched.notified().await; + load.abort(); + assert!(load.await.unwrap_err().is_cancelled()); + + let mut retry_app = app.clone(); + assert_eq!( + post_load(&mut retry_app, "org/adapter-a", false).await, + StatusCode::INTERNAL_SERVER_ERROR + ); + drop(app); + engine_task.abort_and_join().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn failed_unload_restores_adapter() { + let (mut app, engine_task) = test_admin_app_with_engine_script(|dealer, push| { + boxed_test_future(async move { + let load = receive_lora_utility(dealer, "org/adapter-a", false).await; + reply_utility(push, load, true).await; + + let unload = recv_engine_message(dealer).await; + let payload = decode_value(&unload[1]).expect("decode unload"); + let array = payload.as_array().expect("utility payload"); + assert_eq!(array[2], Value::from("remove_lora")); + send_outputs( + push, + UtilityCallOutput { + engine_index: 0, + timestamp: 0.0, + output: UtilityOutput { + call_id: array[1].as_u64().unwrap().into(), + failure_message: Some("rank failed".to_string()), + result: None, + }, + } + .into(), + ) + .await; + + let restore = receive_lora_utility(dealer, "org/adapter-a", false).await; + reply_utility(push, restore, true).await; + }) + }) + .await; + + assert_eq!( + post_load(&mut app, "org/adapter-a", false).await, + StatusCode::OK + ); + assert_eq!( + post_unload(&mut app).await, + StatusCode::INTERNAL_SERVER_ERROR + ); + let models = app + .call(Request::builder().uri("/v1/models").body(Body::empty()).unwrap()) + .await + .unwrap(); + let body = to_bytes(models.into_body(), usize::MAX).await.unwrap(); + let json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(json["data"][1]["id"], "adapter-a"); + drop(app); + engine_task.finish().await; +} diff --git a/rust/src/server/src/state.rs b/rust/src/server/src/state.rs index 722dede8dda0..8fac930ada98 100644 --- a/rust/src/server/src/state.rs +++ b/rust/src/server/src/state.rs @@ -15,7 +15,9 @@ 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::lora::{ + LoadExactLoraError, LoadLoraError, LoraManager, LoraModelResolution, UnloadLoraError, +}; use crate::runtime::build_request_runtime; use crate::server_info::{ServerInfoConfigFormat, ServerInfoSnapshot}; @@ -161,6 +163,11 @@ impl AppState { self.lora_manager.served_lora_requests().await } + /// Whether the dynamic LoRA registry is known to match every engine rank. + pub(crate) fn lora_state_is_consistent(&self) -> bool { + self.lora_manager.is_consistent() + } + /// Resolve the requested model against one dynamic LoRA registry snapshot. pub async fn resolve_model_with_loras(&self, model_name: Option<&str>) -> LoraModelResolution { self.lora_manager.resolve_model(&self.served_model_names, model_name).await @@ -186,6 +193,22 @@ impl AppState { .await } + /// Load one dynamic adapter with an externally assigned ID. + pub async fn load_lora_exact( + &self, + lora_request: LoraRequest, + load_inplace: bool, + ) -> Result<(LoraRequest, bool), LoadExactLoraError> { + self.lora_manager + .load_lora_exact( + self.engine_core_client(), + &self.served_model_names, + lora_request, + load_inplace, + ) + .await + } + /// Remove one dynamic LoRA adapter from the engine and public model /// registry. pub async fn unload_lora( diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs index aa54595acdaa..b5c9b2f4e53d 100644 --- a/rust/src/text/src/lower.rs +++ b/rust/src/text/src/lower.rs @@ -4,11 +4,9 @@ use std::collections::BTreeSet; pub(crate) mod logprobs; -pub(crate) mod sampling; pub(crate) mod token_ids; use logprobs::validate_logprobs; -use sampling::validate_resolved_sampling_params; use token_ids::{validate_prompt_token_ids, validate_vocab_range}; use vllm_engine_core_client::protocol::sampling::{ EngineCoreSamplingParams, RepetitionDetectionParams, @@ -97,6 +95,7 @@ pub fn lower_sampling_params( max_tokens, min_tokens, thinking_token_budget, + routed_experts_prompt_start, logprobs, prompt_logprobs, min_p, @@ -169,6 +168,7 @@ pub fn lower_sampling_params( max_tokens, min_tokens, thinking_token_budget, + routed_experts_prompt_start, logprobs, prompt_logprobs, min_p, @@ -188,7 +188,6 @@ pub fn lower_sampling_params( skip_reading_prefix_cache, extra_args: vllm_xargs, }; - validate_resolved_sampling_params(¶ms)?; validate_vocab_range(¶ms, &sampling_limits)?; Ok(params) } @@ -322,7 +321,7 @@ mod tests { use super::*; use crate::backend::hf::HfTextBackend; use crate::backend::{SamplingHints, TextBackend as _}; - use crate::error::{LogprobsError, SamplingParamsError, TokenIdsError}; + use crate::error::{LogprobsError, TokenIdsError}; use crate::request::{Prompt, TextRequest}; fn stub_tokenizer() -> TestTokenizer { @@ -486,117 +485,17 @@ mod tests { } #[test] - fn lower_sampling_params_rejects_invalid_sampling_ranges() { - let cases = [ - ( - "temperature", - SamplingParams { - temperature: Some(5.0), - ..SamplingParams::default() - }, - ), - ( - "top_p", - SamplingParams { - top_p: Some(0.0), - ..SamplingParams::default() - }, - ), - ( - "min_p", - SamplingParams { - min_p: Some(2.0), - ..SamplingParams::default() - }, - ), - ( - "repetition_penalty", - SamplingParams { - repetition_penalty: Some(0.0), - ..SamplingParams::default() - }, - ), - ( - "frequency_penalty", - SamplingParams { - frequency_penalty: Some(100.0), - ..SamplingParams::default() - }, - ), - ( - "presence_penalty", - SamplingParams { - presence_penalty: Some(100.0), - ..SamplingParams::default() - }, - ), - ]; - - for (expected_parameter, sampling_params) in cases { - let error = - lower_sampling_params_with_limits(sampling_params, sample_sampling_limits()) - .unwrap_err(); - - assert!( - matches!( - error, - Error::SamplingParams(SamplingParamsError::OutOfRange { - parameter, - .. - }) if parameter == expected_parameter - ), - "{expected_parameter} should be rejected" - ); - } - } - - #[test] - fn lower_sampling_params_rejects_non_finite_sampling_values() { - for (expected_parameter, sampling_params) in [ - ( - "temperature", - SamplingParams { - temperature: Some(f32::INFINITY), - ..SamplingParams::default() - }, - ), - ( - "repetition_penalty", - SamplingParams { - repetition_penalty: Some(f32::NAN), - ..SamplingParams::default() - }, - ), - ] { - let error = - lower_sampling_params_with_limits(sampling_params, sample_sampling_limits()) - .unwrap_err(); - - assert!( - matches!( - error, - Error::SamplingParams(SamplingParamsError::NotFinite { - parameter, - .. - }) if parameter == expected_parameter - ), - "{expected_parameter} should reject non-finite values" - ); - } - } - - #[test] - fn lower_sampling_params_accepts_python_compatible_repetition_penalty_above_two() { - let params = lower_sampling_params_with_limits( + fn lower_sampling_params_forwards_routed_experts_prompt_start() { + let lowered = lower_sampling_params_with_limits( SamplingParams { - repetition_penalty: Some(2.5), + routed_experts_prompt_start: 23, ..SamplingParams::default() }, sample_sampling_limits(), ) .unwrap(); - assert_eq!(params.repetition_penalty, 2.5); + assert_eq!(lowered.routed_experts_prompt_start, 23); } #[test] @@ -620,6 +519,7 @@ mod tests { max_tokens: 999997, min_tokens: 0, thinking_token_budget: None, + routed_experts_prompt_start: 0, logprobs: None, prompt_logprobs: None, min_p: 0.0, @@ -673,6 +573,7 @@ mod tests { max_tokens: 999997, min_tokens: 0, thinking_token_budget: None, + routed_experts_prompt_start: 0, logprobs: None, prompt_logprobs: None, min_p: 0.0, @@ -834,6 +735,7 @@ mod tests { max_tokens: 40957, min_tokens: 0, thinking_token_budget: None, + routed_experts_prompt_start: 0, logprobs: None, prompt_logprobs: None, min_p: 0.0, @@ -897,6 +799,7 @@ mod tests { max_tokens: 999997, min_tokens: 0, thinking_token_budget: None, + routed_experts_prompt_start: 0, logprobs: None, prompt_logprobs: None, min_p: 0.0, @@ -968,6 +871,7 @@ mod tests { max_tokens: 32, min_tokens: 2, thinking_token_budget: None, + routed_experts_prompt_start: 0, logprobs: None, prompt_logprobs: None, min_p: 0.1, @@ -1217,6 +1121,7 @@ mod tests { max_tokens: 128, min_tokens: 0, thinking_token_budget: None, + routed_experts_prompt_start: 0, logprobs: None, prompt_logprobs: None, min_p: 0.1, diff --git a/rust/src/text/src/output/decoded.rs b/rust/src/text/src/output/decoded.rs index cb0fe9c18e7f..e4de75259524 100644 --- a/rust/src/text/src/output/decoded.rs +++ b/rust/src/text/src/output/decoded.rs @@ -8,7 +8,8 @@ use futures::{Stream, StreamExt}; use serde::{Deserialize, Serialize}; use tracing::{Level, debug, trace}; use vllm_engine_core_client::AbortCause; -use vllm_engine_core_client::protocol::output::StopReason; +use vllm_engine_core_client::protocol::output::{StopReason, concatenate_routed_experts}; +use vllm_engine_core_client::protocol::tensor::WireNdArray; use vllm_llm::{FinishReason, GenerateOutput, TokenUsage}; use vllm_tokenizer::{DynTokenizer, IncrementalDecoder}; @@ -50,6 +51,8 @@ pub struct Finished { /// Connector-specific encoder cache transfer parameters for disaggregated /// serving. pub ec_transfer_params: Option, + /// Routed-expert IDs concatenated along the token axis. + pub routed_experts: Option, } /// Internal decoded-text event emitted before higher-level assistant @@ -107,9 +110,13 @@ pub async fn decoded_text_event_stream( let mut token_ids = Vec::new(); let mut output_token_count: usize = 0; let mut logprobs: Option = None; + let mut routed_experts_chunks = Vec::new(); while let Some(next) = raw_stream.next().await { - let output = next?; + let mut output = next?; + if let Some(routed_experts) = output.routed_experts.take() { + routed_experts_chunks.push(routed_experts); + } cached_token_count = cached_token_count.max(output.cached_token_count); // If it's the first output, init states and yield `Start` event. @@ -283,6 +290,7 @@ pub async fn decoded_text_event_stream( finish_reason: reason, kv_transfer_params, ec_transfer_params, + routed_experts: concatenate_routed_experts(routed_experts_chunks)?, }), }) .await; diff --git a/rust/src/text/src/output/mod.rs b/rust/src/text/src/output/mod.rs index c3499a9f3539..3a2532ea633d 100644 --- a/rust/src/text/src/output/mod.rs +++ b/rust/src/text/src/output/mod.rs @@ -14,6 +14,7 @@ mod logprobs; use std::sync::Arc; use futures::{StreamExt as _, pin_mut}; +use vllm_engine_core_client::protocol::tensor::WireNdArray; use crate::{Error, FinishReason, Result, TextOutputStream}; @@ -32,6 +33,8 @@ pub struct CollectedTextOutput { /// Connector-specific encoder cache transfer parameters for disaggregated /// serving. pub ec_transfer_params: Option, + /// Routed-expert IDs concatenated along the token axis. + pub routed_experts: Option, } #[allow(clippy::manual_async_fn, reason = "specify `Send` bound")] @@ -84,6 +87,7 @@ impl T { usage: vllm_llm::TokenUsage::default(), kv_transfer_params: None, ec_transfer_params: None, + routed_experts: None, }) }; @@ -93,6 +97,7 @@ impl T { collected.usage = finished.usage; collected.kv_transfer_params = finished.kv_transfer_params; collected.ec_transfer_params = finished.ec_transfer_params; + collected.routed_experts = finished.routed_experts; return Ok(collected); } } @@ -165,6 +170,7 @@ mod tests { finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, + routed_experts: None, }), }), ]); @@ -283,6 +289,7 @@ mod tests { finish_reason: FinishReason::stop_eos(), kv_transfer_params: None, ec_transfer_params: None, + routed_experts: None, }), }), ]); diff --git a/rust/src/text/src/request.rs b/rust/src/text/src/request.rs index c12da81c712a..774d4c831a12 100644 --- a/rust/src/text/src/request.rs +++ b/rust/src/text/src/request.rs @@ -67,6 +67,8 @@ pub struct SamplingParams { /// here; `-1` is normalized to `None` (and other negatives rejected) during /// lowering (see `lower_sampling_params`). pub thinking_token_budget: Option, + /// Number of prompt-token rows to omit from returned routed-expert data. + pub routed_experts_prompt_start: u32, /// Number of log probabilities to return per generated token. /// /// `None` disables sample logprobs. `-1` requests the full vocabulary. @@ -131,6 +133,7 @@ impl Default for SamplingParams { max_tokens: None, min_tokens: None, thinking_token_budget: None, + routed_experts_prompt_start: 0, logprobs: None, prompt_logprobs: None, min_p: None, diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py b/tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py index d24436bbd4bd..c555018243a0 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_mm_serde.py @@ -112,3 +112,16 @@ def test_mm_features_with_kwargs_data(): decoded = decode_mm_kwargs_item(features2.kwargs_data["image"][0]) assert torch.equal(elem.data, decoded["pixel_values"].data) + + +def test_legacy_disagg_mm_serde_import_is_compatible(): + """Existing renderer clients can follow the module's scale-out move.""" + from vllm.entrypoints.serve.disagg.mm_serde import ( + decode_mm_kwargs_item as legacy_decode, + ) + from vllm.entrypoints.serve.disagg.mm_serde import ( + encode_mm_kwargs_item as legacy_encode, + ) + + assert legacy_encode is encode_mm_kwargs_item + assert legacy_decode is decode_mm_kwargs_item diff --git a/tests/v1/engine/test_engine_core.py b/tests/v1/engine/test_engine_core.py index aa2a70559ddd..ee451f08e3f9 100644 --- a/tests/v1/engine/test_engine_core.py +++ b/tests/v1/engine/test_engine_core.py @@ -23,7 +23,7 @@ from vllm.platforms import current_platform from vllm.utils.torch_utils import set_default_torch_num_threads from vllm.v1.engine import EngineCoreRequest -from vllm.v1.engine.core import EngineCore +from vllm.v1.engine.core import EngineCore, select_kv_event_block_size from vllm.v1.executor.abstract import Executor from vllm.v1.executor.uniproc_executor import UniProcExecutor from vllm.v1.kv_cache_interface import KVCacheConfig @@ -44,6 +44,47 @@ _REQUEST_COUNTER = 0 +@pytest.mark.parametrize( + ("metadata", "fallback", "expected"), + [ + ( + [ + {"group_idx": 0, "kind": "sliding_window", "block_size": 4}, + {"group_idx": 1, "kind": "mla_attention", "block_size": 256}, + ], + 4, + 256, + ), + ( + [{"group_idx": 0, "kind": "full_attention", "block_size": 16}], + 16, + 16, + ), + ( + [ + {"group_idx": 0, "kind": "mamba", "block_size": 8}, + { + "group_idx": 1, + "kind": "sink_full_attention", + "block_size": 128, + }, + ], + 8, + 128, + ), + ([{"group_idx": 0, "kind": "mamba", "block_size": 8}], 32, 32), + ([], 16, 16), + ( + [{"group_idx": 0, "kind": "full_attention", "block_size": None}], + 64, + 64, + ), + ], +) +def test_select_kv_event_block_size(metadata, fallback, expected): + assert select_kv_event_block_size(metadata, fallback) == expected + + def make_request() -> EngineCoreRequest: global _REQUEST_COUNTER _REQUEST_COUNTER += 1 diff --git a/tests/v1/engine/test_engine_core_client.py b/tests/v1/engine/test_engine_core_client.py index 64adf7a8b3c8..a8c927788ae1 100644 --- a/tests/v1/engine/test_engine_core_client.py +++ b/tests/v1/engine/test_engine_core_client.py @@ -335,6 +335,9 @@ def test_apply_ready_response_syncs_block_size(): max_num_seqs=256, max_num_batched_tokens=8192, instance_id="test-instance", + kv_event_block_size=256, + supports_lora=False, + max_loras=0, ) ) client._apply_ready_response(payload) diff --git a/vllm/entrypoints/serve/disagg/__init__.py b/vllm/entrypoints/serve/disagg/__init__.py new file mode 100644 index 000000000000..96c6a336ee01 --- /dev/null +++ b/vllm/entrypoints/serve/disagg/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compatibility imports for the former disaggregated serving package.""" diff --git a/vllm/entrypoints/serve/disagg/mm_serde.py b/vllm/entrypoints/serve/disagg/mm_serde.py new file mode 100644 index 000000000000..b9bd6131afd8 --- /dev/null +++ b/vllm/entrypoints/serve/disagg/mm_serde.py @@ -0,0 +1,10 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Compatibility alias for multimodal serialization helpers.""" + +from vllm.entrypoints.scale_out.token_in_token_out.mm_serde import ( + decode_mm_kwargs_item, + encode_mm_kwargs_item, +) + +__all__ = ["decode_mm_kwargs_item", "encode_mm_kwargs_item"] diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py index 514282080221..b0a2f294bcfe 100644 --- a/vllm/v1/engine/__init__.py +++ b/vllm/v1/engine/__init__.py @@ -11,7 +11,6 @@ import numpy as np import torch -from vllm.config.kv_events import KVEventsConfig from vllm.lora.request import LoRARequest from vllm.multimodal.inputs import MultiModalFeatureSpec from vllm.pooling_params import PoolingParams @@ -36,6 +35,8 @@ class EEPNotificationType(enum.Enum): + NEW_CORE_ENGINES_INIT_READY = "NEW_CORE_ENGINES_INIT_READY" + NEW_CORE_ENGINES_WEIGHTS_INIT_READY = "NEW_CORE_ENGINES_WEIGHTS_INIT_READY" RECONFIGURE_FINISHED = "RECONFIGURE_FINISHED" SHUTDOWN_COMPLETE = "SHUTDOWN_COMPLETE" @@ -88,10 +89,16 @@ class EngineCoreReadyResponse: max_num_seqs: int max_num_batched_tokens: int instance_id: str + kv_event_block_size: int + supports_lora: bool + max_loras: int # KV cache capacity (None for encoder-only/attention-free models). kv_cache_size_tokens: int | None = None kv_cache_max_concurrency: float | None = None - kv_events_config: KVEventsConfig | None = None + kv_role: str | None = None + kv_events_publisher: str | None = None + kv_events_endpoint: str | None = None + kv_events_topic: str | None = None class EngineCoreRequest( diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index 11c78c7e7453..7cb5360c4631 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -7,7 +7,7 @@ import threading import time from collections import defaultdict, deque -from collections.abc import Callable, Generator, Sequence +from collections.abc import Callable, Generator from concurrent.futures import Future from contextlib import ExitStack, contextmanager from enum import IntEnum @@ -83,11 +83,15 @@ EngineCoreSentinel, fault_tolerant_wrapper, ) -from vllm.v1.kv_cache_interface import KVCacheConfig, get_kv_cache_spec_kind +from vllm.v1.kv_cache_interface import ( + KVCacheConfig, + KVCacheSpecKind, + get_kv_cache_spec_kind, +) from vllm.v1.metrics.stats import SchedulerIterationDetails, SchedulerStats from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import Request, RequestStatus -from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder, bytestr +from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder from vllm.v1.structured_output import StructuredOutputManager from vllm.v1.utils import compute_iteration_details from vllm.version import __version__ as VLLM_VERSION @@ -99,6 +103,26 @@ _R = TypeVar("_R") # Return type for collective_rpc +_MAIN_ATTENTION_KV_CACHE_KINDS = frozenset( + ( + KVCacheSpecKind.FULL_ATTENTION.value, + KVCacheSpecKind.MLA_ATTENTION.value, + KVCacheSpecKind.SINK_FULL_ATTENTION.value, + ) +) + + +def select_kv_event_block_size( + group_metadata: list[dict[str, int | str | None]], + fallback_block_size: int, +) -> int: + """Select the main-attention block size used by published KV events.""" + for group in group_metadata: + if group.get("kind") in _MAIN_ATTENTION_KV_CACHE_KINDS: + block_size = group.get("block_size") + return int(block_size) if block_size else fallback_block_size + return fallback_block_size + class EngineCore: """Inner loop of vLLM's Engine.""" @@ -327,9 +351,8 @@ def _initialize_kv_caches(self, vllm_config: VllmConfig) -> KVCacheConfig: vllm_config.validate_block_size() + # Initialize kv cache and warmup the execution self.model_executor.initialize_from_config(kv_cache_configs) - if not envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: - self.model_executor.compile_or_warm_up_model() elapsed = time.time() - start compile_time = vllm_config.compilation_config.compilation_time @@ -436,6 +459,13 @@ def get_kv_cache_group_metadata(self) -> list[dict[str, int | str | None]]: ) return metadata + def _kv_event_block_size(self) -> int: + """Return event granularity without changing native cache block size.""" + return select_kv_event_block_size( + self.get_kv_cache_group_metadata(), + self.vllm_config.cache_config.block_size or 0, + ) + def add_request(self, request: Request, request_wave: int = 0): """Add request to the scheduler. @@ -994,7 +1024,9 @@ def _eep_scale_up_before_kv_init(self): raise NotImplementedError def _eep_send_engine_core_notification( - self, notification_type: EEPNotificationType + self, + notification_type: EEPNotificationType, + vllm_config: VllmConfig | None = None, ): raise NotImplementedError @@ -1023,6 +1055,7 @@ def __init__( tensor_queue: Queue | None = None, *, engine_index: int = 0, + logical_data_parallel_size: int | None = None, ): self.input_queue = queue.Queue[tuple[EngineCoreRequestType, Any]]() self.output_queue = queue.Queue[tuple[int, EngineCoreOutputs] | bytes]() @@ -1031,6 +1064,11 @@ def __init__( ) self.engine_index = engine_index + self.logical_data_parallel_size = ( + logical_data_parallel_size + if logical_data_parallel_size is not None + else vllm_config.parallel_config.data_parallel_size + ) identity = self.engine_index.to_bytes(length=2, byteorder="little") self.engines_running = False self.shutdown_state = EngineShutdownState.RUNNING @@ -1069,6 +1107,11 @@ def __init__( self.addresses = addresses self.process_input_queue_block = True + if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: + self._eep_send_engine_core_notification( + EEPNotificationType.NEW_CORE_ENGINES_INIT_READY, + vllm_config=vllm_config, + ) self._init_data_parallel(vllm_config) super().__init__( @@ -1280,6 +1323,7 @@ def run_engine_core(*args, dp_rank: int = 0, local_dp_rank: int = 0, **kwargs): try: vllm_config: VllmConfig = kwargs["vllm_config"] parallel_config: ParallelConfig = vllm_config.parallel_config + logical_data_parallel_size = parallel_config.data_parallel_size data_parallel = parallel_config.data_parallel_size > 1 or dp_rank > 0 if data_parallel: parallel_config.data_parallel_rank_local = local_dp_rank @@ -1315,7 +1359,12 @@ def run_engine_core(*args, dp_rank: int = 0, local_dp_rank: int = 0, **kwargs): parallel_config.data_parallel_size = 1 parallel_config.data_parallel_size_local = 1 parallel_config.data_parallel_rank = 0 - engine_core = EngineCoreProc(*args, engine_index=dp_rank, **kwargs) + engine_core = EngineCoreProc( + *args, + engine_index=dp_rank, + logical_data_parallel_size=logical_data_parallel_size, + **kwargs, + ) assert engine_core is not None @@ -1619,6 +1668,9 @@ def _send_engine_dead(self): def _make_ready_response(self) -> EngineCoreReadyResponse: parallel_config = self.vllm_config.parallel_config scheduler_config = self.vllm_config.scheduler_config + kv_transfer_config = self.vllm_config.kv_transfer_config + kv_events_config = self.vllm_config.kv_events_config + lora_config = self.vllm_config.lora_config return EngineCoreReadyResponse( max_model_len=self.vllm_config.model_config.max_model_len, num_gpu_blocks=self.vllm_config.cache_config.num_gpu_blocks or 0, @@ -1627,7 +1679,7 @@ def _make_ready_response(self) -> EngineCoreReadyResponse: dtype=str(self.vllm_config.model_config.dtype).removeprefix("torch."), vllm_version=VLLM_VERSION, world_size=self.vllm_config.parallel_config.world_size, - data_parallel_size=parallel_config.data_parallel_size, + data_parallel_size=self.logical_data_parallel_size, kv_cache_size_tokens=self.vllm_config.cache_config.kv_cache_size_tokens, kv_cache_max_concurrency=( self.vllm_config.cache_config.kv_cache_max_concurrency @@ -1639,7 +1691,15 @@ def _make_ready_response(self) -> EngineCoreReadyResponse: max_num_seqs=scheduler_config.max_num_seqs, max_num_batched_tokens=scheduler_config.max_num_batched_tokens, instance_id=self.vllm_config.instance_id, - kv_events_config=self.scheduler.get_kv_event_publisher_config(), + kv_role=(kv_transfer_config.kv_role if kv_transfer_config else None), + kv_events_publisher=( + kv_events_config.publisher if kv_events_config else None + ), + kv_events_endpoint=(kv_events_config.endpoint if kv_events_config else None), + kv_events_topic=(kv_events_config.topic if kv_events_config else None), + kv_event_block_size=self._kv_event_block_size(), + supports_lora=lora_config is not None, + max_loras=(lora_config.max_loras if lora_config is not None else 0), ) def process_input_sockets( @@ -1749,11 +1809,10 @@ def process_output_sockets( encoder = MsgpackEncoder() # Send buffers to reuse. reuse_buffers: list[bytearray] = [] - # Payload buffers that can't be reused yet because zmq may still be - # sending them. - # Buffers of the zero-copy tensor/ndarray frames don't need tracking - # here: zmq itself holds a reference to each until it's done with it. - pending = deque[tuple[zmq.MessageTracker, bytearray]]() + # Keep references to outputs and buffers until zmq is finished + # with them (outputs may contain tensors/np arrays whose + # backing buffers were extracted for zero-copy send). + pending = deque[tuple[zmq.MessageTracker, Any, bytearray]]() # We must set linger to ensure the ENGINE_CORE_DEAD # message is sent prior to closing the socket. @@ -1794,38 +1853,20 @@ def process_output_sockets( # Reclaim buffers that zmq is finished with. while pending and pending[-1][0].done: - reclaimed = pending.pop()[1] - if len(reuse_buffers) < max_reuse_bufs: - reuse_buffers.append(reclaimed) + reuse_buffers.append(pending.pop()[2]) buffer = reuse_buffers.pop() if reuse_buffers else bytearray() buffers = encoder.encode_into(outputs, buffer) - tracker = self._send_msg_tracking_payload( - sockets[client_index], buffers + tracker = sockets[client_index].send_multipart( + buffers, copy=False, track=True ) if not tracker.done: - pending.appendleft((tracker, buffer)) + ref = outputs if len(buffers) > 1 else None + pending.appendleft((tracker, ref, buffer)) elif len(reuse_buffers) < max_reuse_bufs: # Limit the number of buffers to reuse. reuse_buffers.append(buffer) - @staticmethod - def _send_msg_tracking_payload( - socket: zmq.Socket, buffers: Sequence[bytestr] - ) -> zmq.MessageTracker: - """Send `buffers` as a zero-copy multipart message, returning a tracker - for the *first* frame. - - Used instead of `Socket.send_multipart()` because we reuse the buffer - passed to `MsgpackEncoder.encode_into()`: `send_multipart()` returns a - tracker for the last frame only. - """ - more_flag = zmq.SNDMORE if len(buffers) > 1 else 0 - tracker = socket.send(buffers[0], more_flag, copy=False, track=True) - if more_flag: - socket.send_multipart(buffers[1:], copy=False) - return tracker - def _handle_request_preproc_error(self, request: EngineCoreRequest) -> None: """Log and return a request-scoped error response for exceptions raised from the add request preprocessing in the input socket processing thread. @@ -2110,16 +2151,12 @@ def run_busy_loop(self): self._maybe_publish_request_counts() if self.eep_scaling_state is not None: - state = self.eep_scaling_state - if state.commit_requested or not state.is_ready_for_switch(): - state.progress() - if state.is_complete(): - if state.worker_type == "removing": + _ = self.eep_scaling_state.progress() + if self.eep_scaling_state.is_complete(): + if self.eep_scaling_state.worker_type == "removing": raise SystemExit self.process_input_queue_block = True self.eep_scaling_state = None - elif not state.commit_requested and state.is_ready_for_switch(): - self.process_input_queue_block = True executed = self._process_engine_step() self._maybe_publish_request_counts() @@ -2189,7 +2226,7 @@ def _has_global_unfinished_reqs(self, local_unfinished: bool) -> bool: def reinitialize_distributed( self, reconfig_request: ReconfigureDistributedRequest - ) -> str: + ) -> None: from copy import deepcopy from vllm.distributed.elastic_ep.elastic_state import ElasticEPScalingState @@ -2221,10 +2258,7 @@ def reinitialize_distributed( == ReconfigureRankType.SHUTDOWN_CURRENT_RANK ) - if self.eep_scaling_state is not None: - raise RuntimeError("Elastic EP reconfiguration is already active") - - state = ElasticEPScalingState( + self.eep_scaling_state = ElasticEPScalingState( model_executor=self.model_executor, engine_core=self, vllm_config=self.vllm_config, @@ -2233,34 +2267,30 @@ def reinitialize_distributed( scale_type="scale_down" if is_scale_down else "scale_up", reconfig_request=reconfig_request, ) - self.eep_scaling_state = state - self.process_input_queue_block = False logger.info( "[Elastic EP] Received reconfiguration request and starting scaling up/down" ) - return state.ready_key - - def commit_prepared_elastic_ep(self) -> None: - state = self.eep_scaling_state - if state is None or state.commit_requested or not state.is_ready_for_switch(): - raise RuntimeError("No prepared Elastic EP reconfiguration is ready") - state.commit_requested = True - self.process_input_queue_block = False - logger.info("[Elastic EP] Committing prepared reconfiguration") def _eep_send_engine_core_notification( - self, notification_type: EEPNotificationType + self, + notification_type: EEPNotificationType, + vllm_config: VllmConfig | None = None, ): """ Send notifications to EngineCoreClient, which can then forward the notifications to other engine core processes. It is used for: - 1) In scale down: removing core engines to notify EngineCoreClient + 1) In scale up: new core engines to notify existing core engines + that they are ready; + 2) In scale down: removing core engines to notify EngineCoreClient so EngineCoreClient can release their ray placement groups; - 2) Both scale up/down: to notify EngineCoreClient that existing + 3) Both scale up/down: to notify EngineCoreClient that existing core engines have already switched to the new parallel setup. """ - dp_rank = self.vllm_config.parallel_config.data_parallel_rank + if vllm_config is None: + dp_rank = self.vllm_config.parallel_config.data_parallel_rank + else: + dp_rank = vllm_config.parallel_config.data_parallel_rank notification_data = (notification_type.value, dp_rank) outputs = EngineCoreOutputs( utility_output=UtilityOutput( @@ -2282,11 +2312,22 @@ def _eep_send_engine_core_notification( ): socket.send_multipart(encoder.encode(outputs)) + def eep_handle_engine_core_notification( + self, notification_type: str | EEPNotificationType + ): + """ + Handle notification received from EngineCoreClient + (forwarded from new core engines). + """ + assert self.eep_scaling_state is not None + if isinstance(notification_type, str): + notification_type = EEPNotificationType(notification_type) + self.eep_scaling_state.handle_notification(notification_type) + def _eep_scale_up_before_kv_init(self): from vllm.distributed.elastic_ep.elastic_state import ElasticEPScalingState - self.ignore_start_dp_wave = True - state = ElasticEPScalingState( + self.eep_scaling_state = ElasticEPScalingState( model_executor=self.model_executor, engine_core=self, vllm_config=self.vllm_config, @@ -2295,10 +2336,7 @@ def _eep_scale_up_before_kv_init(self): scale_type="scale_up", reconfig_request=None, ) - if self.eep_scaling_state is not None: - raise RuntimeError("Elastic EP reconfiguration is already active") - self.eep_scaling_state = state - state.run_pre_kv_init_states() + self.eep_scaling_state.run_pre_kv_init_states() self.process_input_queue_block = False From f9ef944941a586956406b57ee14cb15e0065ebf5 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Sun, 26 Jul 2026 19:23:33 -0700 Subject: [PATCH 2/6] test(grpc): reconcile native sidecar test harness with merged output types --- rust/src/server/src/grpc/tests.rs | 11 ++++++----- rust/src/server/src/routes/inference/generate.rs | 1 + 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index 34f966b7fcef..da2b6670acba 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -1,6 +1,3 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project - use std::future::Future; use std::io; use std::pin::Pin; @@ -120,10 +117,10 @@ fn request_output( stop_reason: None, events: None, kv_transfer_params: None, - ec_transfer_params: None, trace_headers: None, prefill_stats: None, routed_experts: None, + ec_transfer_params: None, num_nans_in_logits: 0, } } @@ -162,6 +159,10 @@ async fn recv_engine_message(dealer: &mut DealerSocket) -> Vec { dealer.recv().await.expect("recv engine message").into_vec() } +fn test_llm(client: EngineCoreClient) -> Llm { + Llm::new(client).with_request_id_randomization(false) +} + #[derive(Clone, Debug)] struct FakeTextBackend; @@ -249,7 +250,7 @@ async fn setup_grpc_service( let engine_health = client.subscribe_health(); let chat = ChatLlm::from_shared_backend( - Llm::new(client), + test_llm(client), Arc::new(FakeTextBackend) as Arc, ); let state = Arc::new(AppState::new(vec!["test-model".to_string()], chat)); diff --git a/rust/src/server/src/routes/inference/generate.rs b/rust/src/server/src/routes/inference/generate.rs index b13194df3af2..0428a5018575 100644 --- a/rust/src/server/src/routes/inference/generate.rs +++ b/rust/src/server/src/routes/inference/generate.rs @@ -504,6 +504,7 @@ mod tests { }, kv_transfer_params: None, ec_transfer_params: None, + routed_experts: None, prompt_token_ids, }; From a4170a35cbe80a3368a60f78ffdc414b0541a7f7 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Mon, 27 Jul 2026 01:27:13 -0700 Subject: [PATCH 3/6] fix(grpc): restore upstream health monitoring and drop fork duplicate --- rust/src/server/src/lib.rs | 27 +-------------------------- rust/src/server/src/routes/lora.rs | 1 - 2 files changed, 1 insertion(+), 27 deletions(-) diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs index dc43ff59ce1c..8e1032700935 100644 --- a/rust/src/server/src/lib.rs +++ b/rust/src/server/src/lib.rs @@ -34,7 +34,6 @@ use hyper_util::rt::{TokioIo, TokioTimer}; use hyper_util::server::graceful::GracefulShutdown; use hyper_util::service::TowerToHyperService; use tokio::net::TcpListener; -use tokio::sync::watch; use tokio::time::{Instant, sleep_until}; use tokio_util::sync::CancellationToken; use tonic::transport::Server as TonicServer; @@ -80,30 +79,6 @@ async fn set_grpc_not_serving(health_reporter: &HealthReporter) { .await; } -async fn wait_until_engine_unhealthy(mut engine_health: watch::Receiver) { - loop { - if !*engine_health.borrow_and_update() { - return; - } - if engine_health.changed().await.is_err() { - return; - } - } -} - -async fn monitor_grpc_health( - health_reporter: HealthReporter, - engine_health: watch::Receiver, - shutdown: CancellationToken, -) { - tokio::select! { - _ = wait_until_engine_unhealthy(engine_health) => {} - _ = shutdown.cancelled() => {} - } - - set_grpc_not_serving(&health_reporter).await; -} - async fn monitor_lora_health( state: Arc, health_reporter: HealthReporter, @@ -418,7 +393,7 @@ where }; let lora_health = monitor_lora_health(state, health_reporter.clone(), shutdown.child_token()); - let engine_health = monitor_grpc_health(health_reporter, engine_health, shutdown); + let engine_health = grpc::monitor_health(health_reporter, engine_health, shutdown); tokio::join!(engine_health, lora_health); } }; diff --git a/rust/src/server/src/routes/lora.rs b/rust/src/server/src/routes/lora.rs index c879e1a0ccd3..488c4ed865eb 100644 --- a/rust/src/server/src/routes/lora.rs +++ b/rust/src/server/src/routes/lora.rs @@ -1,7 +1,6 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project -use std::path::{Component, Path, PathBuf}; use std::sync::Arc; use axum::extract::State; From 017cd3fff751ec653f34d5550ddd199c724fe841 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Tue, 28 Jul 2026 00:55:07 -0700 Subject: [PATCH 4/6] refactor(grpc): adopt upstream control/inference split and share drain state via AppState --- rust/src/server/src/grpc/control.rs | 26 +++---- rust/src/server/src/grpc/inference.rs | 15 ++-- rust/src/server/src/grpc/tests.rs | 69 +++++++++++++++---- rust/src/server/src/grpc/tests/lora.rs | 24 ++++--- rust/src/server/src/lib.rs | 17 ++--- .../server/src/routes/inference/generate.rs | 1 - rust/src/server/src/state.rs | 8 +++ 7 files changed, 103 insertions(+), 57 deletions(-) diff --git a/rust/src/server/src/grpc/control.rs b/rust/src/server/src/grpc/control.rs index c6934e55669f..d87422073f8b 100644 --- a/rust/src/server/src/grpc/control.rs +++ b/rust/src/server/src/grpc/control.rs @@ -9,7 +9,7 @@ use tonic::{Request, Response, Status}; use tonic_health::server::HealthReporter; use vllm_engine_core_client::protocol::handshake::EngineCoreReadyResponse; -use super::{AdmissionGuard, AdmissionState, ControlServer, lora_rpc, pb}; +use super::{AdmissionGuard, ControlServer, lora_rpc, pb}; use crate::lora_path::runtime_lora_allowed_path_prefixes; use crate::state::AppState; @@ -25,7 +25,6 @@ const GRPC_CAPABILITIES: &[&str] = &[ /// gRPC control service backed by the shared application state. pub struct ControlServiceImpl { state: Arc, - admission: Arc, health_reporter: Option, lora_allowed_path_prefixes: Option>, runtime_lora_updating_enabled: bool, @@ -33,23 +32,20 @@ pub struct ControlServiceImpl { impl ControlServiceImpl { pub fn new(state: Arc) -> Self { - Self::with_admission(state, Arc::new(AdmissionState::default()), None) - } - - pub(crate) fn with_admission( - state: Arc, - admission: Arc, - health_reporter: Option, - ) -> Self { Self { state, - admission, - health_reporter, + health_reporter: None, lora_allowed_path_prefixes: runtime_lora_allowed_path_prefixes().map(Arc::from), runtime_lora_updating_enabled: crate::routes::runtime_lora_updating_enabled(), } } + /// Report `NOT_SERVING` on `Drain` so load balancers stop routing here. + pub(crate) fn with_health_reporter(mut self, health_reporter: HealthReporter) -> Self { + self.health_reporter = Some(health_reporter); + self + } + #[cfg(test)] pub(crate) fn with_lora_allowed_path_prefixes(mut self, prefixes: Vec) -> Self { self.lora_allowed_path_prefixes = Some(prefixes.into()); @@ -88,7 +84,7 @@ impl ControlServiceImpl { } fn try_admit(&self) -> Option { - self.admission.try_admit() + self.state.admission().try_admit() } } @@ -168,9 +164,9 @@ impl pb::control_server::Control for ControlServiceImpl { &self, _request: Request, ) -> Result, Status> { - self.admission.begin_drain(); + self.state.admission().begin_drain(); self.report_not_serving().await; - let in_flight = self.admission.in_flight().min(u64::from(u32::MAX)) as u32; + let in_flight = self.state.admission().in_flight().min(u64::from(u32::MAX)) as u32; let state = if in_flight == 0 { pb::DrainState::Complete } else { diff --git a/rust/src/server/src/grpc/inference.rs b/rust/src/server/src/grpc/inference.rs index 3dca3b7fae7b..0afadcf06317 100644 --- a/rust/src/server/src/grpc/inference.rs +++ b/rust/src/server/src/grpc/inference.rs @@ -13,7 +13,7 @@ use tracing::info; use vllm_text::{DecodedTextEvent, Prompt, TextOutputStreamExt as _, TextRequest}; use super::convert::{self, ResponseOpts}; -use super::{AdmissionState, InferenceServer, pb}; +use super::{InferenceServer, pb}; use crate::state::AppState; pub(crate) type InferenceGrpcService = InferenceServer; @@ -21,16 +21,11 @@ pub(crate) type InferenceGrpcService = InferenceServer; /// gRPC inference service backed by the shared application state. pub struct InferenceServiceImpl { state: Arc, - admission: Arc, } impl InferenceServiceImpl { pub fn new(state: Arc) -> Self { - Self::with_admission(state, Arc::new(AdmissionState::default())) - } - - pub(crate) fn with_admission(state: Arc, admission: Arc) -> Self { - Self { state, admission } + Self { state } } async fn prepare_request( @@ -102,7 +97,8 @@ impl pb::inference_server::Inference for InferenceServiceImpl { request: Request, ) -> Result, Status> { let _guard = self - .admission + .state + .admission() .try_admit() .ok_or_else(|| Status::unavailable("gRPC service is draining"))?; let proto_req = request.into_inner(); @@ -153,7 +149,8 @@ impl pb::inference_server::Inference for InferenceServiceImpl { request: Request, ) -> Result, Status> { let guard = self - .admission + .state + .admission() .try_admit() .ok_or_else(|| Status::unavailable("gRPC service is draining"))?; let proto_req = request.into_inner(); diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index da2b6670acba..0bccc9887f3d 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -1,3 +1,6 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + use std::future::Future; use std::io; use std::pin::Pin; @@ -117,10 +120,10 @@ fn request_output( stop_reason: None, events: None, kv_transfer_params: None, + ec_transfer_params: None, trace_headers: None, prefill_stats: None, routed_experts: None, - ec_transfer_params: None, num_nans_in_logits: 0, } } @@ -159,10 +162,6 @@ async fn recv_engine_message(dealer: &mut DealerSocket) -> Vec { dealer.recv().await.expect("recv engine message").into_vec() } -fn test_llm(client: EngineCoreClient) -> Llm { - Llm::new(client).with_request_id_randomization(false) -} - #[derive(Clone, Debug)] struct FakeTextBackend; @@ -250,24 +249,64 @@ async fn setup_grpc_service( let engine_health = client.subscribe_health(); let chat = ChatLlm::from_shared_backend( - test_llm(client), + Llm::new(client), Arc::new(FakeTextBackend) as Arc, ); let state = Arc::new(AppState::new(vec!["test-model".to_string()], chat)); - // Inference and control must share one AdmissionState so `Drain` on the - // control service stops the inference service admitting new work. - let admission = std::sync::Arc::new(super::AdmissionState::default()); ( - InferenceServer::new(InferenceServiceImpl::with_admission( - state.clone(), - admission.clone(), - )), - ControlServer::new(ControlServiceImpl::with_admission(state, admission, None)), + InferenceServer::new(InferenceServiceImpl::new(state.clone())), + ControlServer::new(ControlServiceImpl::new(state)), engine_health, engine_task, ) } +/// Build only the shared `AppState` from a mock engine with a caller-supplied +/// ready response, so a test can configure the services itself (e.g. LoRA +/// path prefixes on the control service). +async fn setup_state_with_ready_and_engine( + engine_id: impl Into, + ready_response: vllm_engine_core_client::protocol::handshake::EngineCoreReadyResponse, + run: F, +) -> ( + Arc, + tokio::sync::watch::Receiver, + 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 = engine_id.into(); + + let engine_task = MockEngineTask::new(spawn_mock_engine_task_with_ready( + handshake_address.clone(), + engine_id.clone(), + ready_response, + run, + )); + + 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 engine_health = client.subscribe_health(); + + let chat = ChatLlm::from_shared_backend( + Llm::new(client), + Arc::new(FakeTextBackend) as Arc, + ); + let state = Arc::new(AppState::new(vec!["test-model".to_string()], chat)); + (state, engine_health, engine_task) +} + /// Spin up a plaintext gRPC server backed by a mock engine. Returns the client, /// the gRPC server task, and the mock engine task. async fn grpc_test_server( @@ -1493,3 +1532,5 @@ async fn grpc_health_watch_closes_on_graceful_shutdown() { .expect("timed out waiting for gRPC server shutdown") .expect("gRPC server task failed"); } + +mod lora; diff --git a/rust/src/server/src/grpc/tests/lora.rs b/rust/src/server/src/grpc/tests/lora.rs index 8c1b0f1a2fa4..3cc884c9f29d 100644 --- a/rust/src/server/src/grpc/tests/lora.rs +++ b/rust/src/server/src/grpc/tests/lora.rs @@ -23,14 +23,22 @@ where let mut ready = default_ready_response(); ready.supports_lora = true; ready.max_loras = 1; - let (service, engine_health, engine_task) = - setup_grpc_service_with_ready_and_engine(engine_id, ready, run).await; - let service = service - .with_lora_allowed_path_prefixes(allowed_path_prefixes) - .with_runtime_lora_updating(runtime_updates_enabled); - let (_generate, control, _health, server_task, engine_task) = - start_grpc_test_server(service, engine_health, engine_task).await; - (control, server_task, engine_task) + let (state, engine_health, engine_task) = + setup_state_with_ready_and_engine(engine_id, ready, run).await; + let control_service = ControlServer::new( + ControlServiceImpl::new(state.clone()) + .with_lora_allowed_path_prefixes(allowed_path_prefixes) + .with_runtime_lora_updating(runtime_updates_enabled), + ); + let inference_service = InferenceServer::new(InferenceServiceImpl::new(state)); + let (channel, server_task) = start_grpc_test_server( + inference_service, + control_service, + engine_health, + tokio_util::sync::CancellationToken::new(), + ) + .await; + (ControlClient::new(channel), server_task, engine_task) } async fn send_utility_result(push: &mut PushSocket, call_id: u64, result: bool) { diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs index 8e1032700935..1b69234c63d5 100644 --- a/rust/src/server/src/lib.rs +++ b/rust/src/server/src/lib.rs @@ -247,16 +247,13 @@ where let engine_health = state.engine_core_client().subscribe_health(); health_reporter.set_serving::().await; health_reporter.set_serving::().await; - let admission = std::sync::Arc::new(grpc::AdmissionState::default()); - let control_service = grpc::ControlGrpcService::new(grpc::ControlServiceImpl::with_admission( - state.clone(), - admission.clone(), - Some(health_reporter.clone()), - )); - let inference_service = grpc::InferenceGrpcService::new( - grpc::InferenceServiceImpl::with_admission(state.clone(), admission), - ) - .max_decoding_message_size(GRPC_MAX_REQUEST_SIZE); + let control_service = grpc::ControlGrpcService::new( + grpc::ControlServiceImpl::new(state.clone()) + .with_health_reporter(health_reporter.clone()), + ); + let inference_service = + grpc::InferenceGrpcService::new(grpc::InferenceServiceImpl::new(state.clone())) + .max_decoding_message_size(GRPC_MAX_REQUEST_SIZE); let svc = TonicServer::builder() .http2_keepalive_interval(Some(GRPC_KEEPALIVE_INTERVAL)) .http2_keepalive_timeout(Some(GRPC_KEEPALIVE_TIMEOUT)) diff --git a/rust/src/server/src/routes/inference/generate.rs b/rust/src/server/src/routes/inference/generate.rs index 0428a5018575..b13194df3af2 100644 --- a/rust/src/server/src/routes/inference/generate.rs +++ b/rust/src/server/src/routes/inference/generate.rs @@ -504,7 +504,6 @@ mod tests { }, kv_transfer_params: None, ec_transfer_params: None, - routed_experts: None, prompt_token_ids, }; diff --git a/rust/src/server/src/state.rs b/rust/src/server/src/state.rs index 8fac930ada98..0fff07366527 100644 --- a/rust/src/server/src/state.rs +++ b/rust/src/server/src/state.rs @@ -55,6 +55,8 @@ pub struct AppState { /// Profiler mode that registers `/start_profile` and `/stop_profile` /// routes when present. pub profiler: Option, + /// gRPC drain/admission state, shared by the inference and control services. + admission: Arc, } impl AppState { @@ -83,9 +85,15 @@ impl AppState { model_path: None, request_runtime: OnceLock::new(), profiler: None, + admission: Arc::new(crate::grpc::AdmissionState::default()), } } + /// Drain/admission state shared by the gRPC inference and control services. + pub(crate) fn admission(&self) -> &Arc { + &self.admission + } + /// Set HTTP/API-server behavior switches. pub fn with_api_server_options(mut self, options: ApiServerOptions) -> Self { self.api_server_options = options; From fb62f11909ec769a0dc0b74e26015cc521aada57 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Sun, 2 Aug 2026 10:19:17 -0700 Subject: [PATCH 5/6] style(rust): format restacked grpc modules --- rust/src/chat/src/lib.rs | 5 +---- rust/src/server/src/lib.rs | 4 +--- 2 files changed, 2 insertions(+), 7 deletions(-) diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index a93fe72a2224..c7795a433a12 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -298,10 +298,7 @@ impl ChatLlm { .backend .multimodal_model_info() .ok_or(Error::UnsupportedMultimodalRenderer)?; - let model_dtype = self - .processor - .model_dtype - .ok_or(Error::UnsupportedMultimodalRenderer)?; + let model_dtype = self.processor.model_dtype.ok_or(Error::UnsupportedMultimodalRenderer)?; let features = info.prepare_multimodal(media, token_ids, model_dtype).await?; Ok(Some(features)) } diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs index 1b69234c63d5..d87a19e7459f 100644 --- a/rust/src/server/src/lib.rs +++ b/rust/src/server/src/lib.rs @@ -66,9 +66,7 @@ const GRPC_KEEPALIVE_TIMEOUT: Duration = Duration::from_secs(20); const GRPC_LORA_HEALTH_POLL_INTERVAL: Duration = Duration::from_secs(1); async fn set_generate_not_serving(health_reporter: &HealthReporter) { - health_reporter - .set_not_serving::() - .await; + health_reporter.set_not_serving::().await; health_reporter.set_service_status("", ServingStatus::NotServing).await; } From e74fc3f1b06258e25519a5d5ed4d9a1b05d9a1cb Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Sun, 2 Aug 2026 10:35:10 -0700 Subject: [PATCH 6/6] fix(grpc): reconcile restacked protocol types --- rust/proto/inference.proto | 1 - .../engine-core-client/src/protocol/output.rs | 16 ++++++++---- rust/src/server/src/grpc/control.rs | 2 +- rust/src/server/src/grpc/convert.rs | 2 +- rust/src/server/src/grpc/tests.rs | 26 +++++-------------- rust/src/server/src/lib.rs | 1 + rust/src/text/src/lower.rs | 3 +++ 7 files changed, 23 insertions(+), 28 deletions(-) diff --git a/rust/proto/inference.proto b/rust/proto/inference.proto index 9493ee757d5a..bdefaf62ce18 100644 --- a/rust/proto/inference.proto +++ b/rust/proto/inference.proto @@ -257,4 +257,3 @@ message PreprocessedMultimodalFeature { // Canonical cache identity derived from modality and kwargs_msgpack. string cache_identifier = 5; } - diff --git a/rust/src/engine-core-client/src/protocol/output.rs b/rust/src/engine-core-client/src/protocol/output.rs index c48fc215ab32..481646b9d139 100644 --- a/rust/src/engine-core-client/src/protocol/output.rs +++ b/rust/src/engine-core-client/src/protocol/output.rs @@ -173,7 +173,8 @@ where ), })? .as_ref() - .to_vec(), + .to_vec() + .into(), }; let expected = shape .checked_numel() @@ -206,12 +207,13 @@ pub fn concatenate_routed_experts( let Some(mut combined) = chunks.next() else { return Ok(None); }; - let WireArrayData::RawView(combined_data) = &mut combined.data else { + let WireArrayData::RawView(initial_data) = &combined.data else { return Err(Error::Decode { target_type: "EngineCoreOutput.routed_experts", message: "unresolved auxiliary payload".to_string(), }); }; + let mut combined_data = initial_data.to_vec(); for chunk in chunks { if chunk.dtype != combined.dtype || chunk.shape.len() != combined.shape.len() @@ -233,8 +235,9 @@ pub fn concatenate_routed_experts( message: "unresolved auxiliary payload".to_string(), }); }; - combined_data.extend(data); + combined_data.extend_from_slice(&data); } + combined.data = WireArrayData::RawView(combined_data.into()); Ok(Some(combined)) } @@ -528,7 +531,7 @@ mod tests { decoded.as_request_batch().unwrap().outputs[0].routed_experts.as_ref().unwrap(); assert_eq!(routed.dtype, "|u1"); assert_eq!(routed.shape, vec![1, 2, 2]); - assert_eq!(routed.data.as_raw_view().unwrap(), &[1, 2, 3, 4]); + assert_eq!(&routed.data.as_raw_view().unwrap()[..], &[1, 2, 3, 4]); } #[test] @@ -541,7 +544,10 @@ mod tests { .unwrap(); assert_eq!(combined.shape, vec![3, 2, 1]); - assert_eq!(combined.data.as_raw_view().unwrap(), &[1, 2, 3, 4, 5, 6]); + assert_eq!( + &combined.data.as_raw_view().unwrap()[..], + &[1, 2, 3, 4, 5, 6] + ); } #[test] diff --git a/rust/src/server/src/grpc/control.rs b/rust/src/server/src/grpc/control.rs index d87422073f8b..554573b350c6 100644 --- a/rust/src/server/src/grpc/control.rs +++ b/rust/src/server/src/grpc/control.rs @@ -236,7 +236,7 @@ impl pb::control_server::Control for ControlServiceImpl { } } -fn kv_event_source(response: &EngineCoreReadyResponse) -> Option { +pub(super) fn kv_event_source(response: &EngineCoreReadyResponse) -> Option { if response.kv_events_publisher.as_deref() != Some("zmq") { return None; } diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index 64a47e744fda..4b3f3142c9ed 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -273,7 +273,7 @@ fn routed_experts_to_proto(tensor: &WireTensor) -> pb::RoutedExpertsTensor { pb::RoutedExpertsTensor { dtype: tensor.dtype.clone(), shape: tensor.shape.iter().map(|&dim| dim as u64).collect(), - data, + data: data.to_vec(), } } diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index 0bccc9887f3d..038636996e6c 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -29,7 +29,6 @@ use vllm_engine_core_client::mock_engine::{ DEFAULT_MOCK_BLOCK_SIZE, DEFAULT_MOCK_MAX_MODEL_LEN, DEFAULT_MOCK_NUM_GPU_BLOCKS, default_ready_response, }; -use vllm_engine_core_client::protocol::handshake::KvEventsConfig; use vllm_engine_core_client::protocol::output::{ EngineCoreFinishReason, EngineCoreOutput, EngineCoreOutputs, RequestBatchOutputs, }; @@ -711,7 +710,7 @@ async fn unary_generate_invalid_sampling_params_returns_invalid_argument() { model: "test-model".to_string(), prompt: Some(pb::generate_request::Prompt::Text("hi".to_string())), sampling: Some(pb::RandomSampling { - top_p: 2.0, + top_p: Some(2.0), ..Default::default() }), ..Default::default() @@ -1372,33 +1371,20 @@ async fn control_aggregates_multi_engine_capacity() { fn kv_event_source_filters_and_exposes_zmq_publisher() { let mut ready = default_ready_response(); ready.data_parallel_rank = 2; - ready.kv_events_config = Some(KvEventsConfig { - enable_kv_cache_events: false, - publisher: "null".to_string(), - endpoint: "tcp://*:5559".to_string(), - replay_endpoint: Some("tcp://*:5560".to_string()), - buffer_steps: 10_000, - hwm: 100_000, - max_queue_size: 100_000, - topic: "kv".to_string(), - }); + ready.kv_events_publisher = Some("null".to_string()); + ready.kv_events_endpoint = Some("tcp://*:5559".to_string()); + ready.kv_events_topic = Some("kv".to_string()); assert!(kv_event_source(&ready).is_none()); - let config = ready.kv_events_config.as_mut().unwrap(); - config.enable_kv_cache_events = true; - config.publisher = "zmq".to_string(); + ready.kv_events_publisher = Some("zmq".to_string()); let source = kv_event_source(&ready).expect("configured ZMQ event source"); assert_eq!(source.transport, "zmq"); - assert_eq!(source.endpoint, "tcp://*:5559"); assert_eq!(source.topic, "kv"); - assert_eq!(source.replay_endpoint, "tcp://*:5560"); assert_eq!(source.data_parallel_rank, Some(2)); assert_eq!(source.encoding, "msgpack"); assert_eq!(source.schema_version, 1); - assert_eq!(source.buffer_steps, 10_000); - assert_eq!(source.hwm, 100_000); - assert_eq!(source.max_queue_size, 100_000); + assert_eq!(source.endpoint_addr.unwrap().port, 5561); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs index d87a19e7459f..d19eca6c1a05 100644 --- a/rust/src/server/src/lib.rs +++ b/rust/src/server/src/lib.rs @@ -138,6 +138,7 @@ async fn build_state(config: &Config) -> Result> { .default_chat_template_kwargs .clone() .unwrap_or_default(), + limit_mm_per_prompt: config.limit_mm_per_prompt.clone(), }, ) .await diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs index b5c9b2f4e53d..2632049e4659 100644 --- a/rust/src/text/src/lower.rs +++ b/rust/src/text/src/lower.rs @@ -4,9 +4,11 @@ use std::collections::BTreeSet; pub(crate) mod logprobs; +pub(crate) mod sampling; pub(crate) mod token_ids; use logprobs::validate_logprobs; +use sampling::validate_resolved_sampling_params; use token_ids::{validate_prompt_token_ids, validate_vocab_range}; use vllm_engine_core_client::protocol::sampling::{ EngineCoreSamplingParams, RepetitionDetectionParams, @@ -188,6 +190,7 @@ pub fn lower_sampling_params( skip_reading_prefix_cache, extra_args: vllm_xargs, }; + validate_resolved_sampling_params(¶ms)?; validate_vocab_range(¶ms, &sampling_limits)?; Ok(params) }