diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 972575b71..ebf388519 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -2656,6 +2656,7 @@ class SchedulerDisaggregationDecodeMixin: self.process_batch_result(batch, result) else: # When the server is idle, do self-check and re-init some states + self._sched_idled = True self.on_idle() # Update last_batch @@ -2707,6 +2708,7 @@ class SchedulerDisaggregationDecodeMixin: self.result_queue.append((batch.copy(), batch_result)) else: batch_result = None + self._sched_idled = True # Process the last batch if self.last_batch: diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 18d069408..abb247f3d 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -659,6 +659,7 @@ class SchedulerDisaggregationPrefillMixin: result = self.run_batch(batch) self.process_batch_result(batch, result) else: + self._sched_idled = True self.on_idle() self.process_disagg_prefill_inflight_queue() @@ -700,6 +701,7 @@ class SchedulerDisaggregationPrefillMixin: self.result_queue.append((batch.copy(), batch_result)) else: batch_result = None + self._sched_idled = True # Process the last batch if self.last_batch: diff --git a/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py b/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py index 62d602523..2619ec4d7 100644 --- a/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py +++ b/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py @@ -98,6 +98,8 @@ class SchedulerMlxOverlapMixin: self.forward_ct += 1 batch.forward_iter = self.forward_ct batch.launch_ts = time.monotonic() + batch.after_idle_gap = self._sched_idled + self._sched_idled = False self.profiler_manager._profile_batch_predicate(batch) def _finalize_mlx_pending_job(self: Scheduler, pending: MlxPendingJob): @@ -273,6 +275,7 @@ class SchedulerMlxOverlapMixin: pending_curr = _launch_fresh(next_batch) self.result_queue.append(pending_curr) else: + self._sched_idled = True self.on_idle() self.last_batch = next_batch diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index 2566554c4..5fdd1fc05 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -184,6 +184,7 @@ class SchedulerPPMixin: # When the server is idle, self-check and re-init some states if server_is_idle: + self._sched_idled = True self.on_idle() @DynamicGradMode() @@ -362,8 +363,10 @@ class SchedulerPPMixin: self.running_batch.batch_is_full = False # When the server is idle, self-check and re-init some states - if server_is_idle and len(self.disagg_prefill_inflight_queue) == 0: - self.on_idle() + if server_is_idle: + self._sched_idled = True + if len(self.disagg_prefill_inflight_queue) == 0: + self.on_idle() @DynamicGradMode() def event_loop_pp_disagg_decode(self: Scheduler): @@ -552,8 +555,10 @@ class SchedulerPPMixin: if get_disagg().disaggregation_decode_enable_offload_kvcache: queue_size += len(self.decode_offload_manager.ongoing_offload) - if server_is_idle and queue_size == 0: - self.on_idle() + if server_is_idle: + self._sched_idled = True + if queue_size == 0: + self.on_idle() def init_pp_loop_state(self: Scheduler): self.pp_loop_size: int = self.ps.pp_size + get_parallel().pp_async_batch_depth diff --git a/python/sglang/srt/multiplex/multiplexing_mixin.py b/python/sglang/srt/multiplex/multiplexing_mixin.py index a52a736bb..d2f4024fa 100644 --- a/python/sglang/srt/multiplex/multiplexing_mixin.py +++ b/python/sglang/srt/multiplex/multiplexing_mixin.py @@ -135,6 +135,7 @@ class SchedulerMultiplexMixin: stream_idx > 0 and running_batch.is_empty() ) if running_batch.is_empty() and self.split_prefill_batch is None: + self._sched_idled = True self.on_idle() if adjust_stream_group: diff --git a/test/registered/unit/hardware_backend/mlx/test_attention_patching.py b/test/registered/unit/hardware_backend/mlx/test_attention_patching.py index 47b64f848..38964a655 100644 --- a/test/registered/unit/hardware_backend/mlx/test_attention_patching.py +++ b/test/registered/unit/hardware_backend/mlx/test_attention_patching.py @@ -1158,6 +1158,7 @@ class TestMlxOverlapScheduler(unittest.TestCase): scheduler.gracefully_exit = False scheduler._engine_paused = False scheduler.forward_ct = 0 + scheduler._sched_idled = False scheduler.profiler_manager = SimpleNamespace( _profile_batch_predicate=lambda batch: None ) diff --git a/test/registered/unit/hardware_backend/mlx/test_scheduler_mixin.py b/test/registered/unit/hardware_backend/mlx/test_scheduler_mixin.py index f94872db1..4a1ec2ea3 100644 --- a/test/registered/unit/hardware_backend/mlx/test_scheduler_mixin.py +++ b/test/registered/unit/hardware_backend/mlx/test_scheduler_mixin.py @@ -32,6 +32,7 @@ class TestMlxLaunchBookkeeping(unittest.TestCase): def _make_scheduler(self): scheduler = MagicMock() scheduler.forward_ct = 0 + scheduler._sched_idled = False result = MagicMock() result.next_token_ids = None scheduler.tp_worker.finalize_mlx_result.return_value = result @@ -111,6 +112,7 @@ class TestOverlapLoopStampsLaunchTs(unittest.TestCase): scheduler = MagicMock() scheduler.forward_ct = 0 + scheduler._sched_idled = False scheduler._prepare_mlx_launch.side_effect = lambda batch: ( SchedulerMlxOverlapMixin._prepare_mlx_launch(scheduler, batch) ) @@ -269,13 +271,16 @@ class TestOverlapLoopGracefulExit(unittest.TestCase): scheduler._engine_paused = False scheduler.waiting_queue = [] scheduler.result_queue = deque() - scheduler.ingest_requests.side_effect = recv_side_effect - # Model handle_shutdown: processing a non-empty recv batch (the - # ShutdownReq) flips the flag; the loop must notice at the top of the - # next iteration instead of polling forever. - scheduler.process_input_requests.side_effect = lambda reqs: ( - setattr(scheduler, "gracefully_exit", True) if reqs else None - ) + recv_requests = MagicMock(side_effect=recv_side_effect) + + def ingest_requests(): + reqs = recv_requests() + # Ingestion dispatches the shutdown request before the loop resumes. + if reqs: + scheduler.gracefully_exit = True + return reqs + + scheduler.ingest_requests.side_effect = ingest_requests plan = MagicMock() plan.batch_to_run = None scheduler.get_next_batch_to_run.return_value = plan diff --git a/test/registered/unit/managers/test_disagg_idle_step_counters.py b/test/registered/unit/managers/test_disagg_idle_step_counters.py new file mode 100644 index 000000000..781f3b0db --- /dev/null +++ b/test/registered/unit/managers/test_disagg_idle_step_counters.py @@ -0,0 +1,390 @@ +"""Idle gaps must not inflate prefill busy time or decode timing samples.""" + +import importlib.util +import sys +import unittest +from collections import deque +from contextlib import nullcontext +from pathlib import Path +from types import MethodType, ModuleType, SimpleNamespace +from unittest.mock import Mock, patch + +from parameterized import parameterized + +from sglang.srt.managers import scheduler as scheduler_module +from sglang.srt.managers.schedule_batch import ScheduleBatch +from sglang.srt.managers.scheduler import Scheduler +from sglang.srt.managers.utils import GenerationBatchResult +from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.srt.speculative.spec_info import SpeculativeAlgorithm +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + +LAUNCH_TIMESTAMPS = (0.0, 0.125, 1.0, 1.125) +PP_MODULE = "sglang.srt.managers.scheduler_pp_mixin" +PDMUX_MODULE = "sglang.srt.multiplex.multiplexing_mixin" + + +class _BeforeModelForward(Exception): + pass + + +def load_mlx_scheduler_module(): + # Run the real Python loop on CPU without importing a Metal runtime. Use a + # private module name so this cannot replace the installed MLX scheduler. + path = ( + Path(scheduler_module.__file__).parents[1] + / "hardware_backend/mlx/scheduler_mixin.py" + ) + spec = importlib.util.spec_from_file_location("_idle_counter_mlx_scheduler", path) + module = importlib.util.module_from_spec(spec) + core = ModuleType("mlx.core") + mlx = ModuleType("mlx") + mlx.core = core + with patch.dict(sys.modules, {spec.name: module, "mlx": mlx, "mlx.core": core}): + spec.loader.exec_module(module) + return module + + +class TestSchedulerIdleStepCounters(CustomTestCase): + @parameterized.expand( + [ + ( + "prefill_normal", + Scheduler.event_loop_normal_disagg_prefill, + ForwardMode.EXTEND, + ), + ( + "prefill_overlap", + Scheduler.event_loop_overlap_disagg_prefill, + ForwardMode.EXTEND, + ), + ( + "decode_normal", + Scheduler.event_loop_normal_disagg_decode, + ForwardMode.DECODE, + ), + ( + "decode_overlap", + Scheduler.event_loop_overlap_disagg_decode, + ForwardMode.DECODE, + ), + ] + ) + def test_disagg_idle_gap(self, name, event_loop, mode): + # One no-batch iteration flushes an overlap result without on_idle; + # two iterations also exercise the fully idle branch. + for idle_iterations in (0, 1, 2): + with self.subTest(idle_iterations=idle_iterations): + batches = self.make_batches(mode) + schedule = batches[:2] + [None] * idle_iterations + batches[2:] + [None] + scheduler = self.make_scheduler(schedule) + self.run_and_check(scheduler, event_loop, mode, idle_iterations > 0) + + @parameterized.expand( + [ + ("unified", Scheduler.event_loop_pp, ForwardMode.EXTEND), + ("prefill", Scheduler.event_loop_pp_disagg_prefill, ForwardMode.EXTEND), + ("decode", Scheduler.event_loop_pp_disagg_decode, ForwardMode.DECODE), + ] + ) + def test_pp_idle_cycles(self, name, event_loop, mode): + for depth in (0, 1): + loop_size = 2 + depth + for pending_transfer in (False,) if name == "unified" else (False, True): + for pattern in ("consecutive", "empty_slot", "idle_cycle"): + with self.subTest( + depth=depth, pending_transfer=pending_transfer, pattern=pattern + ): + batches = self.make_batches(mode) + padding = [None] * (loop_size - 2) + schedule = batches[:2] + padding + if pattern == "idle_cycle": + schedule += [None] * loop_size + schedule += batches[2:] + padding + if pattern == "empty_slot": + schedule = [ + slot + for batch in batches + for slot in [batch] + [None] * (loop_size - 1) + ] + schedule += [None] * loop_size # Drain the last result. + scheduler = self.make_scheduler(schedule) + self.prepare_pp_scheduler(scheduler) + if pending_transfer: + if name == "prefill": + scheduler.disagg_prefill_inflight_queue = [object()] + else: + scheduler.disagg_decode_transfer_queue.queue = [ + object() + ] + parallel = SimpleNamespace( + pp_async_batch_depth=depth, + enable_dsa_prefill_context_parallel=False, + ) + with ( + patch(f"{PP_MODULE}.get_parallel", return_value=parallel), + patch( + f"{PP_MODULE}.get_disagg", + return_value=SimpleNamespace( + disaggregation_decode_enable_offload_kvcache=False + ), + ), + patch(f"{PP_MODULE}.set_time_batch"), + ): + self.run_and_check( + scheduler, event_loop, mode, pattern == "idle_cycle" + ) + # Transfers suppress housekeeping, not the idle flag. + self.assertEqual( + scheduler.on_idle.call_count, + 0 if pending_transfer else 1 + (pattern == "idle_cycle"), + ) + + def test_pdmux_idle_gap(self): + for idle_iterations in (0, 1, 2): + with self.subTest(idle_iterations=idle_iterations): + batches = self.make_batches(ForwardMode.DECODE) + schedule = batches[:2] + [None] * idle_iterations + batches[2:] + [None] + scheduler = self.make_scheduler(schedule) + scheduler.stream_groups = [(Mock(), Mock())] + scheduler.sm_counts = [(80, 52)] + scheduler.split_prefill_batch = None + scheduler.update_split_prefill_batch = lambda sm_count, running_batch: ( + False, + running_batch, + ) + scheduler.update_running_batch = Mock( + side_effect=[batch or ScheduleBatch(reqs=[]) for batch in schedule] + ) + with ( + patch(f"{PDMUX_MODULE}.get_current_stream_idx", return_value=0), + patch(f"{PDMUX_MODULE}.set_pdmux_status"), + patch(f"{PDMUX_MODULE}.torch.cuda.empty_cache"), + patch( + f"{PDMUX_MODULE}.torch.cuda.stream", + side_effect=lambda stream: nullcontext(), + ), + ): + self.run_and_check( + scheduler, + Scheduler.event_loop_pdmux, + ForwardMode.DECODE, + idle_iterations > 0, + ) + + @parameterized.expand([("fresh", False), ("chained", True)]) + def test_mlx_idle_gap(self, name, chained): + module = load_mlx_scheduler_module() + mixin = module.SchedulerMlxOverlapMixin + mode = ForwardMode.DECODE if chained else ForwardMode.EXTEND + for idle_iterations in (0, 1, 2): + with self.subTest(idle_iterations=idle_iterations): + batches = self.make_batches(mode) + if chained: + # Finishing the request drains its already-launched lookahead + # before scheduling fresh work, allowing the loop to go idle. + for batch in batches: + batch.reqs[0].finished.return_value = True + schedule = ( + [batches[0]] + [None] * idle_iterations + [batches[2], None] + ) + else: + schedule = ( + batches[:2] + [None] * idle_iterations + batches[2:] + [None] + ) + scheduler = self.make_scheduler(schedule) + scheduler.gracefully_exit = False + scheduler.future_map = None + scheduler.result_queue = deque() + scheduler._prepare_mlx_launch = MethodType( + mixin._prepare_mlx_launch, scheduler + ) + scheduler._finalize_mlx_pending_job = MethodType( + mixin._finalize_mlx_pending_job, scheduler + ) + scheduler._mlx_batch_chain_safe = MethodType( + mixin._mlx_batch_chain_safe, scheduler + ) + scheduler.profiler_manager = SimpleNamespace( + _profile_batch_predicate=Mock() + ) + launch = SimpleNamespace( + mode="decode" if chained else "extend", + decode=object() if chained else None, + ) + scheduler.tp_worker = SimpleNamespace( + async_forward_batch_generation_mlx=Mock(return_value=launch), + async_chained_decode_mlx=Mock(return_value=launch), + finalize_mlx_result=Mock(return_value=GenerationBatchResult()), + ) + with ( + patch.object(module, "resolve_forward_inputs"), + patch.object( + module, + "get_device", + return_value=SimpleNamespace(mlx_enable_sampling=False), + ), + patch.object( + module.time, "monotonic", side_effect=LAUNCH_TIMESTAMPS + ), + ): + self.run_and_check( + scheduler, + mixin.event_loop_overlap_mlx, + mode, + idle_iterations > 0, + ) + self.assertEqual( + scheduler.tp_worker.async_chained_decode_mlx.call_count, + 2 if chained else 0, + ) + + def make_batches(self, mode): + return [ + ScheduleBatch( + reqs=[ + SimpleNamespace(rid="request", finished=Mock(return_value=False)) + ], + forward_mode=mode, + spec_algorithm=SpeculativeAlgorithm.NONE, + launch_ts=launch_ts, + extend_num_tokens=1024, + ) + for launch_ts in LAUNCH_TIMESTAMPS + ] + + def run_and_check(self, scheduler, event_loop, mode, after_idle): + observed_idle_flags = [] + observed_iters = [] + + def run_batch(batch, pp_proxy_tensors=None): + # Exercise the real timestamp, iteration, and flag handoff. Only + # model execution is stopped, at the scripted pre-forward hook. + with patch( + "sglang.srt.managers.scheduler.time.monotonic", + return_value=batch.launch_ts, + ): + with self.assertRaises(_BeforeModelForward): + Scheduler.run_batch(scheduler, batch, pp_proxy_tensors) + return GenerationBatchResult() + + def process_batch_result(batch, result): + observed_idle_flags.append(batch.after_idle_gap) + observed_iters.append(batch.forward_iter) + scheduler._record_step_counters(batch, result) + + scheduler.run_batch = run_batch + scheduler.process_batch_result = process_batch_result + with self.assertRaises(StopIteration): + event_loop(scheduler) + + self.assertEqual(observed_idle_flags, [False, False, after_idle, False]) + self.assertEqual(scheduler.forward_ct, 4) + self.assertEqual(observed_iters, [1, 2, 3, 4]) + expected_intervals = [ + LAUNCH_TIMESTAMPS[1] - LAUNCH_TIMESTAMPS[0], + LAUNCH_TIMESTAMPS[3] - LAUNCH_TIMESTAMPS[2], + ] + if not after_idle: + expected_intervals.append(LAUNCH_TIMESTAMPS[2] - LAUNCH_TIMESTAMPS[1]) + expected_samples = len(expected_intervals) + expected_busy_us = round(sum(expected_intervals) * 1_000_000) + if mode == ForwardMode.EXTEND: + self.assertEqual(scheduler.total_prefill_busy_us, expected_busy_us) + self.assertEqual( + scheduler.total_prefill_uncached_tokens, expected_samples * 1024 + ) + else: + self.assertEqual(scheduler.decode_moment_totals[0], expected_samples) + self.assertEqual(scheduler.decode_moment_totals[2], expected_busy_us) + + def make_scheduler(self, schedule): + scheduler = Scheduler.__new__(Scheduler) + scheduler._engine_paused = False + scheduler._sched_idled = False + scheduler._prev_step = None + scheduler.forward_ct = 0 + scheduler.processed_tokens_counter = 0 + scheduler.spec_algorithm = SpeculativeAlgorithm.NONE + scheduler.ps = SimpleNamespace(pp_rank=0, attn_tp_rank=0, attn_cp_rank=0) + scheduler._poll_timeout_aborts = Mock(return_value=[]) + scheduler.scheduler_stage_metrics = None + scheduler.metrics_reporter = SimpleNamespace(record_scheduler_active=Mock()) + scheduler.scripted_scheduler_hook = SimpleNamespace( + on_run_batch=Mock(side_effect=_BeforeModelForward) + ) + scheduler.total_prefill_busy_us = 0 + scheduler.total_prefill_uncached_tokens = 0 + scheduler.decode_moment_totals = [0.0] * 6 + scheduler.running_batch = ScheduleBatch(reqs=[]) + scheduler.last_batch = None + scheduler.chunked_req = None + scheduler.waiting_queue = [] + scheduler.enable_staging = False + scheduler.request_receiver = SimpleNamespace( + recv_requests=Mock(side_effect=[[]] * len(schedule) + [StopIteration]) + ) + scheduler.process_input_requests = Mock() + scheduler.disagg_prefill_bootstrap_queue = SimpleNamespace( + pop_bootstrapped=Mock(return_value=[]) + ) + scheduler.disagg_decode_prealloc_queue = SimpleNamespace( + prefetch_prefill_dp_rank_queries=Mock(), queue=[] + ) + scheduler.process_decode_queue = Mock() + scheduler.ngram_embedding_manager = SimpleNamespace( + prepare_for_forward=lambda batch, chunked_req: batch + ) + next_plan = Mock( + side_effect=[ + SimpleNamespace( + running_batch=scheduler.running_batch, batch_to_run=batch + ) + for batch in schedule + ] + ) + scheduler.get_next_disagg_prefill_batch_to_run = next_plan + scheduler.get_next_disagg_decode_batch_to_run = next_plan + scheduler.get_next_batch_to_run = next_plan + scheduler.get_new_batch_prefill = next_plan + scheduler.is_disable_overlap_for_batch = Mock(return_value=False) + scheduler._apply_war_barrier = Mock() + scheduler.process_disagg_prefill_inflight_queue = Mock() + scheduler.launch_batch_sample_if_needed = Mock() + scheduler.on_idle = Mock() + return scheduler + + def prepare_pp_scheduler(self, scheduler): + scheduler.ps.pp_size = 2 + scheduler.pp_group = SimpleNamespace(is_last_rank=True) + scheduler.forward_stream_ctx = nullcontext() + scheduler.forward_stream = Mock() + scheduler.schedule_stream = Mock() + scheduler.device_module = SimpleNamespace(Event=Mock, current_stream=Mock()) + scheduler._pp_recv_proxy_tensors = Mock() + scheduler._pp_commit_comm_work = Mock() + scheduler._pp_prepare_tensor_dict = Mock(return_value={}) + scheduler._pp_commit_send_output_work_and_preprocess_output_tensors = Mock( + return_value=(None, GenerationBatchResult(), Mock()) + ) + scheduler._pp_pd_get_bootstrapped_ids = Mock(return_value=None) + scheduler._pp_pd_get_prefill_transferred_ids = Mock(return_value=None) + scheduler._pp_pd_get_retract_ids = Mock(return_value=None) + scheduler._pp_pd_get_prealloc_ids = Mock(return_value=None) + scheduler._pp_pd_get_decode_transferred_ids = Mock(return_value=None) + scheduler._pp_pd_send_consensus_bootstrapped_ids = Mock(return_value=([], None)) + scheduler._pp_pd_send_consensus_release_ids = Mock(return_value=([], None)) + scheduler.process_prefill_chunk = Mock() + scheduler._process_hicache_events = Mock() + scheduler.dp_attn_adapter = SimpleNamespace( + maybe_prepare_mlp_sync_batch=lambda batch: batch + ) + scheduler.disagg_prefill_inflight_queue = [] + scheduler.disagg_decode_transfer_queue = SimpleNamespace(queue=[]) + + +if __name__ == "__main__": + unittest.main()