minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
fzyzcjy
2026-05-19 09:23:10 +08:00
committed by GitHub
parent 5067da3eb6
commit 170fe57cf0
3 changed files with 111 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/scheduler.py
+10 -49
View File
@@ -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
+65
View File
@@ -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
+36 -24
View File
@@ -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()