[Metrics] Move LoRA request counts to SchedulerStats

SchedulerStats is the right place for this really, just like
the regular running/waiting counts.

Make sure to call LoRARequestStates.update_scheduler_stats()
even where there was no engine core outputs.

Signed-off-by: Mark McLoughlin <[email protected]>

Author: markmc

tests/v1/engine/test_output_processor.py

@@ -22,11 +22,12 @@
 from vllm.v1.engine import (
     EngineCoreEvent,
     EngineCoreEventType,
+    EngineCoreOutputs,
     EngineCoreRequest,
     FinishReason,
 )
 from vllm.v1.engine.output_processor import OutputProcessor, RequestOutputCollector
-from vllm.v1.metrics.stats import IterationStats
+from vllm.v1.metrics.stats import IterationStats, SchedulerStats
 
 
 def _ref_convert_id_to_token(
@@ -940,21 +941,26 @@ def test_lora_request_tracking(log_stats: bool, dummy_test_vectors):
         output_processor.add_request(request, None)
 
     # First iteration: process outputs with QUEUED events
-    outputs = engine_core.get_outputs()
-    for output in outputs:
+    outputs = EngineCoreOutputs(
+        outputs=engine_core.get_outputs(), scheduler_stats=SchedulerStats()
+    )
+    for output in outputs.outputs:
         output.events = [
             EngineCoreEvent.new_event(EngineCoreEventType.QUEUED, engine_core_timestamp)
         ]
 
     iteration_stats = IterationStats() if log_stats else None
-    output_processor.process_outputs(outputs, engine_core_timestamp, iteration_stats)
+    output_processor.process_outputs(
+        outputs.outputs, engine_core_timestamp, iteration_stats
+    )
+    output_processor.update_scheduler_stats(outputs.scheduler_stats)
 
     if log_stats:
         # Verify waiting counts
-        assert iteration_stats.waiting_lora_adapters.get("lora-1") == 1
-        assert iteration_stats.waiting_lora_adapters.get("lora-2") == 1
-        assert iteration_stats.running_lora_adapters.get("lora-1") == 0
-        assert iteration_stats.running_lora_adapters.get("lora-2") == 0
+        assert outputs.scheduler_stats.waiting_lora_adapters.get("lora-1") == 1
+        assert outputs.scheduler_stats.waiting_lora_adapters.get("lora-2") == 1
+        assert outputs.scheduler_stats.running_lora_adapters.get("lora-1") == 0
+        assert outputs.scheduler_stats.running_lora_adapters.get("lora-2") == 0
         # Verify internal state
         assert len(output_processor.lora_states.requests) == 2
         assert "lora-1" in output_processor.lora_states.requests
@@ -965,76 +971,96 @@ def test_lora_request_tracking(log_stats: bool, dummy_test_vectors):
         assert len(output_processor.lora_states.requests) == 0
 
     # Second iteration: process outputs with SCHEDULED events
-    outputs = engine_core.get_outputs()
-    for output in outputs:
+    outputs = EngineCoreOutputs(
+        outputs=engine_core.get_outputs(), scheduler_stats=SchedulerStats()
+    )
+    for output in outputs.outputs:
         output.events = [
             EngineCoreEvent.new_event(
                 EngineCoreEventType.SCHEDULED, engine_core_timestamp
             )
         ]
 
     iteration_stats = IterationStats() if log_stats else None
-    output_processor.process_outputs(outputs, engine_core_timestamp, iteration_stats)
+    output_processor.process_outputs(
+        outputs.outputs, engine_core_timestamp, iteration_stats
+    )
+    output_processor.update_scheduler_stats(outputs.scheduler_stats)
 
     if log_stats:
         # Verify running counts
-        assert iteration_stats.waiting_lora_adapters.get("lora-1") == 0
-        assert iteration_stats.waiting_lora_adapters.get("lora-2") == 0
-        assert iteration_stats.running_lora_adapters.get("lora-1") == 1
-        assert iteration_stats.running_lora_adapters.get("lora-2") == 1
+        assert outputs.scheduler_stats.waiting_lora_adapters.get("lora-1") == 0
+        assert outputs.scheduler_stats.waiting_lora_adapters.get("lora-2") == 0
+        assert outputs.scheduler_stats.running_lora_adapters.get("lora-1") == 1
+        assert outputs.scheduler_stats.running_lora_adapters.get("lora-2") == 1
     else:
         assert iteration_stats is None
         assert len(output_processor.lora_states.requests) == 0
 
     # Third iteration: finish request-0 (lora-1)
-    outputs = engine_core.get_outputs()
+    outputs = EngineCoreOutputs(
+        outputs=engine_core.get_outputs(), scheduler_stats=SchedulerStats()
+    )
     # Find and mark request-0 as finished (it uses lora-1)
-    for output in outputs:
+    for output in outputs.outputs:
         if output.request_id == "request-0":
             output.finish_reason = FinishReason.LENGTH
             break
 
     iteration_stats = IterationStats() if log_stats else None
-    output_processor.process_outputs(outputs, engine_core_timestamp, iteration_stats)
+    output_processor.process_outputs(
+        outputs.outputs, engine_core_timestamp, iteration_stats
+    )
+    output_processor.update_scheduler_stats(outputs.scheduler_stats)
 
     if log_stats:
         # lora-1 should be removed since no requests remain
         assert "lora-1" not in output_processor.lora_states.requests
         # lora-2 should still be running
-        assert iteration_stats.running_lora_adapters.get("lora-2") == 1
+        assert outputs.scheduler_stats.running_lora_adapters.get("lora-2") == 1
         assert len(output_processor.lora_states.requests) == 1
     else:
         assert len(output_processor.lora_states.requests) == 0
 
     # Fourth iteration: finish request-1 (lora-2)
-    outputs = engine_core.get_outputs()
+    outputs = EngineCoreOutputs(
+        outputs=engine_core.get_outputs(), scheduler_stats=SchedulerStats()
+    )
     # Find and mark request-1 as finished (it uses lora-2)
-    for output in outputs:
+    for output in outputs.outputs:
         if output.request_id == "request-1":
             output.finish_reason = FinishReason.LENGTH
             break
 
     iteration_stats = IterationStats() if log_stats else None
-    output_processor.process_outputs(outputs, engine_core_timestamp, iteration_stats)
+    output_processor.process_outputs(
+        outputs.outputs, engine_core_timestamp, iteration_stats
+    )
+    output_processor.update_scheduler_stats(outputs.scheduler_stats)
 
     if log_stats:
         # lora-2 should be removed since no requests remain
         assert "lora-2" not in output_processor.lora_states.requests
-        assert len(iteration_stats.running_lora_adapters) == 0
+        assert len(outputs.scheduler_stats.running_lora_adapters) == 0
         assert len(output_processor.lora_states.requests) == 0
     else:
         assert len(output_processor.lora_states.requests) == 0
 
     # Finish the last request (no LoRA)
-    outputs = engine_core.get_outputs()
+    outputs = EngineCoreOutputs(
+        outputs=engine_core.get_outputs(), scheduler_stats=SchedulerStats()
+    )
     # Find and mark request-2 as finished (it has no LoRA)
-    for output in outputs:
+    for output in outputs.outputs:
         if output.request_id == "request-2":
             output.finish_reason = FinishReason.LENGTH
             break
 
     iteration_stats = IterationStats() if log_stats else None
-    output_processor.process_outputs(outputs, engine_core_timestamp, iteration_stats)
+    output_processor.process_outputs(
+        outputs.outputs, engine_core_timestamp, iteration_stats
+    )
+    output_processor.update_scheduler_stats(outputs.scheduler_stats)
 
     # Verify all requests are finished
     assert output_processor.get_num_unfinished_requests() == 0

vllm/v1/engine/async_llm.py

@@ -519,6 +519,8 @@ async def output_handler():
                             processed_outputs.reqs_to_abort
                         )
 
