22
33import time
44from dataclasses import dataclass , field
5- from typing import TYPE_CHECKING , List
5+ from typing import TYPE_CHECKING , Dict , List , Optional , Set
66
77if TYPE_CHECKING :
88 from vllm .outputs import RequestOutput
99 from vllm .v1 .engine import EngineCoreEvent , EngineCoreOutput , FinishReason
10+ from vllm .v1 .output_processor import RequestState
1011
1112
1213@dataclass
@@ -36,6 +37,12 @@ class SchedulerStats:
3637 default_factory = PrefixCacheStats )
3738
3839
40+ @dataclass
41+ class LoRAStats :
42+ waiting_requests : Set [str ] = field (default_factory = set )
43+ running_requests : Set [str ] = field (default_factory = set )
44+
45+
3946@dataclass
4047class RequestStateStats :
4148 """Stats that need to be tracked across delta updates."""
@@ -76,14 +83,17 @@ def __init__(self):
7683 self .time_per_output_tokens_iter : List [float ] = []
7784 self .queue_times_iter : List [float ] = []
7885 self .prefill_times_iter : List [float ] = []
86+ self .waiting_lora_adapters : Dict [str , int ] = {}
87+ self .running_lora_adapters : Dict [str , int ] = {}
7988
8089 def _time_since (self , start : float ) -> float :
8190 """Calculate an interval relative to this iteration's timestamp."""
8291 return self .iteration_timestamp - start
8392
8493 def update_from_output (self , output : "EngineCoreOutput" ,
8594 engine_core_timestamp : float , is_prefilling : bool ,
86- prompt_len : int , req_stats : RequestStateStats ):
95+ prompt_len : int , req_stats : RequestStateStats ,
96+ lora_stats : Optional [LoRAStats ]):
8797 num_new_generation_tokens = len (output .new_token_ids )
8898
8999 self .num_generation_tokens += num_new_generation_tokens
@@ -105,7 +115,8 @@ def update_from_output(self, output: "EngineCoreOutput",
105115
106116 # Process request-level engine core events
107117 if output .events is not None :
108- self .update_from_events (output .events , is_prefilling , req_stats )
118+ self .update_from_events (output .request_id , output .events ,
119+ is_prefilling , req_stats , lora_stats )
109120
110121 # Process the batch-level "new tokens" engine core event
111122 if is_prefilling :
@@ -123,17 +134,21 @@ def update_from_output(self, output: "EngineCoreOutput",
123134 if num_new_generation_tokens > 0 :
124135 req_stats .last_token_ts = engine_core_timestamp
125136
126- def update_from_events (self , events : List ["EngineCoreEvent" ],
127- is_prefilling : bool , req_stats : RequestStateStats ):
137+ def update_from_events (self , req_id : str , events : List ["EngineCoreEvent" ],
138+ is_prefilling : bool , req_stats : RequestStateStats ,
139+ lora_stats : Optional [LoRAStats ]):
128140 # Avoid circular dependency
129141 from vllm .v1 .engine import EngineCoreEventType
130142 for event in events :
131143 if event .type == EngineCoreEventType .QUEUED :
132144 req_stats .queued_ts = event .timestamp
145+ if lora_stats is not None :
146+ lora_stats .waiting_requests .add (req_id )
133147 elif event .type == EngineCoreEventType .SCHEDULED :
134148 queued_interval = event .timestamp - req_stats .queued_ts
135149 self .queue_times_iter .append (queued_interval )
136150 req_stats .scheduled_ts = event .timestamp
151+ LoRARequestStates .scheduled_request (lora_stats , req_id )
137152
138153 def update_from_finished_request (self , finish_reason : "FinishReason" ,
139154 request_output : "RequestOutput" ,
@@ -151,3 +166,55 @@ def update_from_finished_request(self, finish_reason: "FinishReason",
151166 inference_time = inference_time ,
152167 decode_time = decode_time )
153168 self .finished_requests .append (finished_req )
169+
170+
171+ class LoRARequestStates :
172+ """Per-LoRA request state stats."""
173+
174+ def __init__ (self ):
175+ self .lora_name_to_stats : Dict [str , LoRAStats ] = {}
176+
177+ def get_stats (self , req_state : 'RequestState' ) -> Optional [LoRAStats ]:
178+ if req_state .lora_name is None :
179+ return None
180+ if req_state .lora_name not in self .lora_name_to_stats :
181+ self .lora_name_to_stats [req_state .lora_name ] = LoRAStats ()
182+ return self .lora_name_to_stats [req_state .lora_name ]
183+
184+ def add_request (self , req_state : 'RequestState' ):
185+ if (lora_stats := self .get_stats (req_state )) is not None :
186+ lora_stats .waiting_requests .add (req_state .request_id )
187+
188+ def finish_request (self , req_state : 'RequestState' ):
189+ if req_state .lora_name is None :
190+ return
191+ lora_stats = self .lora_name_to_stats [req_state .lora_name ]
192+ lora_stats .running_requests .remove (req_state .request_id )
193+
194+ def abort_request (self , req_state : 'RequestState' ):
195+ if req_state .lora_name is None :
196+ return
197+ lora_stats = self .lora_name_to_stats [req_state .lora_name ]
198+ lora_stats .waiting_requests .discard (req_state .request_id )
199+ lora_stats .running_requests .discard (req_state .request_id )
200+
201+ # Break the pattern for this lifecycle methods so we can
202+ # call this from IterationStats.update_from_events()
203+ @staticmethod
204+ def scheduled_request (lora_stats : Optional [LoRAStats ], request_id : str ):
205+ if lora_stats is None :
206+ return
207+ lora_stats .waiting_requests .remove (request_id )
208+ lora_stats .running_requests .add (request_id )
209+
210+ def update_iteration_stats (self ,
211+ iteration_stats : Optional [IterationStats ]):
212+ if iteration_stats is None :
213+ return
214+ for lora_name , stats in self .lora_name_to_stats .items ():
215+ if stats .waiting_requests :
216+ iteration_stats .waiting_lora_adapters [lora_name ] = \
217+ len (stats .waiting_requests )
218+ if stats .running_requests :
219+ iteration_stats .running_lora_adapters [lora_name ] = \
220+ len (stats .running_requests )
0 commit comments