From 170fe57cf046d0e5a0ac75ed3685f6e743becdf8 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Tue, 19 May 2026 09:23:10 +0800 Subject: [PATCH] Encapsulate the pending-flush bookkeeping in a small wrapper (#25727) --- python/sglang/srt/managers/scheduler.py | 59 +++-------------- .../scheduler_components/flush_wrapper.py | 65 +++++++++++++++++++ .../managers/test_scheduler_flush_cache.py | 60 ++++++++++------- 3 files changed, 111 insertions(+), 73 deletions(-) create mode 100644 python/sglang/srt/managers/scheduler_components/flush_wrapper.py diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 43dfa30c1..8bccdd219 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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() diff --git a/python/sglang/srt/managers/scheduler_components/flush_wrapper.py b/python/sglang/srt/managers/scheduler_components/flush_wrapper.py new file mode 100644 index 000000000..2a3b603bc --- /dev/null +++ b/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, + ) diff --git a/test/registered/unit/managers/test_scheduler_flush_cache.py b/test/registered/unit/managers/test_scheduler_flush_cache.py index 3693d725c..9a99cb309 100644 --- a/test/registered/unit/managers/test_scheduler_flush_cache.py +++ b/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( - scheduler, FlushCacheReqInput(timeout_s=None) - ) + output = scheduler.flush_wrapper.handle(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( - scheduler, FlushCacheReqInput(timeout_s=5.0) - ) + output = scheduler.flush_wrapper.handle(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): - output = Scheduler.flush_cache_wrapped(scheduler, req) + with patch( + "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( - scheduler, FlushCacheReqInput(timeout_s=timeout) + output = scheduler.flush_wrapper.handle( + 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): - Scheduler._check_pending_flush(scheduler) + with patch( + "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): - Scheduler._check_pending_flush(scheduler) + with patch( + "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()