fix scheduler bug in release2.1 (#3295)

This commit is contained in:
chenjian
2025-08-10 13:55:22 +08:00
committed by GitHub
parent 86ff68be4b
commit c000cff744
3 changed files with 32 additions and 8 deletions

View File

@@ -75,6 +75,7 @@ class ResourceManagerV1(ResourceManager):
self.running: list[Request] = []
self.finish_execution_pool = ThreadPoolExecutor(max_workers=1)
self.lock = threading.Lock()
self.to_be_rescheduled_request_id_set = set()
def allocated_slots(self, request: Request):
return len(request.block_tables) * self.config.cache_config.block_size
@@ -96,6 +97,13 @@ class ResourceManagerV1(ResourceManager):
def _prepare_preempt_task(self, request):
return ScheduledPreemptTask(idx=request.idx, request_id=request.request_id)
def reschedule_preempt_task(self, request_id):
with self.lock:
if request_id in self.to_be_rescheduled_request_id_set and request_id in self.requests:
request = self.requests[request_id]
self.waiting.appendleft(request)
self.to_be_rescheduled_request_id_set.remove(request_id)
def _trigger_preempt(self, request, num_new_blocks, preempted_reqs, scheduled_reqs):
can_schedule = True
@@ -106,7 +114,7 @@ class ResourceManagerV1(ResourceManager):
preempted_req.num_computed_tokens = 0
preempted_req.prefill_block_num = 0
self._free_blocks(preempted_req)
self.waiting.appendleft(preempted_req)
self.to_be_rescheduled_request_id_set.add(preempted_req.request_id)
preempted_reqs.append(preempted_req)
scheduled_reqs.append(self._prepare_preempt_task(preempted_req))
if preempted_req == request:
@@ -381,8 +389,9 @@ class ResourceManagerV1(ResourceManager):
return False
def add_request(self, request: Request) -> None:
self.waiting.append(request)
self.requests[request.request_id] = request
with self.lock:
self.waiting.append(request)
self.requests[request.request_id] = request
def _free_blocks(self, request: Request):
if self.config.cache_config.enable_prefix_caching:
@@ -409,9 +418,14 @@ class ResourceManagerV1(ResourceManager):
if request is None:
# Invalid request ID.
continue
request.status = RequestStatus.FINISHED
self.running.remove(request)
self._free_blocks(request)
if request in self.running: # normally run and finished
self.running.remove(request)
request.status = RequestStatus.FINISHED
self._free_blocks(request)
if request.request_id in self.to_be_rescheduled_request_id_set: # finished after preempted, blocks have been recycled.
self.to_be_rescheduled_request_id_set.remove(request.request_id) # just remove from to_be_rescheduled_request_id_set
if request in self.waiting: # after finished, this request still scheduled from preempted to waiting, unexpected error, should not be here
raise RuntimeError(f"request {request.request_id} scheduled into waiting list, after finished")
self.tasks_list[request.idx] = None
self.stop_flags[request.idx] = True
del self.requests[req_id]