Fix PP is_fully_idle missing in-flight microbatches (#27446)

python/sglang/srt/managers/scheduler.py

@@ -3330,7 +3330,7 @@ class Scheduler(
             and (self.last_batch is None or self.last_batch.is_empty())
             and (self.cur_batch is None or self.cur_batch.is_empty())
             and (not self.enable_overlap or len(self.result_queue) == 0)
-            and (self.ps.pp_size == 1 or all(x.is_empty() for x in self.running_mbs))
+            and self._pp_microbatches_drained()
         )
 
         # Waiting queues: waiting + bootstrapping + preallocation + kv transfer (decode)
@@ -3366,6 +3366,13 @@ class Scheduler(
 
         return idle
 
+    def _pp_microbatches_drained(self) -> bool:
+        if self.ps.pp_size == 1:
+            return True
+        return all(x.is_empty() for x in self.running_mbs) and all(
+            mb is None or mb.is_empty() for mb in self.mbs
+        )
+
     def attach_hicache_storage_wrapped(
         self, recv_req: AttachHiCacheStorageReqInput
     ) -> AttachHiCacheStorageReqOutput:

python/sglang/test/scripted_runtime/context/queries.py

@@ -12,10 +12,15 @@ def _get_all_reqs(ctx: "ScriptedContext") -> Iterator["Req"]:
     if s.chunked_req is not None:
         yield s.chunked_req
     yield from s.waiting_queue
-    if s.running_batch is not None:
-        yield from s.running_batch.reqs
-    if s.last_batch is not None:
-        yield from s.last_batch.reqs
+    if s.ps.pp_size > 1:
+        for mb in (*s.mbs, *s.last_mbs, *s.running_mbs):
+            if mb is not None:
+                yield from mb.reqs
+    else:
+        if s.running_batch is not None:
+            yield from s.running_batch.reqs
+        if s.last_batch is not None:
+            yield from s.last_batch.reqs
 
 
 def list_active_reqs(ctx: "ScriptedContext") -> List["Req"]:

python/sglang/test/scripted_runtime/scheduler_hook.py

@@ -99,16 +99,13 @@ def _reset_engine_state(ctx: ScriptedContext) -> Generator:
     ctx.abort_all()
     for _ in range(RESET_DRAIN_MAX_STEPS):
         yield
-        if (
-            scheduler.chunked_req is None
-            and len(scheduler.waiting_queue) == 0
-            and scheduler.running_batch.is_empty()
-        ):
+        if scheduler.is_fully_idle():
             break
-
-    server_args = scheduler.server_args
-    for _ in range(2 * (server_args.pp_size + server_args.pp_async_batch_depth)):
-        yield
+    else:
+        raise RuntimeError(
+            "scripted_runtime reset: scheduler did not become fully idle "
+            f"within {RESET_DRAIN_MAX_STEPS} steps"
+        )
 
     ctx.flush_cache()
     yield