Refactor PyTorch scheduler ownership and API boundaries

lmdeploy/pytorch/disagg/conn/engine_conn.py

@@ -92,9 +92,7 @@ async def handle_zmq_recv(self, remote_engine_id: str):
                 logger.error(f'invalid zmq request from {remote_engine_id}: {e}')
                 continue
             session_id = req.remote_session_id
-            if session_id in self.engine.scheduler.sessions:
-                self.engine.end_session(session_id=session_id)
-            else:
+            if not self.engine.end_session(session_id=session_id):
                 logger.error(f'invalid free, {remote_engine_id}, {session_id}')
 
     async def zmq_disconnect(self, remote_engine_id: str):

lmdeploy/pytorch/engine/engine.py

@@ -345,7 +345,7 @@ def _on_add_session(self, reqs: list[Request], **kwargs):
             session_id = req.data['session_id']
             resp = req.data.get('response', True)
             resp_type = ResponseType.SESSION_REPEAT
-            if session_id not in self.scheduler.sessions:
+            if self.scheduler.get_session(session_id) is None:
                 self.scheduler.add_session(session_id)
                 resp_type = ResponseType.SUCCESS
             if resp:
@@ -357,8 +357,8 @@ def _on_stop_session(self, reqs: list[Request], **kwargs):
             session_id = req.data['session_id']
             resp = req.data.get('response', True)
             resp_type = ResponseType.SESSION_NOT_EXIST
-            if session_id in self.scheduler.sessions:
-                session = self.scheduler.sessions[session_id]
+            session = self.scheduler.get_session(session_id)
+            if session is not None:
                 stopped_resp_ids = set()
                 for seq in session.sequences.values():
                     if seq.status not in (MessageStatus.STOPPED, MessageStatus.TO_BE_MIGRATED):
@@ -412,10 +412,11 @@ def _on_end_session(self, reqs: list[Request], **kwargs):
             session_id = req.data['session_id']
             resp = req.data.get('response', True)
             resp_type = ResponseType.SESSION_NOT_EXIST
-            if session_id in self.scheduler.sessions:
-                msgs = list(self.scheduler.sessions[session_id].sequences.values())
+            session = self.scheduler.get_session(session_id)
+            if session is not None:
+                msgs = list(session.sequences.values())
                 if len(msgs) > 0 and msgs[0].preserve_cache:
-                    msgs[0].state.finish()
+                    msgs[0].finish()
                 else:
                     self.end_session(session_id)
                 resp_type = ResponseType.SUCCESS
@@ -428,7 +429,7 @@ def _on_add_message(self, reqs: list[Request], **kwargs):
         for req in reqs:
             req_data = req.data
             session_id = req_data['session_id']
-            if self.scheduler and session_id not in self.scheduler.sessions:
+            if self.scheduler and self.scheduler.get_session(session_id) is None:
                 self._response(req.resp, ResponseType.SESSION_NOT_EXIST)
                 continue
             valid_reqs.append(req)
@@ -481,7 +482,7 @@ def __update_max_new_tokens(msg):
         scheduler = self.scheduler
         for req in reqs:
             session_id = req.data['session_id']
-            sess = scheduler.sessions.get(session_id, None)
+            sess = scheduler.get_session(session_id)
             if sess is None:
                 self._response(req.resp, ResponseType.SESSION_NOT_EXIST)
                 continue
@@ -510,7 +511,7 @@ def __update_max_new_tokens(msg):
                     mode=UpdateTokenMode.INPUTS,
                 )
                 msg.sampling_param = sampling_param
-                msg.state.activate()
+                msg.activate()
 
             __update_max_new_tokens(msg)
             msg.resp = req.resp
@@ -576,8 +577,9 @@ def _unblock_new_inputs(self):
     def _cancel_and_end_all_sessions(self):
         """Cancel active responses and remove all scheduler sessions."""
         num_cancelled = 0
-        session_ids = list(self.scheduler.sessions.keys())
-        for session in list(self.scheduler.sessions.values()):
+        sessions = self.scheduler.get_sessions()
+        session_ids = [session.session_id for session in sessions]
+        for session in sessions:
             for seq in list(session.sequences.values()):
                 resp: Response = getattr(seq, 'resp', None)
                 if resp is None or resp.is_done:
@@ -605,7 +607,7 @@ async def sleep(self, level: int = 1):
         # cancel all remain sessions
         self._cancel_and_end_all_sessions()
         await self.executor.sleep(level)
-        self.scheduler.finish_deferred_kv_transfers_after_worker_drain()
+        self.scheduler.finish_kv_transfers_after_worker_drain()
         if self._engine_loop is not None:
             self._engine_loop.reset_runtime_state()
         logger.info('PyTorch engine entered sleep: level=%s, sleeping_tags=%s.', level, sorted(self._sleeping_tags))
