| 1 | // SPDX-License-Identifier: Apache-2.0 |
| 2 | // SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | use std::path::{Component, Path, PathBuf}; |
| 5 | use std::sync::atomic::{AtomicU64, Ordering}; |
| 6 | |
| 7 | use indexmap::IndexMap; |
| 8 | use thiserror::Error; |
| 9 | use thiserror_ext::Macro; |
| 10 | use tokio::sync::{Mutex, RwLock}; |
| 11 | use vllm_engine_core_client::EngineCoreClient; |
| 12 | use vllm_engine_core_client::protocol::lora::LoraRequest; |
| 13 | |
| 14 | use crate::config::LoraModulePath; |
| 15 | |
| 16 | const RUNTIME_LORA_ALLOWED_PATH_PREFIXES_ENV: &str = "VLLM_RUNTIME_LORA_ALLOWED_PATH_PREFIXES"; |
| 17 | |
| 18 | #[derive(Debug, Error, Macro)] |
| 19 | pub(crate) enum LoraPathAccessError { |
| 20 | #[error("{message}")] |
| 21 | InvalidPath { message: String }, |
| 22 | #[error("{message}")] |
| 23 | InvalidConfiguration { message: String }, |
| 24 | } |
| 25 | |
| 26 | fn runtime_lora_allowed_path_prefixes() -> Option<Vec<PathBuf>> { |
| 27 | let prefixes = std::env::var_os(RUNTIME_LORA_ALLOWED_PATH_PREFIXES_ENV)?; |
| 28 | let prefixes: Vec<_> = std::env::split_paths(&prefixes) |
| 29 | .filter(|path| !path.as_os_str().is_empty()) |
| 30 | .collect(); |
| 31 | (!prefixes.is_empty()).then_some(prefixes) |
| 32 | } |
| 33 | |
| 34 | fn looks_like_local_lora_path(lora_path: &str) -> bool { |
| 35 | let path = Path::new(lora_path); |
| 36 | path.is_absolute() |
| 37 | || lora_path.starts_with('~') |
| 38 | || lora_path.starts_with('.') |
| 39 | || path.components().any(|component| matches!(component, Component::ParentDir)) |
| 40 | } |
| 41 | |
| 42 | fn validate_lora_path_access( |
| 43 | lora_path: &str, |
| 44 | allowed_prefixes: Option<&[PathBuf]>, |
| 45 | ) -> Result<Option<String>, LoraPathAccessError> { |
| 46 | let path = Path::new(lora_path); |
| 47 | if !looks_like_local_lora_path(lora_path) && !path.exists() { |
| 48 | return Ok(None); |
| 49 | } |
| 50 | |
| 51 | let Some(allowed_prefixes) = allowed_prefixes else { |
| 52 | return Err(invalid_path!( |
| 53 | "Local LoRA adapter paths require {RUNTIME_LORA_ALLOWED_PATH_PREFIXES_ENV} to be configured." |
| 54 | )); |
| 55 | }; |
| 56 | |
| 57 | if !path.is_absolute() { |
| 58 | return Err(invalid_path!( |
| 59 | "Local LoRA adapter paths must be absolute and under one of the prefixes configured by {RUNTIME_LORA_ALLOWED_PATH_PREFIXES_ENV}." |
| 60 | )); |
| 61 | } |
| 62 | |
| 63 | let canonical_path = path |
| 64 | .canonicalize() |
| 65 | .map_err(|_| invalid_path!("Local LoRA adapter path must exist and be accessible."))?; |
| 66 | let canonical_prefixes = allowed_prefixes |
| 67 | .iter() |
| 68 | .map(|prefix| { |
| 69 | prefix.canonicalize().map_err(|_| { |
| 70 | invalid_configuration!( |
| 71 | "configured {RUNTIME_LORA_ALLOWED_PATH_PREFIXES_ENV} path prefix must exist and be accessible" |
| 72 | ) |
| 73 | }) |
| 74 | }) |
| 75 | .collect::<Result<Vec<_>, _>>()?; |
| 76 | |
| 77 | if !canonical_prefixes.iter().any(|prefix| canonical_path.starts_with(prefix)) { |
| 78 | return Err(invalid_path!( |
| 79 | "Local LoRA adapter path is outside the configured allowed prefixes." |
| 80 | )); |
| 81 | } |
| 82 | |
| 83 | Ok(Some(canonical_path.to_string_lossy().into_owned())) |
| 84 | } |
| 85 | |
| 86 | /// Snapshot of the currently served model names plus the requested LoRA, if |
| 87 | /// the model name resolves to a dynamic adapter. |
| 88 | #[derive(Debug, Clone)] |
| 89 | pub(crate) struct LoraModelResolution { |
| 90 | pub model_names: Vec<String>, |
| 91 | pub lora_request: Option<LoraRequest>, |
| 92 | } |
| 93 | |
| 94 | /// Runtime registry for dynamically loaded LoRA adapters. |
| 95 | pub(crate) struct LoraManager { |
| 96 | /// Dynamically loaded LoRA adapters keyed by public model name, in load order. |
| 97 | requests: RwLock<IndexMap<String, LoraRequest>>, |
| 98 | /// Monotonic adapter id allocator. LoRA ids are one-indexed. |
| 99 | id_counter: AtomicU64, |
| 100 | /// Serialize dynamic LoRA registry updates around engine utility calls. |
| 101 | update_lock: Mutex<()>, |
| 102 | } |
| 103 | |
| 104 | #[derive(Debug, Error)] |
| 105 | #[error("engine was not started with LoRA enabled")] |
| 106 | pub(crate) struct LoraDisabledError; |
| 107 | |
| 108 | #[derive(Debug, Error, Macro)] |
| 109 | pub(crate) enum LoadLoraError { |
| 110 | #[error(transparent)] |
| 111 | Disabled(#[from] LoraDisabledError), |
| 112 | #[error("{message}")] |
| 113 | InvalidAdapter { message: String }, |
| 114 | #[error(transparent)] |
| 115 | PathAccess(#[from] LoraPathAccessError), |
| 116 | #[error("LoRA adapter `{lora_name}` is already loaded")] |
| 117 | AlreadyLoaded { lora_name: String }, |
| 118 | #[error("LoRA adapter `{lora_name}` conflicts with a served base model")] |
| 119 | BaseModelName { lora_name: String }, |
| 120 | #[error("failed to load LoRA adapter `{lora_name}`")] |
| 121 | Engine { |
| 122 | lora_name: String, |
| 123 | #[source] |
| 124 | source: vllm_engine_core_client::Error, |
| 125 | }, |
| 126 | #[error("one or more engine ranks rejected LoRA adapter `{lora_name}`")] |
| 127 | NotLoaded { lora_name: String }, |
| 128 | } |
| 129 | |
| 130 | #[derive(Debug, Error)] |
| 131 | pub(crate) enum UnloadLoraError { |
| 132 | #[error(transparent)] |
| 133 | Disabled(#[from] LoraDisabledError), |
| 134 | #[error("LoRA adapter `{lora_name}` is not loaded")] |
| 135 | NotFound { lora_name: String }, |
| 136 | #[error( |
| 137 | "requested lora_int_id {actual} does not match loaded adapter `{lora_name}` with id {expected}" |
| 138 | )] |
| 139 | IntIdMismatch { |
| 140 | lora_name: String, |
| 141 | expected: u64, |
| 142 | actual: u64, |
| 143 | }, |
| 144 | #[error("failed to unload LoRA adapter `{lora_name}`")] |
| 145 | Engine { |
| 146 | lora_name: String, |
| 147 | #[source] |
| 148 | source: vllm_engine_core_client::Error, |
| 149 | }, |
| 150 | #[error("engine rejected removal of LoRA adapter `{lora_name}` with id {lora_int_id}")] |
| 151 | NotRemoved { lora_name: String, lora_int_id: u64 }, |
| 152 | } |
| 153 | |
| 154 | impl LoraManager { |
| 155 | pub fn new() -> Self { |
| 156 | Self { |
| 157 | requests: RwLock::new(IndexMap::new()), |
| 158 | id_counter: AtomicU64::new(0), |
| 159 | update_lock: Mutex::new(()), |
| 160 | } |
| 161 | } |
| 162 | |
| 163 | /// Snapshot loaded LoRA adapters in load order. |
| 164 | pub async fn served_lora_requests(&self) -> Vec<LoraRequest> { |
| 165 | self.requests.read().await.values().cloned().collect() |
| 166 | } |
| 167 | |
| 168 | /// Resolve the requested model against one consistent LoRA registry |
| 169 | /// snapshot. |
| 170 | pub async fn resolve_model( |
| 171 | &self, |
| 172 | base_model_names: &[String], |
| 173 | model_name: Option<&str>, |
| 174 | ) -> LoraModelResolution { |
| 175 | let requests = self.requests.read().await; |
| 176 | let mut model_names = base_model_names.to_vec(); |
| 177 | model_names.extend(requests.keys().cloned()); |
| 178 | let lora_request = model_name.and_then(|name| requests.get(name).cloned()); |
| 179 | |
| 180 | LoraModelResolution { |
| 181 | model_names, |
| 182 | lora_request, |
| 183 | } |
| 184 | } |
| 185 | |
| 186 | /// Load one dynamic LoRA adapter and register it as a public model name. |
| 187 | pub async fn load_lora( |
| 188 | &self, |
| 189 | engine_core_client: &EngineCoreClient, |
| 190 | base_model_names: &[String], |
| 191 | mut module: LoraModulePath, |
| 192 | load_inplace: bool, |
| 193 | ) -> Result<LoraRequest, LoadLoraError> { |
| 194 | let allowed_prefixes = runtime_lora_allowed_path_prefixes(); |
| 195 | if let Some(canonical_path) = |
| 196 | validate_lora_path_access(&module.path, allowed_prefixes.as_deref())? |
| 197 | { |
| 198 | module.path = canonical_path; |
| 199 | } |
| 200 | self.register(engine_core_client, base_model_names, module, load_inplace).await |
| 201 | } |
| 202 | |
| 203 | /// Load an operator-configured adapter (`--lora-modules`). Unlike the |
| 204 | /// runtime endpoint, the path is trusted as given: no allowed-prefix |
| 205 | /// check applies. |
| 206 | pub async fn load_static_lora( |
| 207 | &self, |
| 208 | engine_core_client: &EngineCoreClient, |
| 209 | base_model_names: &[String], |
| 210 | module: &LoraModulePath, |
| 211 | ) -> Result<LoraRequest, LoadLoraError> { |
| 212 | self.register(engine_core_client, base_model_names, module.clone(), false).await |
| 213 | } |
| 214 | |
| 215 | async fn register( |
| 216 | &self, |
| 217 | engine_core_client: &EngineCoreClient, |
| 218 | base_model_names: &[String], |
| 219 | module: LoraModulePath, |
| 220 | load_inplace: bool, |
| 221 | ) -> Result<LoraRequest, LoadLoraError> { |
| 222 | let LoraModulePath { |
| 223 | name: lora_name, |
| 224 | path: lora_path, |
| 225 | base_model_name, |
| 226 | is_3d_lora_weight, |
| 227 | } = module; |
| 228 | if lora_name.trim().is_empty() { |
| 229 | bail_invalid_adapter!("lora_name must not be empty"); |
| 230 | } |
| 231 | if lora_path.trim().is_empty() { |
| 232 | bail_invalid_adapter!("lora_path must not be empty"); |
| 233 | } |
| 234 | let _guard = self.update_lock.lock().await; |
| 235 | if base_model_names.iter().any(|name| name == &lora_name) { |
| 236 | return Err(LoadLoraError::BaseModelName { lora_name }); |
| 237 | } |
| 238 | let requests = self.requests.read().await; |
| 239 | let existing_lora_int_id = requests.get(&lora_name).map(|request| request.lora_int_id); |
| 240 | if !load_inplace && existing_lora_int_id.is_some() { |
| 241 | return Err(LoadLoraError::AlreadyLoaded { lora_name }); |
| 242 | } |
| 243 | |
| 244 | let lora_int_id = existing_lora_int_id |
| 245 | .unwrap_or_else(|| self.id_counter.fetch_add(1, Ordering::Relaxed) + 1); |
| 246 | let lora_request = LoraRequest { |
| 247 | lora_name: lora_name.clone(), |
| 248 | lora_int_id, |
| 249 | lora_path, |
| 250 | base_model_name, |
| 251 | tensorizer_config_dict: None, |
| 252 | load_inplace, |
| 253 | is_3d_lora_weight, |
| 254 | }; |
| 255 | drop(requests); |
| 256 | |
| 257 | let loaded = engine_core_client.add_lora(&lora_request).await.map_err(|source| { |
| 258 | LoadLoraError::Engine { |
| 259 | lora_name: lora_name.clone(), |
| 260 | source, |
| 261 | } |
| 262 | })?; |
| 263 | if !loaded { |
| 264 | return Err(LoadLoraError::NotLoaded { lora_name }); |
| 265 | } |
| 266 | self.requests.write().await.insert(lora_name, lora_request.clone()); |
| 267 | Ok(lora_request) |
| 268 | } |
| 269 | |
| 270 | /// Remove one dynamic LoRA adapter from the engine and public model |
| 271 | /// registry. |
| 272 | pub async fn unload_lora( |
| 273 | &self, |
| 274 | engine_core_client: &EngineCoreClient, |
| 275 | lora_name: &str, |
| 276 | requested_lora_int_id: Option<u64>, |
| 277 | ) -> Result<LoraRequest, UnloadLoraError> { |
| 278 | let _guard = self.update_lock.lock().await; |
| 279 | let lora_request = self.requests.read().await.get(lora_name).cloned().ok_or_else(|| { |
| 280 | UnloadLoraError::NotFound { |
| 281 | lora_name: lora_name.to_string(), |
| 282 | } |
| 283 | })?; |
| 284 | |
| 285 | if let Some(actual) = requested_lora_int_id |
| 286 | && actual != lora_request.lora_int_id |
| 287 | { |
| 288 | return Err(UnloadLoraError::IntIdMismatch { |
| 289 | lora_name: lora_name.to_string(), |
| 290 | expected: lora_request.lora_int_id, |
| 291 | actual, |
| 292 | }); |
| 293 | } |
| 294 | |
| 295 | let removed = |
| 296 | engine_core_client |
| 297 | .remove_lora(lora_request.lora_int_id) |
| 298 | .await |
| 299 | .map_err(|source| UnloadLoraError::Engine { |
| 300 | lora_name: lora_request.lora_name.clone(), |
| 301 | source, |
| 302 | })?; |
| 303 | if !removed { |
| 304 | return Err(UnloadLoraError::NotRemoved { |
| 305 | lora_name: lora_request.lora_name, |
| 306 | lora_int_id: lora_request.lora_int_id, |
| 307 | }); |
| 308 | } |
| 309 | |
| 310 | Ok(self.requests.write().await.shift_remove(lora_name).unwrap_or(lora_request)) |
| 311 | } |
| 312 | } |
| 313 | |
| 314 | #[cfg(test)] |
| 315 | mod tests { |
| 316 | use std::fs; |
| 317 | use std::path::PathBuf; |
| 318 | use std::time::{SystemTime, UNIX_EPOCH}; |
| 319 | |
| 320 | use super::validate_lora_path_access; |
| 321 | |
| 322 | fn temp_lora_dir(test_name: &str) -> PathBuf { |
| 323 | let suffix = SystemTime::now() |
| 324 | .duration_since(UNIX_EPOCH) |
| 325 | .expect("clock should be after unix epoch") |
| 326 | .as_nanos(); |
| 327 | let path = std::env::temp_dir().join(format!( |
| 328 | "vllm-lora-{test_name}-{}-{suffix}", |
| 329 | std::process::id() |
| 330 | )); |
| 331 | fs::create_dir_all(&path).expect("create temp lora dir"); |
| 332 | path |
| 333 | } |
| 334 | |
| 335 | #[test] |
| 336 | fn lora_path_allows_hf_repo_ids_without_prefixes() { |
| 337 | assert_eq!( |
| 338 | validate_lora_path_access("org/adapter-a", None).expect("hf repo id should be allowed"), |
| 339 | None |
| 340 | ); |
| 341 | } |
| 342 | |
| 343 | #[test] |
| 344 | fn lora_path_rejects_local_paths_without_prefixes() { |
| 345 | assert!(validate_lora_path_access("/tmp/adapter-a", None).is_err()); |
| 346 | assert!(validate_lora_path_access("./adapter-a", None).is_err()); |
| 347 | assert!(validate_lora_path_access("~/adapter-a", None).is_err()); |
| 348 | assert!(validate_lora_path_access("subdir/../../../etc/sensitive", None).is_err()); |
| 349 | } |
| 350 | |
| 351 | #[test] |
| 352 | fn lora_path_rejects_existing_bare_relative_paths_without_prefixes() { |
| 353 | let root = |
| 354 | PathBuf::from("target").join(format!("vllm-lora-relative-{}", std::process::id())); |
| 355 | let adapter = root.join("adapter-a"); |
| 356 | fs::create_dir_all(&adapter).expect("create relative adapter dir"); |
| 357 | |
| 358 | assert!( |
| 359 | validate_lora_path_access(adapter.to_str().expect("utf-8 temp path"), None).is_err() |
| 360 | ); |
| 361 | |
| 362 | fs::remove_dir_all(root).ok(); |
| 363 | } |
| 364 | |
| 365 | #[test] |
| 366 | fn lora_path_allows_absolute_paths_under_configured_prefixes() { |
| 367 | let root = temp_lora_dir("allowed-prefix"); |
| 368 | let allowed = root.join("allowed"); |
| 369 | let adapter = allowed.join("adapter-a"); |
| 370 | fs::create_dir_all(&adapter).expect("create adapter dir"); |
| 371 | |
| 372 | let prefixes = [allowed]; |
| 373 | let resolved = |
| 374 | validate_lora_path_access(adapter.to_str().expect("utf-8 temp path"), Some(&prefixes)) |
| 375 | .expect("path under configured prefix should be allowed"); |
| 376 | assert_eq!( |
| 377 | resolved.as_deref(), |
| 378 | Some( |
| 379 | adapter |
| 380 | .canonicalize() |
| 381 | .expect("canonical adapter") |
| 382 | .to_str() |
| 383 | .expect("utf-8 temp path") |
| 384 | ) |
| 385 | ); |
| 386 | |
| 387 | fs::remove_dir_all(root).ok(); |
| 388 | } |
| 389 | |
| 390 | #[test] |
| 391 | fn lora_path_rejects_parent_escape_from_configured_prefixes() { |
| 392 | let root = temp_lora_dir("parent-escape"); |
| 393 | let allowed = root.join("allowed"); |
| 394 | let private_adapter = root.join("private").join("adapter-a"); |
| 395 | fs::create_dir_all(&allowed).expect("create allowed dir"); |
| 396 | fs::create_dir_all(&private_adapter).expect("create private adapter dir"); |
| 397 | |
| 398 | let escaped = allowed.join("../private/adapter-a"); |
| 399 | let prefixes = [allowed]; |
| 400 | assert!( |
| 401 | validate_lora_path_access(escaped.to_str().expect("utf-8 temp path"), Some(&prefixes)) |
| 402 | .is_err() |
| 403 | ); |
| 404 | |
| 405 | fs::remove_dir_all(root).ok(); |
| 406 | } |
| 407 | } |