From 29dd3a36c0250c7e6206e6a10c78ba09e09b0561 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Mon, 4 May 2026 13:57:13 -0700 Subject: [PATCH] Refactor device timer installation and rename prefill prealloc to bootstrap (#24341) --- python/sglang/srt/managers/io_struct.py | 4 +-- python/sglang/srt/managers/scheduler.py | 13 +--------- .../scheduler_runtime_checker_mixin.py | 2 +- .../srt/observability/metrics_collector.py | 11 ++++---- .../observability/scheduler_metrics_mixin.py | 25 ++++++++++++++----- 5 files changed, 29 insertions(+), 26 deletions(-) diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 64506f461..3a1838792 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -1947,8 +1947,8 @@ class DisaggregationMetrics: """PD disaggregation metrics.""" mode: str # "prefill", "decode", or "null" - not a metric - prefill_prealloc_queue_reqs: int = field( - default=0, metadata={"metric": ("gauge", "Prefill prealloc queue requests")} + prefill_bootstrap_queue_reqs: int = field( + default=0, metadata={"metric": ("gauge", "Prefill bootstrap queue requests")} ) prefill_inflight_queue_reqs: int = field( default=0, metadata={"metric": ("gauge", "Prefill inflight queue requests")} diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index de2226f66..56a752b75 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -431,6 +431,7 @@ class Scheduler( # Launch a model worker and draft model worker if using speculative decoding self.init_model_worker() + self.install_device_timer_on_runners() if (t := envs.SGLANG_TEST_STUCK_SCHEDULER_INIT.get()) > 0: time.sleep(t) @@ -709,18 +710,6 @@ class Scheduler( else: self.model_worker = self.draft_worker - # Install device timer on model runners for fwd occupancy tracking - if hasattr(self, "forward_pass_device_timer"): - timer = self.forward_pass_device_timer - self.tp_worker.model_runner.device_timer = timer - if self.draft_worker is not None: - dw = getattr(self.draft_worker, "draft_worker", None) - if dw is not None: - if hasattr(dw, "draft_runner"): - dw.draft_runner.device_timer = timer - for r in getattr(dw, "draft_runner_list", []): - r.device_timer = timer - # Get token and memory info from the model worker ( self.max_total_num_tokens, diff --git a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py index 47f9a4b31..9c05b6080 100644 --- a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py +++ b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py @@ -512,7 +512,7 @@ class SchedulerRuntimeCheckerMixin: ) self.stats.num_grammar_queue_reqs = len(self.grammar_manager) if self.disaggregation_mode == DisaggregationMode.PREFILL: - self.stats.num_prefill_prealloc_queue_reqs = QueueCount.from_reqs( + self.stats.num_prefill_bootstrap_queue_reqs = QueueCount.from_reqs( self.disagg_prefill_bootstrap_queue.queue, priority_enabled ) self.stats.num_prefill_inflight_queue_reqs = QueueCount.from_reqs( diff --git a/python/sglang/srt/observability/metrics_collector.py b/python/sglang/srt/observability/metrics_collector.py index 3e149c0f0..5ba61b37b 100644 --- a/python/sglang/srt/observability/metrics_collector.py +++ b/python/sglang/srt/observability/metrics_collector.py @@ -107,7 +107,7 @@ class SchedulerStats: num_paused_reqs: int = 0 # PD disaggregation - num_prefill_prealloc_queue_reqs: QueueCount = field(default_factory=QueueCount) + num_prefill_bootstrap_queue_reqs: QueueCount = field(default_factory=QueueCount) num_prefill_inflight_queue_reqs: QueueCount = field(default_factory=QueueCount) num_decode_prealloc_queue_reqs: QueueCount = field(default_factory=QueueCount) num_decode_transfer_queue_reqs: QueueCount = field(default_factory=QueueCount) @@ -340,9 +340,9 @@ class SchedulerMetricsCollector: # ================================================================= # PD disaggregation # ================================================================= - self.num_prefill_prealloc_queue_reqs = Gauge( - name="sglang:num_prefill_prealloc_queue_reqs", - documentation="The number of requests in the prefill prealloc queue.", + self.num_prefill_bootstrap_queue_reqs = Gauge( + name="sglang:num_prefill_bootstrap_queue_reqs", + documentation="The number of requests in the prefill bootstrap queue.", labelnames=labels.keys(), multiprocess_mode="mostrecent", ) @@ -1080,7 +1080,8 @@ class SchedulerMetricsCollector: # PD disaggregation self._log_gauge_queue_count( - self.num_prefill_prealloc_queue_reqs, stats.num_prefill_prealloc_queue_reqs + self.num_prefill_bootstrap_queue_reqs, + stats.num_prefill_bootstrap_queue_reqs, ) self._log_gauge_queue_count( self.num_prefill_inflight_queue_reqs, stats.num_prefill_inflight_queue_reqs diff --git a/python/sglang/srt/observability/scheduler_metrics_mixin.py b/python/sglang/srt/observability/scheduler_metrics_mixin.py index 78925a440..07cac0545 100644 --- a/python/sglang/srt/observability/scheduler_metrics_mixin.py +++ b/python/sglang/srt/observability/scheduler_metrics_mixin.py @@ -178,6 +178,19 @@ class SchedulerMetricsMixin: enable_metrics=self.enable_metrics ) + def install_device_timer_on_runners(self: Scheduler): + if not hasattr(self, "forward_pass_device_timer"): + return + timer = self.forward_pass_device_timer + self.tp_worker.model_runner.device_timer = timer + if self.draft_worker is not None: + dw = getattr(self.draft_worker, "draft_worker", None) + if dw is not None: + if hasattr(dw, "draft_runner"): + dw.draft_runner.device_timer = timer + for r in getattr(dw, "draft_runner_list", []): + r.device_timer = timer + def init_kv_events(self: Scheduler, kv_events_config: Optional[str]): self.enable_kv_cache_events = bool( kv_events_config and self.attn_tp_rank == 0 and self.attn_cp_rank == 0 @@ -373,7 +386,7 @@ class SchedulerMetricsMixin: ) if self.disaggregation_mode == DisaggregationMode.PREFILL: - msg += f"#prealloc-req: {len(self.disagg_prefill_bootstrap_queue.queue)}, " + msg += f"#bootstrap-req: {len(self.disagg_prefill_bootstrap_queue.queue)}, " msg += f"#inflight-req: {len(self.disagg_prefill_inflight_queue)}, " if ( @@ -438,7 +451,7 @@ class SchedulerMetricsMixin: # PD disaggregation if self.disaggregation_mode == DisaggregationMode.PREFILL: - self.stats.num_prefill_prealloc_queue_reqs = QueueCount.from_reqs( + self.stats.num_prefill_bootstrap_queue_reqs = QueueCount.from_reqs( self.disagg_prefill_bootstrap_queue.queue, priority_enabled ) self.stats.num_prefill_inflight_queue_reqs = QueueCount.from_reqs( @@ -615,7 +628,7 @@ class SchedulerMetricsMixin: # PD disaggregation if self.disaggregation_mode == DisaggregationMode.PREFILL: - self.stats.num_prefill_prealloc_queue_reqs = QueueCount.from_reqs( + self.stats.num_prefill_bootstrap_queue_reqs = QueueCount.from_reqs( self.disagg_prefill_bootstrap_queue.queue, priority_enabled ) self.stats.num_prefill_inflight_queue_reqs = QueueCount.from_reqs( @@ -861,7 +874,7 @@ class SchedulerMetricsMixin: disaggregation = None if include_all or "disagg" in include: mode_str = "null" - prefill_prealloc = 0 + prefill_bootstrap = 0 prefill_inflight = 0 decode_prealloc = 0 decode_transfer = 0 @@ -869,7 +882,7 @@ class SchedulerMetricsMixin: if self.disaggregation_mode == DisaggregationMode.PREFILL: mode_str = "prefill" - prefill_prealloc = len(self.disagg_prefill_bootstrap_queue.queue) + prefill_bootstrap = len(self.disagg_prefill_bootstrap_queue.queue) prefill_inflight = len(self.disagg_prefill_inflight_queue) elif self.disaggregation_mode == DisaggregationMode.DECODE: mode_str = "decode" @@ -881,7 +894,7 @@ class SchedulerMetricsMixin: disaggregation = DisaggregationMetrics( mode=mode_str, - prefill_prealloc_queue_reqs=prefill_prealloc, + prefill_bootstrap_queue_reqs=prefill_bootstrap, prefill_inflight_queue_reqs=prefill_inflight, decode_prealloc_queue_reqs=decode_prealloc, decode_transfer_queue_reqs=decode_transfer,