fix unfinished

Author: grimoire

lmdeploy/pytorch/engine/engine_loop.py

@@ -2,6 +2,7 @@
 import asyncio
 import logging
 import time
+from collections.abc import Callable
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Optional
 
@@ -55,25 +56,32 @@ def clear(self):
 class RunableEventAsync:
     """Awaitable async runable event."""
 
-    def __init__(self, scheduler: 'Scheduler'):
+    def __init__(self, scheduler: 'Scheduler', extra_runable_checker: Callable[[], bool] | None = None):
         self.scheduler = scheduler
+        self.extra_runable_checker = extra_runable_checker
         self.event = asyncio.Event()
 
+    def has_unfinished(self):
+        """Check whether scheduler or engine-local state has runnable work."""
+        if self.scheduler.has_unfinished():
+            return True
+        return self.extra_runable_checker is not None and self.extra_runable_checker()
+
     async def wait(self):
         """Wait event."""
         await self.event.wait()
 
     def set(self):
         """Set event."""
-        if self.scheduler.has_unfinished():
+        if self.has_unfinished():
             self.event.set()
         else:
             self.event.clear()
 
 
-def build_runable_event(scheduler: 'Scheduler'):
+def build_runable_event(scheduler: 'Scheduler', extra_runable_checker: Callable[[], bool] | None = None):
     """Build runable event."""
-    return RunableEventAsync(scheduler)
+    return RunableEventAsync(scheduler, extra_runable_checker)
 
 
 @dataclass
@@ -128,7 +136,9 @@ def __init__(self,
         self.resp_queue = asyncio.Queue()
         self.forward_event = CounterEvent()
         self.migration_event = asyncio.Event()
-        self.has_runable_event = RunableEventAsync(self.scheduler)
+        # Active long-context chunks are owned by InputsMaker, not the
+        # scheduler WAITING/READY queues, so include them in the runnable gate.
+        self.has_runable_event = RunableEventAsync(self.scheduler, self.inputs_maker.has_pending_long_context_chunk)
         # Sleep uses a small handshake with the scheduling loops:
         # 1. sleep() sets _sleep_requested and waits for main/migration drain events.
         # 2. main_loop and migration_loop reach safe boundaries, acknowledge
@@ -383,13 +393,12 @@ def __get_logprobs(batched_outputs: 'BatchedOutputs'):
 
     async def _main_loop_try_send_next_inputs(self):
         """Try send next inputs."""
-        scheduler = self.scheduler
-        if not scheduler.has_unfinished():
+        if not self.has_runable_event.has_unfinished():
             await self.has_runable_event.wait()
         if self._sleep_requested:
             return None, None
 
-        scheduler.collect_migration_done()
+        self.scheduler.collect_migration_done()
         return await self.inputs_maker.send_next_inputs()
 
     @staticmethod

lmdeploy/pytorch/engine/inputs_maker.py

@@ -346,6 +346,11 @@ def _has_pending_last_long_context_chunk(self):
         left."""
         return self.long_context_chunker.enabled() and self.long_context_chunker.is_last_chunk()
 
+    def has_pending_long_context_chunk(self):
+        """Check whether engine-local long-context chunk work can run."""
+        self.long_context_chunker.check_enable()
+        return self.long_context_chunker.enabled()
+
     def _should_defer_long_context_chunk(self, prefill: bool):
         """Check whether the active long-context chunk should yield this
         loop."""

tests/pytorch/engine/test_inputs_maker.py

@@ -7,7 +7,7 @@
 
 import lmdeploy.pytorch.engine.inputs_maker as inputs_maker_module
 from lmdeploy.pytorch.disagg.config import EngineRole
-from lmdeploy.pytorch.engine.engine_loop import EngineLoop
+from lmdeploy.pytorch.engine.engine_loop import EngineLoop, RunableEventAsync
 from lmdeploy.pytorch.engine.inputs_maker import (
     InputsMakerAsync,
     InputsMakerConfig,
@@ -206,6 +206,38 @@ async def get_output_async(self):
     assert not block_trie.pinned
 
 
+def test_engine_loop_treats_pending_long_context_chunk_as_runnable():
+    events = []
+
+    class _Scheduler:
+
+        def has_unfinished(self):
+            return False
+
+        def collect_migration_done(self):
+            events.append('collect_migration_done')
+
+    class _InputsMaker:
+
+        def has_pending_long_context_chunk(self):
+            return True
+
+        async def send_next_inputs(self):
+            events.append('send_next_inputs')
+            return 'forward_inputs', ['long-seq']
+
+    loop = EngineLoop.__new__(EngineLoop)
+    loop.scheduler = _Scheduler()
+    loop.inputs_maker = _InputsMaker()
+    loop.has_runable_event = RunableEventAsync(loop.scheduler, loop.inputs_maker.has_pending_long_context_chunk)
+    loop._sleep_requested = False
+
+    result = asyncio.run(asyncio.wait_for(loop._main_loop_try_send_next_inputs(), timeout=1.0))
+
+    assert result == ('forward_inputs', ['long-seq'])
+    assert events == ['collect_migration_done', 'send_next_inputs']
+
+
 def _make_policy_maker(long_seq, decode_seq=None):
     maker = InputsMakerAsync.__new__(InputsMakerAsync)
     maker.config = SimpleNamespace(role=EngineRole.Decode)