@@ -714,8 +716,9 @@ def start_loop(self):
 
     def end_session(self, session_id: int):
         """End session."""
-        if session_id in self.scheduler.sessions:
-            has_multimodal = self._has_multimodal_session(self.scheduler.sessions[session_id])
+        session = self.scheduler.get_session(session_id)
+        if session is not None:
+            has_multimodal = self._has_multimodal_session(session)
             self.scheduler.end_session(session_id)
             self._maybe_trim_multimodal_session(has_multimodal)
             return True

lmdeploy/pytorch/engine/engine_loop.py

@@ -24,6 +24,7 @@
     from lmdeploy.pytorch.engine.model_agent import BatchedOutputs
     from lmdeploy.pytorch.model_inputs import ModelInputs, ModelInputsDelta
     from lmdeploy.pytorch.paging import Scheduler
+    from lmdeploy.pytorch.paging.block_trie.checkpoint_lifecycle import StateCheckpointLifecycle
     from lmdeploy.pytorch.strategies.base.sequence import SequenceStrategy
 
     from .engine import Engine, SeqList
@@ -117,13 +118,15 @@ class EngineLoop:
     def __init__(self,
                  req_manager: 'RequestManager',
                  scheduler: 'Scheduler',
+                 state_checkpoints: 'StateCheckpointLifecycle',
                  executor: 'ExecutorBase',
                  seq_strategy: 'SequenceStrategy',
                  inputs_maker: 'InputsMakerAsync',
                  config: EngineLoopConfig,
                  engine_conn: Optional['EngineP2PConnection'] = None):
         self.req_manager = req_manager
         self.scheduler = scheduler
+        self.state_checkpoints = state_checkpoints
         self.executor = executor
         self.seq_strategy = seq_strategy
         self.inputs_maker = inputs_maker
@@ -323,7 +326,7 @@ def __get_logprobs(batched_outputs: 'BatchedOutputs'):
             seq.append_routed_experts(all_routed_experts)
             seq.append_logits(logits)
             seq.append_ce_loss(ce_loss, finish=False)
-            self.scheduler.block_trie.cache_routed_experts_for_seq(seq)
+            self.scheduler.cache_routed_experts([seq])
             return dict()
 
         new_token_timestamp = batched_outputs.new_token_timestamp
@@ -337,7 +340,7 @@ def __get_logprobs(batched_outputs: 'BatchedOutputs'):
                                          batched_outputs=batched_outputs,
                                          model_inputs=model_inputs,
                                          delta=delta)
-        self.scheduler.block_trie.cache_routed_experts(running)
+        self.scheduler.cache_routed_experts(running)
 
         # generate output
         outputs: dict[int, InferOutput] = dict()
@@ -354,7 +357,7 @@ def __get_logprobs(batched_outputs: 'BatchedOutputs'):
                 continue
             session_id = msg.session_id
             if msg.resp_cache:
-                cache_block_ids = self.scheduler.block_manager.get_block_table(msg).tolist()
+                cache_block_ids = self.scheduler.get_block_tables([msg])[0].tolist()
             else:
                 cache_block_ids = None
 
@@ -403,14 +406,12 @@ async def _main_loop_try_send_next_inputs(self):
         if self._sleep_requested:
             return None, None
 
-        self.scheduler.collect_migration_done()
         return await self.inputs_maker.send_next_inputs()
 
     async def _prefetch_next_inputs(self):
-        """Collect migration completions before prefetching the next batch."""
+        """Prefetch the next batch unless sleep has started."""
         if self._sleep_requested:
             return None, None
-        self.scheduler.collect_migration_done()
         return await self.inputs_maker.prefetch_next_inputs()
 
     async def _wait_for_schedulable_prefill(self):
@@ -422,22 +423,11 @@ async def _wait_for_schedulable_prefill(self):
             # warning or adding the full pressure backoff to TTFT.
             await asyncio.sleep(0.001)
             return
+        cache_usage = scheduler.schedule_metrics.cache_usage
         logger.warning(f'no next prefill running request, Maybe cache is full, '
-                       f'free gpu cache blocks: {scheduler.block_manager.get_num_free_gpu_blocks()}, '
-                       f'total gpu cache blocks: {scheduler.block_manager.num_gpu_blocks}')
+                       f'gpu cache usage: {cache_usage:.1%}')
         await asyncio.sleep(0.1)
 
