From 1f9444a34ff4ebfba4d65c68971bb5306a11aa92 Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Tue, 25 Aug 2026 00:13:19 -0400 Subject: [PATCH] [Rust Frontend][gRPC] Add LoRA lifecycle control (#52840) Co-authored-by: Bugen Zhao Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Signed-off-by: Connor Carpenter Signed-off-by: Bugen Zhao --- rust/Cargo.lock | 1 + rust/proto/control.proto | 30 +++ rust/proto/inference.proto | 3 + .../engine-core-client/src/protocol/lora.rs | 27 ++- rust/src/server/Cargo.toml | 1 + rust/src/server/src/grpc/control.rs | 92 ++++++++- rust/src/server/src/grpc/inference.rs | 14 ++ rust/src/server/src/grpc/tests.rs | 174 ++++++++++++++++++ rust/src/server/src/lora.rs | 81 +++++--- rust/src/server/src/routes/lora.rs | 85 ++++----- rust/src/server/src/routes/tests.rs | 17 +- rust/src/server/src/state.rs | 26 ++- 12 files changed, 470 insertions(+), 81 deletions(-) diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 2adb17dc51..0ad12b46d5 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -6025,6 +6025,7 @@ dependencies = [ "socket2", "subtle", "tempfile", + "thiserror 2.0.18", "thiserror-ext", "tls-listener", "tokio", diff --git a/rust/proto/control.proto b/rust/proto/control.proto index 45eac602ec..3050035fdc 100644 --- a/rust/proto/control.proto +++ b/rust/proto/control.proto @@ -5,9 +5,19 @@ syntax = "proto3"; package vllm; service Control { + // Server and model discovery. rpc GetServerInfo (GetServerInfoRequest) returns (ServerInfo) {} rpc GetModelInfo (GetModelInfoRequest) returns (ModelInfo) {} + + // Request lifecycle. rpc Abort (AbortRequest) returns (AbortResponse) {} + + // LoRA lifecycle. + rpc LoadLora (LoadLoraRequest) returns (LoadLoraResponse) {} + rpc UnloadLora (UnloadLoraRequest) returns (UnloadLoraResponse) {} + rpc ListLoras (ListLorasRequest) returns (ListLorasResponse) {} + + // KV event discovery. rpc GetKvEventSources (GetKvEventSourcesRequest) returns (GetKvEventSourcesResponse) {} // Reinforcement-learning lifecycle and weight updates. @@ -79,6 +89,26 @@ message AbortRequest { message AbortResponse {} +// ====================================================================================== +// LoRA lifecycle +// ====================================================================================== + +message LoraAdapter { + int64 lora_id = 1; + string lora_name = 2; + string source_path = 3; +} + +message LoadLoraRequest { + string lora_name = 1; + string source_path = 2; +} +message LoadLoraResponse { LoraAdapter adapter = 1; } +message UnloadLoraRequest { string lora_name = 1; } +message UnloadLoraResponse { LoraAdapter adapter = 1; } +message ListLorasRequest {} +message ListLorasResponse { repeated LoraAdapter adapters = 1; } + // ====================================================================================== // Reinforcement-learning control // ====================================================================================== diff --git a/rust/proto/inference.proto b/rust/proto/inference.proto index 897319ca9c..c854a9abdf 100644 --- a/rust/proto/inference.proto +++ b/rust/proto/inference.proto @@ -50,6 +50,9 @@ message GenerateRequest { // Multimodal inputs aligned with placeholder markers in token_ids. repeated MediaItem media = 14; + + // Loaded LoRA name; empty selects the base model. + string lora_name = 15; } message RandomSampling { diff --git a/rust/src/engine-core-client/src/protocol/lora.rs b/rust/src/engine-core-client/src/protocol/lora.rs index 0f66d36cbe..4621a91da3 100644 --- a/rust/src/engine-core-client/src/protocol/lora.rs +++ b/rust/src/engine-core-client/src/protocol/lora.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project use serde_tuple::{Deserialize_tuple, Serialize_tuple}; +use thiserror::Error; use crate::protocol::OpaqueValue; @@ -24,6 +25,16 @@ pub struct LoraRequest { pub is_3d_lora_weight: bool, } +#[derive(Debug, Error, PartialEq, Eq)] +pub enum LoraRequestError { + #[error("lora_name must not be empty")] + EmptyName, + #[error("lora_int_id must be nonzero")] + ZeroId, + #[error("lora_path must not be empty")] + EmptyPath, +} + impl LoraRequest { pub fn new( lora_name: String, @@ -31,8 +42,18 @@ impl LoraRequest { lora_path: String, load_inplace: bool, is_3d_lora_weight: bool, - ) -> Self { - Self { + ) -> Result { + if lora_name.trim().is_empty() { + return Err(LoraRequestError::EmptyName); + } + if lora_int_id == 0 { + return Err(LoraRequestError::ZeroId); + } + if lora_path.trim().is_empty() { + return Err(LoraRequestError::EmptyPath); + } + + Ok(Self { lora_name, lora_int_id, lora_path, @@ -40,6 +61,6 @@ impl LoraRequest { tensorizer_config_dict: None, load_inplace, is_3d_lora_weight, - } + }) } } diff --git a/rust/src/server/Cargo.toml b/rust/src/server/Cargo.toml index cc0d278f27..6f5e95e93f 100644 --- a/rust/src/server/Cargo.toml +++ b/rust/src/server/Cargo.toml @@ -28,6 +28,7 @@ serde_with.workspace = true sha2.workspace = true socket2.workspace = true subtle.workspace = true +thiserror.workspace = true thiserror-ext.workspace = true tls-listener.workspace = true tokio.workspace = true diff --git a/rust/src/server/src/grpc/control.rs b/rust/src/server/src/grpc/control.rs index ba4ee7b997..99b74ded82 100644 --- a/rust/src/server/src/grpc/control.rs +++ b/rust/src/server/src/grpc/control.rs @@ -6,12 +6,14 @@ use std::sync::Arc; use serde_json::Value as JsonValue; use thiserror_ext::AsReport as _; use tokio::sync::Mutex; -use tonic::{Request, Response, Status}; +use tonic::{Code, Request, Response, Status}; use vllm_engine_core_client::EngineCoreClient; use vllm_engine_core_client::protocol::handshake::EngineCoreReadyResponse; +use vllm_engine_core_client::protocol::lora::LoraRequest; use vllm_engine_core_client::protocol::utility::PauseMode as EnginePauseMode; use super::{ControlServer, pb}; +use crate::lora::{LoadLoraError, LoraDisabledError, UnloadLoraError}; use crate::state::AppState; pub(crate) type ControlGrpcService = ControlServer; @@ -134,6 +136,41 @@ fn weight_version(value: String) -> Result { Ok(value) } +fn lora_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 load_lora_status(error: LoadLoraError) -> Status { + let code = match &error { + LoadLoraError::Disabled(_) => Code::FailedPrecondition, + LoadLoraError::InvalidRequest(_) => Code::InvalidArgument, + LoadLoraError::AlreadyLoaded { .. } | LoadLoraError::BaseModelName { .. } => { + Code::AlreadyExists + } + LoadLoraError::Engine { .. } | LoadLoraError::NotLoaded { .. } => Code::Internal, + }; + Status::new(code, error.to_report_string()) +} + +fn unload_lora_status(error: UnloadLoraError) -> Status { + let code = match &error { + UnloadLoraError::Disabled(_) => Code::FailedPrecondition, + UnloadLoraError::NotFound { .. } => Code::NotFound, + UnloadLoraError::IntIdMismatch { .. } + | UnloadLoraError::Engine { .. } + | UnloadLoraError::NotRemoved { .. } => Code::Internal, + }; + Status::new(code, error.to_report_string()) +} + +fn list_loras_status(error: LoraDisabledError) -> Status { + Status::failed_precondition(error.to_report_string()) +} + #[tonic::async_trait] impl pb::control_server::Control for ControlServiceImpl { async fn get_server_info( @@ -201,6 +238,59 @@ impl pb::control_server::Control for ControlServiceImpl { Ok(Response::new(pb::AbortResponse {})) } + async fn load_lora( + &self, + request: Request, + ) -> Result, Status> { + let request = request.into_inner(); + let adapter = self + .state + .load_lora(request.lora_name, request.source_path, false, false) + .await + .map_err(load_lora_status)?; + Ok(Response::new(pb::LoadLoraResponse { + adapter: Some(lora_to_proto(&adapter)), + })) + } + + async fn unload_lora( + &self, + request: Request, + ) -> Result, Status> { + let name = request.into_inner().lora_name; + if name.trim().is_empty() { + return Err(Status::invalid_argument("lora_name is required")); + } + + let adapter = self + .state + .list_loras() + .await + .map_err(list_loras_status)? + .into_iter() + .find(|adapter| adapter.lora_name == name) + .ok_or_else(|| Status::not_found(format!("LoRA adapter `{name}` is not loaded")))?; + let adapter = self + .state + .unload_lora(&name, Some(adapter.lora_int_id)) + .await + .map_err(unload_lora_status)?; + Ok(Response::new(pb::UnloadLoraResponse { + adapter: Some(lora_to_proto(&adapter)), + })) + } + + async fn list_loras( + &self, + _request: Request, + ) -> Result, Status> { + let mut adapters = self.state.list_loras().await.map_err(list_loras_status)?; + adapters.sort_by(|left, right| left.lora_name.cmp(&right.lora_name)); + Ok(Response::new(pb::ListLorasResponse { + adapters: adapters.iter().map(lora_to_proto).collect(), + })) + } + async fn get_kv_event_sources( &self, _request: Request, diff --git a/rust/src/server/src/grpc/inference.rs b/rust/src/server/src/grpc/inference.rs index 8eebdb7210..d5403f03c7 100644 --- a/rust/src/server/src/grpc/inference.rs +++ b/rust/src/server/src/grpc/inference.rs @@ -58,12 +58,14 @@ impl InferenceServiceImpl { } else { proto_request.model.clone() }; + let lora_name = proto_request.lora_name.clone(); let media_count = proto_request.media.len(); let request_span = info_span!( "grpc_inference", %request_id, rpc, %model, + %lora_name, media_count, ?data_parallel_rank, ); @@ -76,6 +78,18 @@ impl InferenceServiceImpl { text_request.arrival_time = Some(arrival_time); text_request.data_parallel_rank = data_parallel_rank; + if !lora_name.is_empty() { + if !self.state.engine_core_client().ready_response().supports_lora { + return Err(Status::failed_precondition( + "engine was not started with LoRA enabled", + )); + } + let resolution = self.state.resolve_model_with_loras(Some(&lora_name)).await; + text_request.lora_request = Some(resolution.lora_request.ok_or_else(|| { + Status::not_found(format!("LoRA adapter `{lora_name}` is not loaded")) + })?); + } + let media = convert::media_parts_from_request(media)?; if !media.is_empty() { let Prompt::TokenIds(mut token_ids) = text_request.prompt else { diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index 6a1074b524..0db360f14b 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -151,6 +151,35 @@ async fn send_outputs(push: &mut PushSocket, outputs: EngineCoreOutputs) { .expect("send outputs"); } +async fn reply_utility_bool( + dealer: &mut DealerSocket, + push: &mut PushSocket, + expected_method: &str, + result: bool, +) { + let frames = recv_engine_message(dealer).await; + assert_eq!(frames[0].as_ref(), &[0x03]); + let payload = decode_value(&frames[1]).expect("decode utility payload"); + let fields = payload.as_array().expect("utility payload array"); + let call_id = fields[1].as_u64().expect("utility call id"); + assert_eq!(fields[2].as_str(), Some(expected_method)); + send_outputs( + push, + UtilityCallOutput { + output: UtilityOutput { + call_id: call_id.into(), + failure_message: None, + result: Some(UtilityResultEnvelope::without_type_info(Value::Boolean( + result, + ))), + }, + ..Default::default() + } + .into(), + ) + .await; +} + async fn recv_engine_message(dealer: &mut DealerSocket) -> Vec { dealer.recv().await.expect("recv engine message").into_vec() } @@ -1558,6 +1587,151 @@ async fn control_reports_server_and_model_info() { server_task.abort(); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn control_lora_lifecycle_selects_adapter_for_generation() { + let mut ready = default_ready_response(); + ready.supports_lora = true; + ready.max_loras = 4; + let (inference_service, control_service, engine_health, engine_task) = + setup_grpc_service_with_engine_script( + b"engine-grpc-lora".to_vec(), + ready, + Arc::new(FakeTextBackend), + |dealer, push| { + boxed_test_future(async move { + reply_utility_bool(dealer, push, "add_lora", true).await; + + let frames = recv_engine_message(dealer).await; + assert_eq!(frames[0].as_ref(), &[0x00]); + let request: EngineCoreRequest = + rmp_serde::from_slice(&frames[1]).expect("decode generation request"); + let lora = request.lora_request.expect("generation LoRA request"); + assert_eq!(lora.lora_name, "adapter"); + assert_eq!(lora.lora_int_id, 1); + send_outputs( + push, + engine_outputs_for_request( + &request.request_id, + vec![(vec![b'!' as u32], Some(EngineCoreFinishReason::Stop))], + ), + ) + .await; + + reply_utility_bool(dealer, push, "remove_lora", true).await; + }) + }, + ) + .await; + let (channel, server_task) = start_grpc_test_server( + inference_service, + control_service, + engine_health, + tokio_util::sync::CancellationToken::new(), + ) + .await; + let mut control_client = ControlClient::new(channel.clone()); + let mut inference_client = InferenceClient::new(channel); + let loaded = control_client + .load_lora(pb::LoadLoraRequest { + lora_name: "adapter".to_string(), + source_path: "adapter".to_string(), + }) + .await + .expect("load LoRA") + .into_inner(); + assert_eq!( + loaded.adapter.as_ref().map(|adapter| adapter.lora_id), + Some(1) + ); + + let duplicate = control_client + .load_lora(pb::LoadLoraRequest { + lora_name: "adapter".to_string(), + source_path: "adapter".to_string(), + }) + .await + .expect_err("duplicate LoRA name should be rejected"); + assert_eq!(duplicate.code(), tonic::Code::AlreadyExists); + + let listed = control_client + .list_loras(pb::ListLorasRequest {}) + .await + .expect("list LoRAs") + .into_inner(); + assert_eq!(listed.adapters.len(), 1); + assert_eq!(listed.adapters[0].lora_name, "adapter"); + + inference_client + .generate(pb::GenerateRequest { + request_id: "grpc-lora-request".to_string(), + model: "test-model".to_string(), + prompt: Some(pb::generate_request::Prompt::Text("hello".to_string())), + stopping: Some(pb::StoppingCriteria { + max_new_tokens: 1, + ..Default::default() + }), + lora_name: "adapter".to_string(), + ..Default::default() + }) + .await + .expect("generate with LoRA"); + + let unloaded = control_client + .unload_lora(pb::UnloadLoraRequest { + lora_name: "adapter".to_string(), + }) + .await + .expect("unload LoRA") + .into_inner(); + assert_eq!( + unloaded.adapter.as_ref().map(|adapter| adapter.lora_id), + Some(1) + ); + assert!( + control_client + .list_loras(pb::ListLorasRequest {}) + .await + .expect("list LoRAs after unload") + .into_inner() + .adapters + .is_empty() + ); + + engine_task.await.expect("mock engine task"); + server_task.abort(); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn control_list_loras_requires_lora_enabled_engine() { + let (inference_service, control_service, engine_health, engine_task) = + setup_grpc_service_with_engine_script( + b"engine-grpc-lora-disabled".to_vec(), + default_ready_response(), + Arc::new(FakeTextBackend), + |_, _| boxed_test_future(async move {}), + ) + .await; + let (channel, server_task) = start_grpc_test_server( + inference_service, + control_service, + engine_health, + tokio_util::sync::CancellationToken::new(), + ) + .await; + let mut control_client = ControlClient::new(channel); + let status = control_client + .list_loras(pb::ListLorasRequest {}) + .await + .expect_err("list should require LoRA support"); + assert_eq!(status.code(), tonic::Code::FailedPrecondition); + assert_eq!(status.message(), "engine was not started with LoRA enabled"); + + engine_task.await.expect("mock engine task"); + server_task.abort(); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn control_forwards_weight_update_without_pause_guard() { diff --git a/rust/src/server/src/lora.rs b/rust/src/server/src/lora.rs index 896ba32aac..5cca58c57d 100644 --- a/rust/src/server/src/lora.rs +++ b/rust/src/server/src/lora.rs @@ -4,9 +4,10 @@ use std::sync::atomic::{AtomicU64, Ordering}; use indexmap::IndexMap; +use thiserror::Error; use tokio::sync::{Mutex, RwLock}; use vllm_engine_core_client::EngineCoreClient; -use vllm_engine_core_client::protocol::lora::LoraRequest; +use vllm_engine_core_client::protocol::lora::{LoraRequest, LoraRequestError}; /// Snapshot of the currently served model names plus the requested LoRA, if /// the model name resolves to a dynamic adapter. @@ -26,29 +27,52 @@ pub(crate) struct LoraManager { update_lock: Mutex<()>, } -#[derive(Debug)] +#[derive(Debug, Error)] +#[error("engine was not started with LoRA enabled")] +pub(crate) struct LoraDisabledError; + +#[derive(Debug, Error)] pub(crate) enum LoadLoraError { + #[error(transparent)] + Disabled(#[from] LoraDisabledError), + #[error(transparent)] + InvalidRequest(#[from] LoraRequestError), + #[error("LoRA adapter `{lora_name}` is already loaded")] AlreadyLoaded { lora_name: String }, + #[error("LoRA adapter `{lora_name}` conflicts with a served base model")] BaseModelName { lora_name: String }, - Engine(vllm_engine_core_client::Error), + #[error("failed to load LoRA adapter `{lora_name}`")] + Engine { + lora_name: String, + #[source] + source: vllm_engine_core_client::Error, + }, + #[error("one or more engine ranks rejected LoRA adapter `{lora_name}`")] NotLoaded { lora_name: String }, } -#[derive(Debug)] +#[derive(Debug, Error)] pub(crate) enum UnloadLoraError { - NotFound { - lora_name: String, - }, + #[error(transparent)] + Disabled(#[from] LoraDisabledError), + #[error("LoRA adapter `{lora_name}` is not loaded")] + NotFound { lora_name: String }, + #[error( + "requested lora_int_id {actual} does not match loaded adapter `{lora_name}` with id {expected}" + )] IntIdMismatch { lora_name: String, expected: u64, actual: u64, }, - Engine(vllm_engine_core_client::Error), - NotRemoved { + #[error("failed to unload LoRA adapter `{lora_name}`")] + Engine { lora_name: String, - lora_int_id: u64, + #[source] + source: vllm_engine_core_client::Error, }, + #[error("engine rejected removal of LoRA adapter `{lora_name}` with id {lora_int_id}")] + NotRemoved { lora_name: String, lora_int_id: u64 }, } impl LoraManager { @@ -97,16 +121,13 @@ impl LoraManager { 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) { + let requests = self.requests.read().await; + let existing_lora_int_id = requests.get(&lora_name).map(|request| request.lora_int_id); + if !load_inplace && existing_lora_int_id.is_some() { return Err(LoadLoraError::AlreadyLoaded { lora_name }); } - let lora_int_id = self - .requests - .read() - .await - .get(&lora_name) - .map(|request| request.lora_int_id) + let lora_int_id = existing_lora_int_id .unwrap_or_else(|| self.id_counter.fetch_add(1, Ordering::Relaxed) + 1); let lora_request = LoraRequest::new( lora_name.clone(), @@ -114,12 +135,16 @@ impl LoraManager { lora_path, load_inplace, is_3d_lora_weight, - ); + ) + .map_err(LoadLoraError::InvalidRequest)?; + drop(requests); - let loaded = engine_core_client - .add_lora(&lora_request) - .await - .map_err(LoadLoraError::Engine)?; + let loaded = engine_core_client.add_lora(&lora_request).await.map_err(|source| { + LoadLoraError::Engine { + lora_name: lora_name.clone(), + source, + } + })?; if !loaded { return Err(LoadLoraError::NotLoaded { lora_name }); } @@ -152,10 +177,14 @@ impl LoraManager { }); } - let removed = engine_core_client - .remove_lora(lora_request.lora_int_id) - .await - .map_err(UnloadLoraError::Engine)?; + let removed = + engine_core_client + .remove_lora(lora_request.lora_int_id) + .await + .map_err(|source| UnloadLoraError::Engine { + lora_name: lora_request.lora_name.clone(), + source, + })?; if !removed { return Err(UnloadLoraError::NotRemoved { lora_name: lora_request.lora_name, diff --git a/rust/src/server/src/routes/lora.rs b/rust/src/server/src/routes/lora.rs index dad817c69a..c434e9fd6f 100644 --- a/rust/src/server/src/routes/lora.rs +++ b/rust/src/server/src/routes/lora.rs @@ -108,17 +108,45 @@ fn validate_lora_path_access( Ok(Some(canonical_path.to_string_lossy().into_owned())) } +fn load_lora_api_error(error: LoadLoraError) -> ApiError { + let message = error.to_report_string(); + match error { + LoadLoraError::Disabled(_) => ApiError::invalid_request(message, None), + LoadLoraError::InvalidRequest(_) => ApiError::invalid_request(message, None), + LoadLoraError::AlreadyLoaded { .. } => ApiError::invalid_request( + format!( + "{message}. If you want to load the adapter in place, set 'load_inplace' to true." + ), + Some("lora_name"), + ), + LoadLoraError::BaseModelName { .. } => { + ApiError::invalid_request(message, Some("lora_name")) + } + LoadLoraError::Engine { .. } | LoadLoraError::NotLoaded { .. } => { + ApiError::server_error(message) + } + } +} + +fn unload_lora_api_error(error: UnloadLoraError) -> ApiError { + let message = error.to_report_string(); + match error { + UnloadLoraError::Disabled(_) => ApiError::invalid_request(message, None), + UnloadLoraError::NotFound { lora_name } => ApiError::model_not_found(lora_name), + UnloadLoraError::IntIdMismatch { .. } => { + ApiError::invalid_request(message, Some("lora_int_id")) + } + UnloadLoraError::Engine { .. } | UnloadLoraError::NotRemoved { .. } => { + ApiError::server_error(message) + } + } +} + /// Dynamically load one LoRA adapter and expose it as an OpenAI model id. pub async fn load_lora_adapter( State(state): State>, ValidatedJson(request): ValidatedJson, ) -> Result { - if request.lora_name.is_empty() || request.lora_path.is_empty() { - return Err(ApiError::invalid_request( - "Both 'lora_name' and 'lora_path' must be provided.".to_string(), - None, - )); - } let allowed_prefixes = runtime_lora_allowed_path_prefixes(); let lora_path = validate_lora_path_access(&request.lora_path, allowed_prefixes.as_deref())? .unwrap_or(request.lora_path); @@ -132,25 +160,7 @@ pub async fn load_lora_adapter( request.is_3d_lora_weight, ) .await - .map_err(|error| match error { - 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." - ), - Some("lora_name"), - ), - LoadLoraError::BaseModelName { lora_name } => ApiError::invalid_request( - format!("The lora adapter name '{lora_name}' conflicts with a served base model."), - Some("lora_name"), - ), - LoadLoraError::Engine(error) => ApiError::server_error(format!( - "failed to load LoRA adapter '{lora_name}': {}", - error.to_report_string() - )), - LoadLoraError::NotLoaded { lora_name } => ApiError::server_error(format!( - "failed to load LoRA adapter '{lora_name}': engine rejected the adapter" - )), - })?; + .map_err(load_lora_api_error)?; Ok(format!( "Success: LoRA adapter '{lora_name}' added successfully." @@ -172,30 +182,7 @@ pub async fn unload_lora_adapter( let lora_request = state .unload_lora(&request.lora_name, request.lora_int_id) .await - .map_err(|error| match error { - UnloadLoraError::NotFound { lora_name } => ApiError::model_not_found(lora_name), - UnloadLoraError::IntIdMismatch { - lora_name, - expected, - actual, - } => ApiError::invalid_request( - format!( - "The requested lora_int_id {actual} does not match loaded adapter '{lora_name}' with id {expected}." - ), - Some("lora_int_id"), - ), - UnloadLoraError::Engine(error) => ApiError::server_error(format!( - "failed to 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}" - )), - })?; + .map_err(unload_lora_api_error)?; Ok(format!( "Success: LoRA adapter '{}' removed successfully.", diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index d57792b0bc..fb10e7cb47 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -28,6 +28,7 @@ use vllm_chat::{ }; use vllm_engine_core_client::mock_engine::default_ready_response; use vllm_engine_core_client::protocol::decode_value; +use vllm_engine_core_client::protocol::handshake::EngineCoreReadyResponse; use vllm_engine_core_client::protocol::logprobs::{ Logprobs, MaybeWireLogprobs, PositionLogprobs, TokenLogprob, }; @@ -833,6 +834,19 @@ where } async fn test_admin_app_with_engine_script(script: F) -> (axum::Router, 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 = 4; + test_admin_app_with_ready_and_engine_script(ready, script).await +} + +async fn test_admin_app_with_ready_and_engine_script( + ready: EngineCoreReadyResponse, + script: F, +) -> (axum::Router, MockEngineTask) where F: for<'a> FnOnce(&'a mut DealerSocket, &'a mut PushSocket) -> TestFuture<'a> + Send + 'static, { @@ -840,9 +854,10 @@ where let handshake_address = ipc.handshake_endpoint(); let engine_id = b"engine-openai-admin".to_vec(); - let engine_task = MockEngineTask::new(spawn_mock_engine_task( + let engine_task = MockEngineTask::new(spawn_mock_engine_task_with_ready( handshake_address.clone(), engine_id.clone(), + ready, move |dealer, push| script(dealer, push), )); diff --git a/rust/src/server/src/state.rs b/rust/src/server/src/state.rs index 722dede8dd..b13270a40c 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::{ + LoadLoraError, LoraDisabledError, LoraManager, LoraModelResolution, UnloadLoraError, +}; use crate::runtime::build_request_runtime; use crate::server_info::{ServerInfoConfigFormat, ServerInfoSnapshot}; @@ -161,12 +163,22 @@ impl AppState { self.lora_manager.served_lora_requests().await } + /// Snapshot loaded LoRA adapters for the lifecycle API. + /// + /// Returns error if the engine was started without LoRA support. + pub async fn list_loras(&self) -> Result, LoraDisabledError> { + self.ensure_lora_enabled()?; + Ok(self.lora_manager.served_lora_requests().await) + } + /// 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 } /// Load one dynamic LoRA adapter and register it as a public model name. + /// + /// Returns error if the engine was started without LoRA support. pub async fn load_lora( &self, lora_name: String, @@ -174,6 +186,7 @@ impl AppState { load_inplace: bool, is_3d_lora_weight: bool, ) -> Result { + self.ensure_lora_enabled()?; self.lora_manager .load_lora( self.engine_core_client(), @@ -188,16 +201,27 @@ impl AppState { /// Remove one dynamic LoRA adapter from the engine and public model /// registry. + /// + /// Returns error if the engine was started without LoRA support. pub async fn unload_lora( &self, lora_name: &str, lora_int_id: Option, ) -> Result { + self.ensure_lora_enabled()?; self.lora_manager .unload_lora(self.engine_core_client(), lora_name, lora_int_id) .await } + fn ensure_lora_enabled(&self) -> Result<(), LoraDisabledError> { + self.engine_core_client() + .ready_response() + .supports_lora + .then_some(()) + .ok_or(LoraDisabledError) + } + /// Return a reference to the underlying engine core client for utility /// calls. pub(crate) fn engine_core_client(&self) -> &EngineCoreClient {