@@ -97,13 +97,13 @@ def _prepare_decode_task(self, request):
97
97
98
98
def _prepare_preempt_task (self , request ):
99
99
return ScheduledPreemptTask (idx = request .idx , request_id = request .request_id )
100
-
100
+
101
101
def reschedule_preempt_task (self , request_id ):
102
102
with self .lock :
103
103
if request_id in self .to_be_rescheduled_request_id_set and request_id in self .requests :
104
104
request = self .requests [request_id ]
105
105
self .waiting .appendleft (request )
106
- self .to_be_rescheduled_request_id_set .remove (request_id )
106
+ self .to_be_rescheduled_request_id_set .remove (request_id )
107
107
108
108
def _trigger_preempt (self , request , num_new_blocks , preempted_reqs , scheduled_reqs ):
109
109
can_schedule = True
@@ -421,9 +421,15 @@ def finish_requests(self, request_ids: Union[str, Iterable[str]]):
421
421
self .running .remove (request )
422
422
request .status = RequestStatus .FINISHED
423
423
self ._free_blocks (request )
424
- if request .request_id in self .to_be_rescheduled_request_id_set : # finished after preempted, blocks have been recycled.
425
- self .to_be_rescheduled_request_id_set .remove (request .request_id ) # just remove from to_be_rescheduled_request_id_set
426
- if request in self .waiting : # after finished, this request still scheduled from preempted to waiting, unexpected error, should not be here
424
+ if (
425
+ request .request_id in self .to_be_rescheduled_request_id_set
426
+ ): # finished after preempted, blocks have been recycled.
427
+ self .to_be_rescheduled_request_id_set .remove (
428
+ request .request_id
429
+ ) # just remove from to_be_rescheduled_request_id_set
430
+ if (
431
+ request in self .waiting
432
+ ): # after finished, this request still scheduled from preempted to waiting, unexpected error, should not be here
427
433
raise RuntimeError (f"request { request .request_id } scheduled into waiting list, after finished" )
428
434
429
435
self .tasks_list [request .idx ] = None
0 commit comments