Refactor device timer installation and rename prefill prealloc to bootstrap (#24341)

This commit is contained in:
Lianmin Zheng
2026-05-04 13:57:13 -07:00
committed by GitHub
parent e93bb638bf
commit 29dd3a36c0
5 changed files with 29 additions and 26 deletions
+2 -2
View File
@@ -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")}
+1 -12
View File
@@ -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,
@@ -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(
@@ -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
@@ -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,