Refactor device timer installation and rename prefill prealloc to bootstrap (#24341)
This commit is contained in:
@@ -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")}
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user