-    def _publish_forward_checkpoints(self, running: 'SeqList', has_state_checkpoint_save: bool):
-        """Publish per-forward prefix-cache ownership before prefetching."""
-        state_checkpoints = self.scheduler.block_trie.state_checkpoints
-        if has_state_checkpoint_save:
-            state_checkpoints.publish_saves(running, pin_saves=True)
-        state_checkpoints.unpin_restores(running)
-
-    def _release_forward_save_pins(self, running: 'SeqList'):
-        """Unpin producers after the forward output/event boundary."""
-        self.scheduler.block_trie.state_checkpoints.unpin_saves(running)
-
     def _finish_forward_output(self,
                                out: 'BatchedOutputs | None',
                                running: 'SeqList',
@@ -446,16 +436,14 @@ def _finish_forward_output(self,
         """Apply connector progress and publish model outputs."""
         if out is None:
             return
-        # A connector polling step intentionally has no token output. Consume
-        # its transfer completions first so newly loaded requests become
-        # schedulable even when no model forward ran in this executor step.
+        # Connector-only polls have no token output; apply completions before
+        # returning.
         self.scheduler.update_connector_output(out.kv_connector_output)
         if out.next_token_ids is None:
             return
         step_outputs = self._make_infer_outputs(out, running=running, model_inputs=model_inputs, delta=delta)
-        # Sequence history is advanced by _make_infer_outputs. Only now can the
-        # scheduler prove that a prefill reached its reserved target and release
-        # the soft block reservation used while admitting external KV loads.
+        # _make_infer_outputs advances history; only then can soft reservations
+        # be released.
         self.scheduler.release_completed_prefill_reservations(running)
         self.resp_queue.put_nowait(step_outputs)
 
@@ -481,11 +469,14 @@ async def _main_loop_get_outputs(
         # for GPU output; save checkpoints keep a producer pin until the output
         # event boundary so prefetch cannot evict/reuse their destination slots.
         if has_model_work:
-            self._publish_forward_checkpoints(running, has_state_checkpoint_save)
+            self.state_checkpoints.finish_forward_dispatch(
+                running,
+                has_save_plan=has_state_checkpoint_save,
+            )
         forward_inputs, next_running = await self._prefetch_next_inputs()
         out = await self.executor.get_output_async()
         if has_model_work:
-            self._release_forward_save_pins(running)
+            self.state_checkpoints.unpin_saves(running)
         self._finish_forward_output(out, running, model_inputs, delta)
         # out might come from shared memory, need to explicitly delete to release memory in time
         del out
@@ -528,7 +519,7 @@ async def main_loop(self):
                 running=next_running,
                 forward_inputs=forward_inputs,
             )
-            self.inputs_maker.deactivate_evict_seqs()
+            self.inputs_maker.preempt_invalid_decode_seqs()
             has_runable_event.set()
 
     def update_running_migration(self, running: 'SeqList', next_token_ids: np.ndarray, stopped: torch.Tensor,
@@ -546,7 +537,7 @@ def update_running_migration(self, running: 'SeqList', next_token_ids: np.ndarra
             if stop:
                 update_token = _EMPTY_TOKEN
                 msg.update_token_ids(update_token, model_meta=model_meta, mode=UpdateTokenMode.PREFILL)
-                msg.state.finish()
+                msg.finish()
 
     async def _migration_loop_migrate(self, migration_ready: 'SeqList'):
         """Migration loop migrate."""
@@ -558,7 +549,7 @@ async def _migration_loop_migrate(self, migration_ready: 'SeqList'):
             migration_execution_requests: list[tuple[int, list[tuple[int, int]]]] = []
             migration_request = msg.migration_request
             prefill_block_ids = migration_request.remote_block_ids
-            decode_block_ids = list(self.scheduler.block_manager.get_block_table(msg=msg))
+            decode_block_ids = list(self.scheduler.get_block_tables([msg])[0])
 
             assert len(prefill_block_ids) == len(decode_block_ids), (
                 f'#prefill block ids ({len(prefill_block_ids)}) must equal to '
@@ -615,7 +606,7 @@ async def migration_loop(self):
                 await self._sleep_resume_event.wait()
                 continue
 
-            migration_ready = self.scheduler._schedule_migration()
+            migration_ready = self.scheduler.schedule_migration()
             if not migration_ready and not self.scheduler.has_migration_waiting():
                 await self.migration_event.wait()
             elif migration_ready:
@@ -692,6 +683,7 @@ def build_engine_loop(engine: 'Engine'):
     return EngineLoop(
         req_manager=engine.req_manager,
         scheduler=engine.scheduler,
+        state_checkpoints=engine.scheduler.state_checkpoints,
         executor=engine.executor,
         seq_strategy=engine.seq_strategy,
         inputs_maker=inputs_maker,