Encapsulate the pending-flush bookkeeping in a small wrapper (#25727)

python/sglang/srt/managers/scheduler.py

@@ -103,7 +103,6 @@ from sglang.srt.managers.io_struct import (
     ExpertDistributionReqOutput,
     ExpertDistributionReqType,
     FlushCacheReqInput,
-    FlushCacheReqOutput,
     FreezeGCReq,
     GetInternalStateReq,
     GetInternalStateReqOutput,
@@ -165,6 +164,9 @@ from sglang.srt.managers.scheduler_components.batch_result_processor import (
 from sglang.srt.managers.scheduler_components.dp_attn import (
     SchedulerDPAttnAdapter,
 )
+from sglang.srt.managers.scheduler_components.flush_wrapper import (
+    SchedulerFlushWrapper,
+)
 from sglang.srt.managers.scheduler_components.idle_sleeper import IdleSleeper
 from sglang.srt.managers.scheduler_components.invariant_checker import (
     SchedulerInvariantChecker,
@@ -991,7 +993,11 @@ class Scheduler(
         self.last_batch: Optional[ScheduleBatch] = None
         self.forward_ct = 0
         self.return_health_check_ipcs: Deque[Optional[str]] = deque()
-        self._pending_flush: Optional[Tuple[FlushCacheReqInput, float]] = None
+        self.flush_wrapper = SchedulerFlushWrapper(
+            flush_cache=self.flush_cache,
+            is_fully_idle=self.is_fully_idle,
+            ipc_channels=self.ipc_channels,
+        )
         self.session_controller = SessionController(self.tree_cache)
         self.forward_sleep_time = None
         self._engine_paused = False
@@ -1355,7 +1361,7 @@ class Scheduler(
                 (TokenizedEmbeddingReqInput, self.handle_embedding_request),
                 (BatchTokenizedGenerateReqInput, self.handle_batch_generate_request),
                 (BatchTokenizedEmbeddingReqInput, self.handle_batch_embedding_request),
-                (FlushCacheReqInput, self.flush_cache_wrapped),
+                (FlushCacheReqInput, self.flush_wrapper.handle),
                 (ClearHiCacheReqInput, self.clear_hicache_storage_wrapped),
                 (AttachHiCacheStorageReqInput, self.attach_hicache_storage_wrapped),
                 (DetachHiCacheStorageReqInput, self.detach_hicache_storage_wrapped),
@@ -1629,7 +1635,7 @@ class Scheduler(
                     if self.ipc_channels.recv_from_rpc is not None:
                         self.ipc_channels.recv_from_rpc.send_pyobj(output)
 
-        self._check_pending_flush()
+        self.flush_wrapper.check_pending()
         if self.external_corpus_manager is not None:
             self.external_corpus_manager.check_pending_load()
 
@@ -3010,32 +3016,6 @@ class Scheduler(
                 )
             )
 
-    def _check_pending_flush(self):
-        if self._pending_flush is None:
-            return
-
-        pending_req, deadline = self._pending_flush
-
-        if self.is_fully_idle():
-            success = self.flush_cache()
-            self._pending_flush = None
-            self.ipc_channels.send_to_tokenizer.send_output(
-                FlushCacheReqOutput(success=success), pending_req
-            )
-            return
-
-        if time.monotonic() >= deadline:
-            logging.warning(
-                "Deferred flush_cache timed out while waiting for idle state."
-            )
-            self._pending_flush = None
-            self.ipc_channels.send_to_tokenizer.send_output(
-                FlushCacheReqOutput(
-                    success=False, message="Timed out waiting for idle state."
-                ),
-                pending_req,
-            )
-
     def add_external_corpus(
         self, recv_req: AddExternalCorpusReqInput
     ) -> Optional[AddExternalCorpusReqOutput]:
@@ -3066,25 +3046,6 @@ class Scheduler(
             )
         return self.external_corpus_manager.list(recv_req)
 
-    def flush_cache_wrapped(
-        self, recv_req: FlushCacheReqInput
-    ) -> Optional[FlushCacheReqOutput]:
-        if self._pending_flush is not None:
-            return FlushCacheReqOutput(
-                success=False,
-                message="Another flush_cache is already in progress.",
-            )
-
-        timeout_s = float(recv_req.timeout_s or 0.0)
-        if timeout_s <= 0.0:
-            return FlushCacheReqOutput(success=self.flush_cache())
-
-        if self.is_fully_idle():
-            return FlushCacheReqOutput(success=self.flush_cache())
-
-        self._pending_flush = (recv_req, time.monotonic() + timeout_s)
-        return None
-
     def clear_hicache_storage_wrapped(self, recv_req: ClearHiCacheReqInput):
         if self.enable_hierarchical_cache:
             self.tree_cache.clear_storage_backend()

python/sglang/srt/managers/scheduler_components/flush_wrapper.py

@@ -0,0 +1,65 @@
+import logging
+import time
+from typing import Callable, Optional, Tuple
+
+from sglang.srt.managers.io_struct import FlushCacheReqInput, FlushCacheReqOutput
+from sglang.srt.managers.scheduler_components.ipc_channels import (
+    SchedulerIpcChannels,
+)
+
+
+class SchedulerFlushWrapper:
+    def __init__(
+        self,
+        *,
+        flush_cache: Callable[[], bool],
+        is_fully_idle: Callable[[], bool],
+        ipc_channels: SchedulerIpcChannels,
+    ) -> None:
+        self._flush_cache = flush_cache
+        self._is_fully_idle = is_fully_idle
+        self._ipc_channels = ipc_channels
+        self._pending: Optional[Tuple[FlushCacheReqInput, float]] = None
+
+    def handle(self, recv_req: FlushCacheReqInput) -> Optional[FlushCacheReqOutput]:
+        if self._pending is not None:
+            return FlushCacheReqOutput(
+                success=False,
+                message="Another flush_cache is already in progress.",
+            )
+
+        timeout_s = float(recv_req.timeout_s or 0.0)
+        if timeout_s <= 0.0:
+            return FlushCacheReqOutput(success=self._flush_cache())
+
+        if self._is_fully_idle():
+            return FlushCacheReqOutput(success=self._flush_cache())
+
+        self._pending = (recv_req, time.monotonic() + timeout_s)
+        return None
+
+    def check_pending(self) -> None:
+        if self._pending is None:
+            return
+
+        pending_req, deadline = self._pending
+
+        if self._is_fully_idle():
+            success = self._flush_cache()
+            self._pending = None
+            self._ipc_channels.send_to_tokenizer.send_output(
+                FlushCacheReqOutput(success=success), pending_req
+            )
+            return
+
+        if time.monotonic() >= deadline:
+            logging.warning(
+                "Deferred flush_cache timed out while waiting for idle state."
+            )
+            self._pending = None
+            self._ipc_channels.send_to_tokenizer.send_output(
+                FlushCacheReqOutput(
+                    success=False, message="Timed out waiting for idle state."
+                ),
+                pending_req,
+            )

test/registered/unit/managers/test_scheduler_flush_cache.py

@@ -8,6 +8,9 @@ maybe_stub_sgl_kernel()
 
 from sglang.srt.managers.io_struct import FlushCacheReqInput
 from sglang.srt.managers.scheduler import Scheduler
+from sglang.srt.managers.scheduler_components.flush_wrapper import (
+    SchedulerFlushWrapper,
+)
 
 register_cpu_ci(est_time=14, suite="base-a-test-cpu")
 
@@ -15,10 +18,14 @@ register_cpu_ci(est_time=14, suite="base-a-test-cpu")
 class TestSchedulerFlushCache(unittest.TestCase):
     def _new_scheduler(self) -> Scheduler:
         scheduler = Scheduler.__new__(Scheduler)
-        scheduler._pending_flush = None
         scheduler.ipc_channels = MagicMock()
         scheduler.flush_cache = MagicMock(return_value=True)
         scheduler.is_fully_idle = MagicMock(return_value=False)
+        scheduler.flush_wrapper = SchedulerFlushWrapper(
+            flush_cache=scheduler.flush_cache,
+            is_fully_idle=scheduler.is_fully_idle,
+            ipc_channels=scheduler.ipc_channels,
+        )
         return scheduler
 
     def test_immediate_flush_no_timeout(self):
@@ -26,9 +33,7 @@ class TestSchedulerFlushCache(unittest.TestCase):
         scheduler = self._new_scheduler()
         scheduler.flush_cache.return_value = False
 
-        output = Scheduler.flush_cache_wrapped(
+        output = scheduler.flush_wrapper.handle(FlushCacheReqInput(timeout_s=None))
-            scheduler, FlushCacheReqInput(timeout_s=None)
-        )
 
         self.assertFalse(output.success)
         scheduler.flush_cache.assert_called_once()
@@ -38,9 +43,7 @@ class TestSchedulerFlushCache(unittest.TestCase):
         scheduler = self._new_scheduler()
         scheduler.is_fully_idle.return_value = True
 
-        output = Scheduler.flush_cache_wrapped(
+        output = scheduler.flush_wrapper.handle(FlushCacheReqInput(timeout_s=5.0))
-            scheduler, FlushCacheReqInput(timeout_s=5.0)
-        )
 
         self.assertTrue(output.success)
         scheduler.flush_cache.assert_called_once()
@@ -50,22 +53,25 @@ class TestSchedulerFlushCache(unittest.TestCase):
         scheduler = self._new_scheduler()
         req = FlushCacheReqInput(timeout_s=3.0)
 
-        with patch("sglang.srt.managers.scheduler.time.monotonic", return_value=10.0):
+        with patch(
-            output = Scheduler.flush_cache_wrapped(scheduler, req)
+            "sglang.srt.managers.scheduler_components.flush_wrapper.time.monotonic",
+            return_value=10.0,
+        ):
+            output = scheduler.flush_wrapper.handle(req)
 
         self.assertIsNone(output)
-        pending_req, deadline = scheduler._pending_flush
+        pending_req, deadline = scheduler.flush_wrapper._pending
         self.assertIs(pending_req, req)
         self.assertEqual(deadline, 13.0)
 
     def test_rejects_when_already_pending(self):
         """Any new request is rejected while another is pending."""
         scheduler = self._new_scheduler()
-        scheduler._pending_flush = (FlushCacheReqInput(timeout_s=10.0), 999.0)
+        scheduler.flush_wrapper._pending = (FlushCacheReqInput(timeout_s=10.0), 999.0)
 
         for timeout in [None, 5.0]:
-            output = Scheduler.flush_cache_wrapped(
+            output = scheduler.flush_wrapper.handle(
-                scheduler, FlushCacheReqInput(timeout_s=timeout)
+                FlushCacheReqInput(timeout_s=timeout)
             )
             self.assertFalse(output.success)
             self.assertIn("already in progress", output.message)
@@ -76,11 +82,11 @@ class TestSchedulerFlushCache(unittest.TestCase):
         scheduler = self._new_scheduler()
         scheduler.is_fully_idle.return_value = True
         req = FlushCacheReqInput(timeout_s=1.0)
-        scheduler._pending_flush = (req, 111.0)
+        scheduler.flush_wrapper._pending = (req, 111.0)
 
-        Scheduler._check_pending_flush(scheduler)
+        scheduler.flush_wrapper.check_pending()
 
-        self.assertIsNone(scheduler._pending_flush)
+        self.assertIsNone(scheduler.flush_wrapper._pending)
         scheduler.flush_cache.assert_called_once()
         out = scheduler.ipc_channels.send_to_tokenizer.send_output.call_args.args[0]
         self.assertTrue(out.success)
@@ -88,12 +94,15 @@ class TestSchedulerFlushCache(unittest.TestCase):
     def test_pending_flush_expires_on_timeout(self):
         scheduler = self._new_scheduler()
         req = FlushCacheReqInput(timeout_s=1.0)
-        scheduler._pending_flush = (req, 99.0)
+        scheduler.flush_wrapper._pending = (req, 99.0)
 
-        with patch("sglang.srt.managers.scheduler.time.monotonic", return_value=100.0):
+        with patch(
-            Scheduler._check_pending_flush(scheduler)
+            "sglang.srt.managers.scheduler_components.flush_wrapper.time.monotonic",
+            return_value=100.0,
+        ):
+            scheduler.flush_wrapper.check_pending()
 
-        self.assertIsNone(scheduler._pending_flush)
+        self.assertIsNone(scheduler.flush_wrapper._pending)
         scheduler.flush_cache.assert_not_called()
         out = scheduler.ipc_channels.send_to_tokenizer.send_output.call_args.args[0]
         self.assertFalse(out.success)
@@ -101,12 +110,15 @@ class TestSchedulerFlushCache(unittest.TestCase):
     def test_pending_flush_survives_before_deadline(self):
         scheduler = self._new_scheduler()
         req = FlushCacheReqInput(timeout_s=5.0)
-        scheduler._pending_flush = (req, 101.0)
+        scheduler.flush_wrapper._pending = (req, 101.0)
 
-        with patch("sglang.srt.managers.scheduler.time.monotonic", return_value=100.0):
+        with patch(
-            Scheduler._check_pending_flush(scheduler)
+            "sglang.srt.managers.scheduler_components.flush_wrapper.time.monotonic",
+            return_value=100.0,
+        ):
+            scheduler.flush_wrapper.check_pending()
 
-        self.assertIsNotNone(scheduler._pending_flush)
+        self.assertIsNotNone(scheduler.flush_wrapper._pending)
         scheduler.ipc_channels.send_to_tokenizer.send_output.assert_not_called()