Fix set status

AlibabaPAI · Nov 7, 2024 · 4acb199 · 4acb199
1 parent be947b0
commit 4acb199
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/llumnix/backends/vllm/llm_engine.py b/llumnix/backends/vllm/llm_engine.py
@@ -357,11 +357,9 @@ def commit_dst_request(self, backend_request: SequenceGroupLlumnix) -> None:
             "The status of request migrated to dst instance should be  \
              RequestStatus.WAITING_MIGRATING or RequestStatus.RUNNING_MIGRATING"
         if backend_request.status == RequestStatus.WAITING_MIGRATING:
-            self.engine.scheduler.set_status(backend_request, status_to=SequenceStatus.WAITING)
             self.add_waiting_request(backend_request)
         elif backend_request.status == RequestStatus.RUNNING_MIGRATING:
             backend_request.reset_status()
-            self.engine.scheduler.set_status(backend_request, status_to=SequenceStatus.RUNNING)
             self.add_running_request(backend_request)
 
     async def send_blocks(self, dst_ray_actor: "ray.actor.ActorHandle", src_blocks: List[int], dst_blocks: List[int]) -> None:

diff --git a/llumnix/backends/vllm/scheduler.py b/llumnix/backends/vllm/scheduler.py
@@ -145,9 +145,11 @@ def pre_alloc(self,
         return blocks
 
     def add_running_request(self, backend_request: LlumnixRequest) -> None:
+        self._set_status(backend_request, status_to=SequenceStatus.RUNNING)
         self.running.append(backend_request)
 
     def add_waiting_request(self, backend_request: LlumnixRequest) -> None:
+        self._set_status(backend_request, status_to=SequenceStatus.WAITING)
         # pylint: disable=E0203
         self.waiting.append(backend_request)
         fcfs_policy = PolicyFactory.get_policy(policy_name="fcfs")
@@ -162,12 +164,12 @@ def _allocate_and_set_running(self, seq_group: SequenceGroup) -> None:
         # Change seq status to running, but request status is still waiting_migrating.
         if seq_group.status == RequestStatus.WAITING_MIGRATING:
             # For the waiting request migrated in, blocks have already been allocated when pre alloc.
-            self.set_status(seq_group, status_to=SequenceStatus.RUNNING)
+            self._set_status(seq_group, status_to=SequenceStatus.RUNNING)
             seq_group.reset_status()
         else:
             super()._allocate_and_set_running(seq_group)
 
-    def set_status(self,
+    def _set_status(self,
                    seq_group: SequenceGroup,
                    status_to: SequenceStatus,
                    status_from: SequenceStatus = None):