From 5be9c9f7c68d9ee6895af167c520e4bf211593bf Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Fri, 10 Jul 2026 08:55:06 +0800 Subject: [PATCH] Localize cur_batch field in Scheduler to avoid field-based state access (#29407) --- python/sglang/srt/disaggregation/decode.py | 6 +-- python/sglang/srt/disaggregation/prefill.py | 6 +-- .../hardware_backend/mlx/scheduler_mixin.py | 4 +- python/sglang/srt/managers/scheduler.py | 20 ++++----- .../scheduler_components/invariant_checker.py | 8 ++-- .../sglang/srt/managers/scheduler_pp_mixin.py | 41 +++++++++++-------- .../test_scripted_core_4gpu.py | 5 ++- .../test_decode_queue_cleanup.py | 2 +- .../mlx/test_attention_patching.py | 2 +- .../test_scheduler_pause_generation.py | 14 +++---- 10 files changed, 60 insertions(+), 48 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index b14d2c7df..f4d3410e0 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -1901,7 +1901,7 @@ class SchedulerDisaggregationDecodeMixin: # Get the next batch to run batch = self.get_next_disagg_decode_batch_to_run() - self.cur_batch = batch + self.cur_batch_for_debug = batch # Launch the current batch if batch: @@ -1935,7 +1935,7 @@ class SchedulerDisaggregationDecodeMixin: # Get the next batch to run batch = self.get_next_disagg_decode_batch_to_run() - self.cur_batch = batch + self.cur_batch_for_debug = batch # overlap + spec + grammar is unsupported (would desync DP ranks). disable_overlap_for_batch = self.is_disable_overlap_for_batch(batch) @@ -1958,7 +1958,7 @@ class SchedulerDisaggregationDecodeMixin: # Run sample of the current batch # It depends on the result of the last batch (e.g., grammar), so we run it after the last batch is processed. - self.launch_batch_sample_if_needed(batch_result) + self.launch_batch_sample_if_needed(batch_result, batch) # Update last_batch self.last_batch = batch diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 94dcf0930..882b2ffbb 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -502,7 +502,7 @@ class SchedulerDisaggregationPrefillMixin: # Get the next batch to run batch = self.get_next_disagg_prefill_batch_to_run() - self.cur_batch = batch + self.cur_batch_for_debug = batch # Launch the current batch if batch: @@ -536,7 +536,7 @@ class SchedulerDisaggregationPrefillMixin: # Get the next batch to run batch = self.get_next_disagg_prefill_batch_to_run() - self.cur_batch = batch + self.cur_batch_for_debug = batch # Launch the current batch if batch: @@ -559,7 +559,7 @@ class SchedulerDisaggregationPrefillMixin: # Run sample of the current batch # It depends on the result of the last batch (e.g., grammar), so we run it after the last batch is processed. - self.launch_batch_sample_if_needed(batch_result) + self.launch_batch_sample_if_needed(batch_result, batch) # Update last_batch self.last_batch = batch diff --git a/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py b/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py index 58ab0ff7f..5e640d529 100644 --- a/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py +++ b/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py @@ -234,7 +234,7 @@ class SchedulerMlxOverlapMixin: ): pending_curr = pending_next pending_next = None - self.cur_batch = pending_curr.schedule_batch + self.cur_batch_for_debug = pending_curr.schedule_batch self.last_batch = pending_curr.schedule_batch if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get(): self.invariant_checker.self_check_during_busy() @@ -247,7 +247,7 @@ class SchedulerMlxOverlapMixin: self.result_queue.popleft() pending_next = None next_batch = self.get_next_batch_to_run() - self.cur_batch = next_batch + self.cur_batch_for_debug = next_batch if next_batch: pending_curr = _launch_fresh(next_batch) self.result_queue.append(pending_curr) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 5bb43c0be..a168de47b 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -319,7 +319,7 @@ class Scheduler( self.is_initializing = True # init_soft_watchdog starts a daemon thread that reads these on its first tick. self.forward_ct: int = 0 - self.cur_batch: Optional[ScheduleBatch] = None + self.cur_batch_for_debug: Optional[ScheduleBatch] = None self.init_soft_watchdog(server_args) # Parse args @@ -965,7 +965,7 @@ class Scheduler( # The running decoding batch for continuous batching self.running_batch: ScheduleBatch = ScheduleBatch(reqs=[], batch_is_full=False) # The current forward batch - self.cur_batch: Optional[ScheduleBatch] = None + self.cur_batch_for_debug: Optional[ScheduleBatch] = None # The last forward batch self.last_batch: Optional[ScheduleBatch] = None self.forward_ct = 0 @@ -1542,7 +1542,7 @@ class Scheduler( # Get the next batch to run batch = self.get_next_batch_to_run() - self.cur_batch = batch + self.cur_batch_for_debug = batch # Launch the current batch if batch: @@ -1583,7 +1583,7 @@ class Scheduler( # Get the next batch to run batch = self.get_next_batch_to_run() - self.cur_batch = batch + self.cur_batch_for_debug = batch disable_overlap_for_batch = self.is_disable_overlap_for_batch(batch) # If we do not need to overlap the current batch with the last batch, @@ -1617,7 +1617,7 @@ class Scheduler( # Run sample of the current batch # It depends on the result of the last batch (e.g., grammar), so we run it after the last batch is processed. if self.is_generation: - self.launch_batch_sample_if_needed(batch_result) + self.launch_batch_sample_if_needed(batch_result, batch) # Update last_batch self.last_batch = batch @@ -3415,7 +3415,7 @@ class Scheduler( self.future_map.stash(future_indices, payload) def launch_batch_sample_if_needed( - self, batch_result: GenerationBatchResult + self, batch_result: GenerationBatchResult, cur_batch: ScheduleBatch ) -> Union[GenerationBatchResult]: # TODO(lsyin): make the delayed sample a default behavior after # unifying the forward_batch_generation interface (related to spec V2). @@ -3429,8 +3429,8 @@ class Scheduler( # Delay-sample is non-spec only; relays the sampled bonus tokens. self._relay_forward_payload(batch_result.future_indices, batch_result) batch_result.copy_to_cpu( - return_logprob=self.cur_batch.return_logprob, - return_hidden_states=self.cur_batch.return_hidden_states, + return_logprob=cur_batch.return_logprob, + return_hidden_states=cur_batch.return_hidden_states, ) # Release the closure and large GPU tensors that are no longer needed. @@ -3731,7 +3731,7 @@ class Scheduler( def flush_cache(self, empty_cache: bool = True): """Flush memory pools (e.g., KV cache, Mamba cache) and optionally empty device allocator cache.""" if self.is_fully_idle(): - self.cur_batch = None + self.cur_batch_for_debug = None self.last_batch = None self.tree_cache.reset() self.req_to_token_pool.clear() @@ -4021,7 +4021,7 @@ class Scheduler( self.running_batch.merge_batch(self.last_batch) self.last_batch = None - self.cur_batch = None + self.cur_batch_for_debug = None if recv_req.mode == "retract" and not self.running_batch.is_empty(): self.running_batch.filter_batch() diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index 633baf8ca..bdd1b7019 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/python/sglang/srt/managers/scheduler_components/invariant_checker.py @@ -470,14 +470,16 @@ def create_scheduler_watchdog( scheduler.pool_stats_observer.get_pool_stats(), ) return ( - f"{scheduler.cur_batch.batch_size()=}\n" - f"{scheduler.cur_batch.reqs=}\n" + "\n".join(messages) + f"{scheduler.cur_batch_for_debug.batch_size()=}\n" + f"{scheduler.cur_batch_for_debug.reqs=}\n" + "\n".join(messages) ) return WatchdogRaw( debug_name="Scheduler", get_counter=lambda: scheduler.forward_ct, - is_active=lambda: scheduler.is_initializing or scheduler.cur_batch is not None, + is_active=lambda: ( + scheduler.is_initializing or scheduler.cur_batch_for_debug is not None + ), watchdog_timeout=watchdog_timeout, soft=soft, dump_info=dump_info, diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index 3292c7a09..18dbf1aa5 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -110,8 +110,9 @@ class SchedulerPPMixin: with torch.profiler.record_function("get_next_batch_to_run"): self.mbs[mb_id] = self.get_next_batch_to_run() self.running_mbs[mb_id] = self.running_batch - self.cur_batch: Optional[ScheduleBatch] = self.mbs[mb_id] - if self.cur_batch: + cur_batch: Optional[ScheduleBatch] = self.mbs[mb_id] + self.cur_batch_for_debug = cur_batch + if cur_batch: server_is_idle = False pp_proxy_tensors = self._pp_recv_proxy_tensors() next_pp_outputs = None @@ -125,9 +126,10 @@ class SchedulerPPMixin: ) ) self._pp_commit_comm_work(self.send_proxy_work) - if self.cur_batch: + if cur_batch: result, self.launch_event = self._pp_launch_batch( mb_id, + cur_batch, pp_proxy_tensors, self.mb_metadata, self.last_rank_comm_queue, @@ -148,7 +150,7 @@ class SchedulerPPMixin: ) self.last_mbs[next_mb_id] = self.mbs[next_mb_id] if not self.pp_group.is_last_rank: - if self.cur_batch: + if cur_batch: self.device_module.current_stream().wait_event( self.launch_event ) @@ -252,8 +254,9 @@ class SchedulerPPMixin: self.mbs[mb_id] = batch self.running_mbs[mb_id] = self.running_batch - self.cur_batch: Optional[ScheduleBatch] = self.mbs[mb_id] - if self.cur_batch: + cur_batch: Optional[ScheduleBatch] = self.mbs[mb_id] + self.cur_batch_for_debug = cur_batch + if cur_batch: server_is_idle = False pp_proxy_tensors = self._pp_recv_proxy_tensors() @@ -265,9 +268,10 @@ class SchedulerPPMixin: ) ) self._pp_commit_comm_work(self.send_proxy_work) - if self.cur_batch: + if cur_batch: result, self.launch_event = self._pp_launch_batch( mb_id, + cur_batch, pp_proxy_tensors, self.mb_metadata, self.last_rank_comm_queue, @@ -325,7 +329,7 @@ class SchedulerPPMixin: send_transfer_work = self._pp_send_pyobj_to_next_stage( transferred_rids, async_send=True ) - if self.cur_batch: + if cur_batch: self.device_module.current_stream().wait_event( self.launch_event ) @@ -402,11 +406,12 @@ class SchedulerPPMixin: self.mbs[mb_id] = batch self.running_mbs[mb_id] = self.running_batch - self.cur_batch: Optional[ScheduleBatch] = self.mbs[mb_id] - if self.cur_batch: + cur_batch: Optional[ScheduleBatch] = self.mbs[mb_id] + self.cur_batch_for_debug = cur_batch + if cur_batch: server_is_idle = False pp_proxy_tensors = None - if not self.cur_batch.forward_mode.is_prebuilt(): + if not cur_batch.forward_mode.is_prebuilt(): pp_proxy_tensors = self._pp_recv_proxy_tensors() # early send output if possible @@ -419,9 +424,10 @@ class SchedulerPPMixin: ) self._pp_commit_comm_work(self.send_proxy_work) - if self.cur_batch: + if cur_batch: result, self.launch_event = self._pp_launch_batch( mb_id, + cur_batch, pp_proxy_tensors, self.mb_metadata, self.last_rank_comm_queue, @@ -508,7 +514,7 @@ class SchedulerPPMixin: send_transfer_work = self._pp_send_pyobj_to_next_stage( transferred_rids, async_send=True ) - if self.cur_batch and not self.cur_batch.forward_mode.is_prebuilt(): + if cur_batch and not cur_batch.forward_mode.is_prebuilt(): self.device_module.current_stream().wait_event( self.launch_event ) @@ -1240,6 +1246,7 @@ class SchedulerPPMixin: def _pp_launch_batch( self: Scheduler, mb_id: int, + cur_batch: ScheduleBatch, pp_proxy_tensors: PPProxyTensors, mb_metadata: List[Optional[PPBatchMetadata]], last_rank_comm_queue: deque, @@ -1248,13 +1255,13 @@ class SchedulerPPMixin: with self.forward_stream_ctx: self.forward_stream.wait_stream(self.schedule_stream) set_time_batch( - self.cur_batch.reqs, + cur_batch.reqs, "set_run_batch_cpu_start_time", trace_only=True, ) - result = self.run_batch(self.cur_batch, pp_proxy_tensors) + result = self.run_batch(cur_batch, pp_proxy_tensors) set_time_batch( - self.cur_batch.reqs, + cur_batch.reqs, "set_run_batch_cpu_end_time", trace_only=True, attrs={"pp_mb_id": mb_id}, @@ -1270,7 +1277,7 @@ class SchedulerPPMixin: ( event, PPProxyTensors( - self._pp_prepare_tensor_dict(result, self.cur_batch) + self._pp_prepare_tensor_dict(result, cur_batch) ), ) ) diff --git a/test/registered/chunked_prefill/test_scripted_core_4gpu.py b/test/registered/chunked_prefill/test_scripted_core_4gpu.py index c1e14f0c8..559962cd9 100644 --- a/test/registered/chunked_prefill/test_scripted_core_4gpu.py +++ b/test/registered/chunked_prefill/test_scripted_core_4gpu.py @@ -77,7 +77,10 @@ class TestScriptedPpChunkSweep(ScriptedTestCase): scheduler.chunked_req is None and len(scheduler.waiting_queue) == 0 and all(x.is_empty() for x in scheduler.running_mbs) - and (scheduler.cur_batch is None or scheduler.cur_batch.is_empty()) + and ( + scheduler.cur_batch_for_debug is None + or scheduler.cur_batch_for_debug.is_empty() + ) and (scheduler.last_batch is None or scheduler.last_batch.is_empty()) ) if in_flight and queues_clear: diff --git a/test/registered/unit/disaggregation/test_decode_queue_cleanup.py b/test/registered/unit/disaggregation/test_decode_queue_cleanup.py index 6cb9f53da..afeb1c875 100644 --- a/test/registered/unit/disaggregation/test_decode_queue_cleanup.py +++ b/test/registered/unit/disaggregation/test_decode_queue_cleanup.py @@ -193,7 +193,7 @@ class TestDecodeQueueCleanup(CustomTestCase): scheduler.dllm_manager = MagicMock() scheduler.dllm_manager.any_staging_reqs.return_value = False scheduler.last_batch = None - scheduler.cur_batch = None + scheduler.cur_batch_for_debug = None scheduler.enable_overlap = False scheduler.ps = SimpleNamespace(pp_size=1) scheduler.running_mbs = [] 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 f31df9571..167a24886 100644 --- a/test/registered/unit/hardware_backend/mlx/test_attention_patching.py +++ b/test/registered/unit/hardware_backend/mlx/test_attention_patching.py @@ -1135,7 +1135,7 @@ class TestMlxOverlapScheduler(unittest.TestCase): scheduler.waiting_queue = [] scheduler.result_queue = deque() scheduler.future_map = SimpleNamespace() - scheduler.cur_batch = None + scheduler.cur_batch_for_debug = None scheduler.last_batch = None scheduler.tp_worker = SimpleNamespace( async_forward_batch_generation_mlx=fake_forward diff --git a/test/registered/unit/managers/test_scheduler_pause_generation.py b/test/registered/unit/managers/test_scheduler_pause_generation.py index 33c24f443..596714afd 100644 --- a/test/registered/unit/managers/test_scheduler_pause_generation.py +++ b/test/registered/unit/managers/test_scheduler_pause_generation.py @@ -26,7 +26,7 @@ class TestSchedulerPauseGeneration(unittest.TestCase): scheduler._engine_paused = False scheduler.enable_overlap = False scheduler.last_batch = None - scheduler.cur_batch = None + scheduler.cur_batch_for_debug = None scheduler.chunked_req = None scheduler.running_batch = MagicMock() scheduler.running_batch.reqs = [] @@ -59,11 +59,11 @@ class TestSchedulerPauseGeneration(unittest.TestCase): """in_place pause should only set _engine_paused and return.""" scheduler = self._new_scheduler() scheduler.last_batch = MagicMock() - scheduler.cur_batch = MagicMock() + scheduler.cur_batch_for_debug = MagicMock() scheduler.chunked_req = MagicMock() original_last_batch = scheduler.last_batch - original_cur_batch = scheduler.cur_batch + original_cur_batch = scheduler.cur_batch_for_debug original_chunked_req = scheduler.chunked_req scheduler.pause_generation(PauseGenerationReqInput(mode="in_place")) @@ -71,7 +71,7 @@ class TestSchedulerPauseGeneration(unittest.TestCase): self.assertTrue(scheduler._engine_paused) # All state must be preserved — no mutation self.assertIs(scheduler.last_batch, original_last_batch) - self.assertIs(scheduler.cur_batch, original_cur_batch) + self.assertIs(scheduler.cur_batch_for_debug, original_cur_batch) self.assertIs(scheduler.chunked_req, original_chunked_req) def test_inplace_does_not_drain_overlap_queue(self): @@ -99,17 +99,17 @@ class TestSchedulerPauseGeneration(unittest.TestCase): scheduler.running_batch.merge_batch.assert_not_called() def test_abort_clears_state(self): - """abort mode should clear last_batch and cur_batch.""" + """abort mode should clear last_batch and cur_batch_for_debug.""" scheduler = self._new_scheduler() scheduler.last_batch = MagicMock() scheduler.last_batch.forward_mode.is_extend.return_value = False - scheduler.cur_batch = MagicMock() + scheduler.cur_batch_for_debug = MagicMock() scheduler.pause_generation(PauseGenerationReqInput(mode="abort")) self.assertTrue(scheduler._engine_paused) self.assertIsNone(scheduler.last_batch) - self.assertIsNone(scheduler.cur_batch) + self.assertIsNone(scheduler.cur_batch_for_debug) def test_retract_clears_running_batch(self): """retract mode should retract all requests from running_batch."""