执行摘要
此 PR 通过为 RequestRegistry 引入一个简单的 active_lora_requests 计数器,在无 LoRA 请求时跳过构造 LoRA 适配器集合的全量扫描,从而优化 Rust 前端的指标上报路径。变更集中于一个文件,逻辑清晰,测试覆盖完整,风险低。
功能与动机
Rust 前端在每个引擎输出批次后会调用 lora_adapter_states() 以更新 vllm:lora_requests_info 指标。该函数遍历所有在途请求(O(N))以收集 LoRA 适配器名称。在无需 LoRA 的部署中,此扫描完全多余,仍会线性消耗 CPU 时间。PR 引入计数器,在活跃 LoRA 请求数为 0 时立即返回空集合,消除不必要的开销。
实现拆解
-
字段新增与初始化:在 RequestRegistry 结构体中添加 active_lora_requests: usize,初始化值为 0。
-
register 方法优化:提前构造 lora 对象,若 lora.is_some() 则递增计数器,并将该对象直接插入 TrackedRequest。
-
lora_adapter_states 守卫条件:方法开头检查 self.active_lora_requests == 0,若成立则直接返回 (BTreeSet::new(), BTreeSet::new())。
-
remove 方法递减:当移除的请求带有 LoRA 时,计数器减 1。
-
close 方法清零:关闭时清零计数器,确保状态一致。
-
测试覆盖:新增 registry_counts_only_active_lora_requests 和 registry_clears_lora_count_on_close 两项测试,验证计数器在无 LoRA、多 LoRA、移除、关闭场景下的正确性。
rust/src/engine-core-client/src/client/state.rs
唯一修改的文件,包含所有核心逻辑变更:新增 active_lora_requests 计数器、优化 register、remove、close 和 lora_adapter_states 方法,以及新增测试。
// File: rust/src/engine-core-client/src/client/state.rs
/// 在 `RequestRegistry` 结构体中新增 `active_lora_requests` 计数器。
#[derive(Debug)]
pub struct RequestRegistry {
closed: bool,
requests: HashMap<String, TrackedRequest>,
active_lora_requests: usize, // <-- 新增:精确跟踪活跃 LoRA 请求数量
routing_per_engine: BTreeMap<EngineId, EngineRoutingState>,
}
impl RequestRegistry {
pub fn new(engines: &[ConnectedEngine]) -> Self {
Self {
closed: false,
requests: HashMap::default(),
active_lora_requests: 0, // 初始化为 0
routing_per_engine: engines
.iter()
.map(|engine| (engine.engine_id.clone(), EngineRoutingState::default()))
.collect(),
}
}
/// 注册新请求时,若携带 LoRA 则递增计数器。
pub fn register(
&mut self,
request_id: String,
lora_name: Option<String>,
data_parallel_rank: Option<u32>,
) -> Result<(EngineId, OutputReceiver)> {
// ... 省略重复检查与引擎选择 ...
let lora = lora_name.map(|adapter_name| LoraRequestState {
adapter_name,
phase: LoraPhase::Waiting,
});
if lora.is_some() {
self.active_lora_requests += 1; // <-- 递增
}
self.requests.insert(
request_id,
TrackedRequest {
sender: tx,
engine_id: engine_id.clone(),
lora, // 使用提前构建的 lora
},
);
// ...
}
/// 移除请求时,若该请求携带 LoRA 则递减计数器。
pub fn remove(&mut self, request_id: &str) -> Option<(OutputSender, EngineId)> {
let tracked = self.requests.remove(request_id)?;
if tracked.lora.is_some() {
self.active_lora_requests -= 1; // <-- 递减
}
// ...
}
/// 关闭注册表时,清除计数器。
pub fn close(&mut self) -> Vec<OutputSender> {
self.closed = true;
self.active_lora_requests = 0; // <-- 清零
// ...
}
/// 核心优化:当没有活跃 LoRA 请求时,立即返回空集合,避免 O(N) 扫描。
pub fn lora_adapter_states(&self) -> (BTreeSet<String>, BTreeSet<String>) {
if self.active_lora_requests == 0 { // <-- 守卫条件
return (BTreeSet::new(), BTreeSet::new());
}
let mut running = BTreeSet::new();
let mut waiting = BTreeSet::new();
for lora in self.requests.values().filter_map(|tracked| tracked.lora.as_ref()) {
match lora.phase {
LoraPhase::Running => running.insert(lora.adapter_name.clone()),
LoraPhase::Waiting => waiting.insert(lora.adapter_name.clone()),
};
}
(running, waiting)
}
}
评论区精华
审查者 BugenZhao 在批准时表示:
"LGTM. Ultimately, we should land something like #45411 to avoid the need to maintain the state here, but this still seems like a good quick fix. Thanks!"
这表明当前方案被认可为一个合理的短期优化,长期方向则可能是更根本的重构。
风险与影响
- 风险:极低。改动仅涉及一个文件,逻辑简单,且通过测试验证。潜在风险在于若未来引入并行访问
RequestRegistry,计数器可能不准确,但当前设计假定单线程访问。
- 影响:在无 LoRA 请求的部署中,
lora_adapter_states 调用变为 O(1),消除了每次引擎输出批次后的 O(N) 扫描。对于高并发场景可显著降低 CPU 开销。
关联脉络
此 PR 与 #45411 存在演进关系:当前方案通过手工维护计数器快速优化,而 #45411 可能从架构层面消除对手工维护计数的需求。两者都是 Rust 前端 LoRA 指标路径的持续优化的一部分。
参与讨论