[Rust Frontend] Avoid LoRA registry scans without active LoRA requests (#47040)
Signed-off-by: reidliu41 <reid201711@gmail.com>
Reid committed
Jun 30, 2026 at 12:58 UTC
a16dbd5b8572d4128be9f10b9dcff4999b594b25
1 file changed
+80
-4
rust/src/engine-core-client/src/client/state.rs
+80
-4
@@ -100,6 +100,7 @@ impl EngineRoutingState {
100
pub struct RequestRegistry {
101
closed: bool,
102
requests: HashMap<String, TrackedRequest>,
103
+ active_lora_requests: usize,
104
routing_per_engine: BTreeMap<EngineId, EngineRoutingState>,
105
}
106
@@ -108,6 +109,7 @@ impl RequestRegistry {
109
Self {
110
closed: false,
111
requests: HashMap::default(),
112
+ active_lora_requests: 0,
113
routing_per_engine: engines
114
.iter()
115
.map(|engine| (engine.engine_id.clone(), EngineRoutingState::default()))
@@ -133,15 +135,19 @@ impl RequestRegistry {
135
136
let engine_id = self.choose_engine_for_request(data_parallel_rank)?;
137
let (tx, rx) = mpsc::unbounded_channel();
138
+ let lora = lora_name.map(|adapter_name| LoraRequestState {
139
+ adapter_name,
140
+ phase: LoraPhase::Waiting,
141
+ });
142
+ if lora.is_some() {
143
+ self.active_lora_requests += 1;
144
+ }
145
self.requests.insert(
146
request_id,
147
TrackedRequest {
148
sender: tx,
149
engine_id: engine_id.clone(),
141
- lora: lora_name.map(|adapter_name| LoraRequestState {
142
- adapter_name,
143
- phase: LoraPhase::Waiting,
144
- }),
150
+ lora,
151
},
152
);
153
@@ -230,6 +236,10 @@ impl RequestRegistry {
236
/// Snapshot the adapter names of tracked LoRA requests as
237
/// (running, waiting) sets. Feeds the `vllm:lora_requests_info` gauge.
238
pub fn lora_adapter_states(&self) -> (BTreeSet<String>, BTreeSet<String>) {
239
+ if self.active_lora_requests == 0 {
240
+ return (BTreeSet::new(), BTreeSet::new());
241
+ }
242
+
243
let mut running = BTreeSet::new();
244
let mut waiting = BTreeSet::new();
245
for lora in self.requests.values().filter_map(|tracked| tracked.lora.as_ref()) {
@@ -283,6 +293,7 @@ impl RequestRegistry {
293
}
294
295
self.closed = true;
296
+ self.active_lora_requests = 0;
297
std::mem::take(&mut self.requests)
298
.into_values()
299
.map(|tracked| tracked.sender)
@@ -322,6 +333,9 @@ impl RequestRegistry {
333
#[must_use]
334
pub fn remove(&mut self, request_id: &str) -> Option<(OutputSender, EngineId)> {
335
let tracked = self.requests.remove(request_id)?;
336
+ if tracked.lora.is_some() {
337
+ self.active_lora_requests -= 1;
338
+ }
339
self.routing_per_engine
340
.get_mut(&tracked.engine_id)
341
.expect("request registry must track all known engines")
@@ -359,6 +373,11 @@ impl RequestRegistry {
373
pub fn is_closed(&self) -> bool {
374
self.closed
375
}
376
+
377
+ #[cfg(test)]
378
+ fn active_lora_requests(&self) -> usize {
379
+ self.active_lora_requests
380
+ }
381
}
382
383
/// Internal registry for tracking active utility calls and their waiting
@@ -574,6 +593,63 @@ mod tests {
593
);
594
}
595
596
+ #[test]
597
+ fn registry_counts_only_active_lora_requests() {
598
+ let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]);
599
+
600
+ registry.register("req-plain".to_string(), None, None).unwrap();
601
+ assert_eq!(registry.active_lora_requests(), 0);
602
+ assert_eq!(
603
+ registry.lora_adapter_states(),
604
+ (adapter_names(&[]), adapter_names(&[]))
605
+ );
606
+
607
+ registry
608
+ .register(
609
+ "req-lora-a".to_string(),
610
+ Some("adapter-a".to_string()),
611
+ None,
612
+ )
613
+ .unwrap();
614
+ registry
615
+ .register(
616
+ "req-lora-b".to_string(),
617
+ Some("adapter-b".to_string()),
618
+ None,
619
+ )
620
+ .unwrap();
621
+ assert_eq!(registry.active_lora_requests(), 2);
622
+
623
+ drop(registry.remove("req-plain"));
624
+ assert_eq!(registry.active_lora_requests(), 2);
625
+
626
+ drop(registry.finish_many(&["req-lora-a".to_string()]));
627
+ assert_eq!(registry.active_lora_requests(), 1);
628
+
629
+ drop(registry.abort_many(&["req-lora-b".to_string()], 0.0));
630
+ assert_eq!(registry.active_lora_requests(), 0);
631
+ assert_eq!(
632
+ registry.lora_adapter_states(),
633
+ (adapter_names(&[]), adapter_names(&[]))
634
+ );
635
+ }
636
+
637
+ #[test]
638
+ fn registry_clears_lora_count_on_close() {
639
+ let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]);
640
+ registry
641
+ .register("req-lora".to_string(), Some("adapter-a".to_string()), None)
642
+ .unwrap();
643
+
644
+ assert_eq!(registry.active_lora_requests(), 1);
645
+ drop(registry.close());
646
+ assert_eq!(registry.active_lora_requests(), 0);
647
+ assert_eq!(
648
+ registry.lora_adapter_states(),
649
+ (adapter_names(&[]), adapter_names(&[]))
650
+ );
651
+ }
652
+
653
#[test]
654
fn registry_drops_lora_tracking_on_abort() {
655
let mut registry = RequestRegistry::new(&[connected_engine(EngineId::from(b"engine-0"))]);