+                    output_processor.update_scheduler_stats(outputs.scheduler_stats)
+
                     # 4) Logging.
                     # TODO(rob): make into a coroutine and launch it in
                     # background thread once Prometheus overhead is non-trivial.

vllm/v1/engine/llm_engine.py

@@ -301,6 +301,7 @@ def step(self) -> list[RequestOutput] | list[PoolingRequestOutput]:
             engine_core_timestamp=outputs.timestamp,
             iteration_stats=iteration_stats,
         )
+        self.output_processor.update_scheduler_stats(outputs.scheduler_stats)
 
         # 3) Abort any reqs that finished due to stop strings.
         self.engine_core.abort_requests(processed_outputs.reqs_to_abort)

vllm/v1/engine/output_processor.py

@@ -22,7 +22,12 @@
 from vllm.v1.engine.detokenizer import IncrementalDetokenizer
 from vllm.v1.engine.logprobs import LogprobsProcessor
 from vllm.v1.engine.parallel_sampling import ParentRequest
-from vllm.v1.metrics.stats import IterationStats, LoRARequestStates, RequestStateStats
+from vllm.v1.metrics.stats import (
+    IterationStats,
+    LoRARequestStates,
+    RequestStateStats,
+    SchedulerStats,
+)
 
 
 class RequestOutputCollector:
