test(hicache): cover buffer pipeline idle flush gate

python/sglang/srt/managers/scheduler.py

@@ -4888,9 +4888,8 @@ def is_fully_idle(self, for_health_check=False) -> bool:
                     idle &= len(tc.ongoing_prefetch) == 0
                     idle &= len(tc.ongoing_backup) == 0
                     if get_memory().hicache_host_memory_mode == "buffer_only":
-                        # Queued writes, staged prefetches, and in-flight
-                        # storage writes still hold host staging
-                        # (buffer-mode unified tree only).
+                        # Buffer-mode transfers have additional pending and
+                        # in-flight state owned by the unified tree pipeline.
                         idle &= tc.buffer_pipeline.is_idle()
 
         return idle

test/registered/unit/managers/test_scheduler_buffer_mode_idle.py

@@ -0,0 +1,123 @@
+"""Buffer-mode transfers must drain before destructive cache administration."""
+
+import unittest
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import maybe_stub_sgl_kernel
+
+maybe_stub_sgl_kernel()
+
+from sglang.srt.disaggregation.utils import DisaggregationMode
+from sglang.srt.managers.scheduler import Scheduler
+from sglang.srt.mem_cache.buffer_mode.pipeline import BufferModePipeline
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+class TestSchedulerBufferModeIdle(unittest.TestCase):
+    def setUp(self):
+        self.pipeline = BufferModePipeline.__new__(BufferModePipeline)
+        self.pipeline.reset()
+
+        self.scheduler = Scheduler.__new__(Scheduler)
+        self.scheduler.running_batch = SimpleNamespace(is_empty=lambda: True, reqs=[])
+        self.scheduler.chunked_req = None
+        self.scheduler.dllm_manager = SimpleNamespace(any_staging_reqs=lambda: False)
+        self.scheduler.last_batch = None
+        self.scheduler.enable_overlap = True
+        self.scheduler.result_queue = []
+        self.scheduler.ps = SimpleNamespace(pp_size=1)
+        self.scheduler.waiting_queue = []
+        self.scheduler._engine_paused = False
+        self.scheduler.disaggregation_mode = DisaggregationMode.NULL
+        self.scheduler.grammar_manager = MagicMock(grammar_queue=[])
+        self.scheduler.enable_hisparse = False
+        self.scheduler.enable_hierarchical_cache = True
+        # These cache-mode maps intentionally stay empty: buffer-mode transfers
+        # are owned by the pipeline, not by the similarly named cache maps.
+        self.scheduler.tree_cache = SimpleNamespace(
+            ongoing_write_through={},
+            ongoing_load_back={},
+            ongoing_prefetch={},
+            ongoing_backup={},
+            enable_storage=True,
+            buffer_pipeline=self.pipeline,
+            reset=MagicMock(),
+        )
+        self.scheduler.req_to_token_pool = MagicMock()
+        self.scheduler.token_to_kv_pool_allocator = MagicMock()
+        self.scheduler.metrics_reporter = MagicMock(is_stats_logging_rank=False)
+        self.scheduler.draft_worker = None
+
+        memory = patch(
+            "sglang.srt.managers.scheduler.get_memory",
+            return_value=SimpleNamespace(hicache_host_memory_mode="buffer_only"),
+        )
+        memory.start()
+        self.addCleanup(memory.stop)
+
+    def test_empty_pipeline_is_idle(self):
+        self.assertTrue(self.pipeline.is_idle())
+        self.assertTrue(self.scheduler.is_fully_idle())
+
+    def test_each_buffer_stage_blocks_flush_until_drained(self):
+        stages = (
+            self.pipeline.pending_hit_allocs,
+            self.pipeline.staged_prefetches,
+            self.pipeline.ongoing_buffer_load_back,
+            self.pipeline.pending_write_queue,
+            self.pipeline.inflight_backup_node_ids,
+            self.pipeline.ongoing_write_through,
+            self.pipeline.ongoing_backup,
+        )
+        names = (
+            "pending H2D allocation",
+            "staged prefetch",
+            "H2D",
+            "queued write",
+            "in-flight backup",
+            "D2H",
+            "storage write",
+        )
+        for name, stage in zip(names, stages):
+            with self.subTest(stage=name):
+                if (
+                    stage is self.pipeline.pending_hit_allocs
+                    or stage is self.pipeline.pending_write_queue
+                ):
+                    stage.append(object())
+                elif stage is self.pipeline.inflight_backup_node_ids:
+                    stage.add(1)
+                else:
+                    stage[1] = object()
+
+                try:
+                    self.assertFalse(self.scheduler.flush_cache(empty_cache=False))
+                    self.assertFalse(self.pipeline.is_idle())
+                    self.assertFalse(self.scheduler.is_fully_idle())
+                    self.scheduler.tree_cache.reset.assert_not_called()
+                    self.scheduler.token_to_kv_pool_allocator.clear.assert_not_called()
+                finally:
+                    stage.clear()
+                    self.scheduler.tree_cache.reset.reset_mock()
+                    self.scheduler.token_to_kv_pool_allocator.clear.reset_mock()
+
+                self.assertTrue(self.pipeline.is_idle())
+                self.assertTrue(self.scheduler.flush_cache(empty_cache=False))
+                self.scheduler.tree_cache.reset.assert_called_once()
+                self.scheduler.token_to_kv_pool_allocator.clear.assert_called_once()
+                self.scheduler.tree_cache.reset.reset_mock()
+                self.scheduler.token_to_kv_pool_allocator.clear.reset_mock()
+
+    def test_inflight_transfers_do_not_change_health_check(self):
+        self.pipeline.ongoing_write_through[1] = object()
+        self.pipeline.ongoing_buffer_load_back[-1] = object()
+
+        self.assertFalse(self.scheduler.is_fully_idle())
+        self.assertTrue(self.scheduler.is_fully_idle(for_health_check=True))
+
+
+if __name__ == "__main__":
+    unittest.main()