@@ -477,13 +482,15 @@ def process_outputs(
                 )
                 if self.tracer:
                     self.do_tracing(engine_core_output, req_state, iteration_stats)
-        self.lora_states.update_iteration_stats(iteration_stats)
 
         return OutputProcessorOutput(
             request_outputs=request_outputs,
             reqs_to_abort=reqs_to_abort,
         )
 
+    def update_scheduler_stats(self, scheduler_stats: SchedulerStats | None):
+        self.lora_states.update_scheduler_stats(scheduler_stats)
+
     def do_tracing(
         self,
         engine_core_output: EngineCoreOutput,

vllm/v1/metrics/loggers.py

@@ -974,6 +974,20 @@ def record(
                     scheduler_stats.kv_connector_stats, engine_idx
                 )
 
+            if self.gauge_lora_info is not None:
+                running_lora_adapters = ",".join(
+                    scheduler_stats.running_lora_adapters.keys()
+                )
+                waiting_lora_adapters = ",".join(
+                    scheduler_stats.waiting_lora_adapters.keys()
+                )
+                lora_info_labels = {
+                    self.labelname_running_lora_adapters: running_lora_adapters,
+                    self.labelname_waiting_lora_adapters: waiting_lora_adapters,
+                    self.labelname_max_lora: self.max_lora,
+                }
+                self.gauge_lora_info.labels(**lora_info_labels).set_to_current_time()
+
         if mm_cache_stats is not None:
             self.counter_mm_cache_queries[engine_idx].inc(mm_cache_stats.queries)
             self.counter_mm_cache_hits[engine_idx].inc(mm_cache_stats.hits)
@@ -1037,20 +1051,6 @@ def record(
                     finished_request.max_tokens_param
                 )
 
-        if self.gauge_lora_info is not None:
-            running_lora_adapters = ",".join(
-                iteration_stats.running_lora_adapters.keys()
-            )
-            waiting_lora_adapters = ",".join(
-                iteration_stats.waiting_lora_adapters.keys()
-            )
-            lora_info_labels = {
-                self.labelname_running_lora_adapters: running_lora_adapters,
-                self.labelname_waiting_lora_adapters: waiting_lora_adapters,
-                self.labelname_max_lora: self.max_lora,
-            }
-            self.gauge_lora_info.labels(**lora_info_labels).set_to_current_time()
-
     def record_sleep_state(self, sleep: int = 0, level: int = 0):
         awake = 1
         discard_all = 0

vllm/v1/metrics/stats.py

@@ -170,6 +170,9 @@ class SchedulerStats:
 
     num_corrupted_reqs: int = 0
 
+    waiting_lora_adapters: dict[str, int] = field(default_factory=dict)
+    running_lora_adapters: dict[str, int] = field(default_factory=dict)
+
 
 @dataclass
 class RequestStateStats:
@@ -219,8 +222,6 @@ def __init__(self):
         self.n_params_iter: list[int] = []
         self.time_to_first_tokens_iter: list[float] = []
         self.inter_token_latencies_iter: list[float] = []
-        self.waiting_lora_adapters: dict[str, int] = {}
-        self.running_lora_adapters: dict[str, int] = {}
 
     def __repr__(self) -> str:
         field_to_value_str = ", ".join(f"{k}={v}" for k, v in vars(self).items())
@@ -393,10 +394,9 @@ def request_running(self, req_id: str, lora_name: str | None):
     def request_finished(self, req_id: str, lora_name: str | None):
         self._request_update(req_id, lora_name, waiting=False, running=False)
 
-    def update_iteration_stats(self, iteration_stats: IterationStats | None):
-        if not self.log_stats:
+    def update_scheduler_stats(self, scheduler_stats: SchedulerStats | None):
+        if not self.log_stats or scheduler_stats is None:
             return
-        assert iteration_stats is not None
         for lora_name, stats in self.requests.items():
-            iteration_stats.waiting_lora_adapters[lora_name] = len(stats.waiting)
-            iteration_stats.running_lora_adapters[lora_name] = len(stats.running)
+            scheduler_stats.waiting_lora_adapters[lora_name] = len(stats.waiting)
+            scheduler_stats.running_lora_adapters[lora_name] = len(stats.running)