[srt] Add sglang:weight_load_duration_seconds gauge with source label (#27363)

This commit is contained in:
Leon Gao
2026-06-08 23:41:54 +08:00
committed by GitHub
parent 40030d8af8
commit eb646c7b78
3 changed files with 83 additions and 37 deletions
+1
View File
@@ -1578,6 +1578,7 @@ class Scheduler(
memory_saver_adapter=self.memory_saver_adapter, memory_saver_adapter=self.memory_saver_adapter,
flush_cache=self.flush_cache, flush_cache=self.flush_cache,
is_fully_idle=self.is_fully_idle, is_fully_idle=self.is_fully_idle,
metrics_collector=self.metrics_collector,
) )
def init_lora_drainer(self) -> None: def init_lora_drainer(self) -> None:
@@ -2,9 +2,11 @@ from __future__ import annotations
import hashlib import hashlib
import logging import logging
import time
import traceback import traceback
from contextlib import contextmanager
from dataclasses import dataclass, field from dataclasses import dataclass, field
from typing import Any, Callable, Dict, Tuple from typing import Any, Callable, Dict, Iterator, Optional, Tuple
import torch import torch
@@ -75,9 +77,25 @@ class SchedulerWeightUpdaterManager:
memory_saver_adapter: Any memory_saver_adapter: Any
flush_cache: Callable[..., bool] flush_cache: Callable[..., bool]
is_fully_idle: Callable[..., bool] is_fully_idle: Callable[..., bool]
metrics_collector: Optional[Any] = None
offload_tags: set = field(default_factory=set) offload_tags: set = field(default_factory=set)
stashed_model_static_state: Any = None stashed_model_static_state: Any = None
@contextmanager
def _observe_weight_load(self, source: str) -> Iterator[None]:
# Edge-trigger weight_load_duration_seconds at the end of each
# update_weights_from_* call. Engine is paused during the update so
# the periodic log_stats path can't carry this.
# `source` distinguishes disk vs distributed vs tensor vs ipc.
t0 = time.perf_counter()
try:
yield
finally:
if self.metrics_collector is not None:
self.metrics_collector.observe_weight_load(
time.perf_counter() - t0, source
)
def flush_cache_after_weight_update(self, recv_req) -> None: def flush_cache_after_weight_update(self, recv_req) -> None:
if recv_req.flush_cache: if recv_req.flush_cache:
flush_cache_success = self.flush_cache( flush_cache_success = self.flush_cache(
@@ -87,15 +105,16 @@ class SchedulerWeightUpdaterManager:
def update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput): def update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput):
"""In-place update of the weights from disk.""" """In-place update of the weights from disk."""
success, message = self.tp_worker.update_weights_from_disk(recv_req) with self._observe_weight_load("disk"):
tp_success = success success, message = self.tp_worker.update_weights_from_disk(recv_req)
if success and self.draft_worker is not None: tp_success = success
success, message = self.draft_worker.update_weights_from_disk(recv_req) if success and self.draft_worker is not None:
if tp_success: success, message = self.draft_worker.update_weights_from_disk(recv_req)
self.flush_cache_after_weight_update(recv_req) if tp_success:
if not success: self.flush_cache_after_weight_update(recv_req)
logger.error(message) if not success:
return UpdateWeightFromDiskReqOutput(success, message, 0) logger.error(message)
return UpdateWeightFromDiskReqOutput(success, message, 0)
def init_weights_update_group(self, recv_req: InitWeightsUpdateGroupReqInput): def init_weights_update_group(self, recv_req: InitWeightsUpdateGroupReqInput):
"""Initialize the online model parameter update group.""" """Initialize the online model parameter update group."""
@@ -115,39 +134,42 @@ class SchedulerWeightUpdaterManager:
recv_req: UpdateWeightsFromDistributedReqInput, recv_req: UpdateWeightsFromDistributedReqInput,
) -> Tuple[bool, str]: ) -> Tuple[bool, str]:
"""Update the online model parameter.""" """Update the online model parameter."""
success, message = self.tp_worker.update_weights_from_distributed(recv_req) with self._observe_weight_load("distributed"):
if success: success, message = self.tp_worker.update_weights_from_distributed(recv_req)
self.flush_cache_after_weight_update(recv_req) if success:
else: self.flush_cache_after_weight_update(recv_req)
logger.error(message) else:
return UpdateWeightsFromDistributedReqOutput(success, message) logger.error(message)
return UpdateWeightsFromDistributedReqOutput(success, message)
def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput): def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput):
"""Update the online model parameter from tensors.""" """Update the online model parameter from tensors."""
if recv_req.disable_draft_model: with self._observe_weight_load("tensor"):
worker = self.tp_worker if recv_req.disable_draft_model:
else: worker = self.tp_worker
worker = self.draft_worker or self.tp_worker else:
success, message = worker.update_weights_from_tensor(recv_req) worker = self.draft_worker or self.tp_worker
if success: success, message = worker.update_weights_from_tensor(recv_req)
self.flush_cache_after_weight_update(recv_req) if success:
else: self.flush_cache_after_weight_update(recv_req)
logger.error(message) else:
torch.distributed.barrier(group=self.tp_cpu_group) logger.error(message)
return UpdateWeightsFromTensorReqOutput(success, message) torch.distributed.barrier(group=self.tp_cpu_group)
return UpdateWeightsFromTensorReqOutput(success, message)
def update_weights_from_ipc(self, recv_req: UpdateWeightsFromIPCReqInput): def update_weights_from_ipc(self, recv_req: UpdateWeightsFromIPCReqInput):
"""Update the online model parameter from IPC for checkpoint-engine integration.""" """Update the online model parameter from IPC for checkpoint-engine integration."""
success, message = self.tp_worker.update_weights_from_ipc(recv_req) with self._observe_weight_load("ipc"):
tp_success = success success, message = self.tp_worker.update_weights_from_ipc(recv_req)
if success and self.draft_worker is not None: tp_success = success
success, message = self.draft_worker.update_weights_from_ipc(recv_req) if success and self.draft_worker is not None:
if tp_success: success, message = self.draft_worker.update_weights_from_ipc(recv_req)
self.flush_cache_after_weight_update(recv_req) if tp_success:
if not success: self.flush_cache_after_weight_update(recv_req)
logger.error(message) if not success:
torch.distributed.barrier(group=self.tp_cpu_group) logger.error(message)
return UpdateWeightsFromIPCReqOutput(success, message) torch.distributed.barrier(group=self.tp_cpu_group)
return UpdateWeightsFromIPCReqOutput(success, message)
def get_weights_by_name(self, recv_req: GetWeightsByNameReqInput): def get_weights_by_name(self, recv_req: GetWeightsByNameReqInput):
parameter = self.tp_worker.get_weights_by_name(recv_req) parameter = self.tp_worker.get_weights_by_name(recv_req)
@@ -393,6 +393,21 @@ class SchedulerMetricsCollector(_StatLoggerDIMixin):
multiprocess_mode="mostrecent", multiprocess_mode="mostrecent",
) )
# =================================================================
# Weight update
# =================================================================
self.weight_load_duration_seconds = Gauge(
name="sglang:weight_load_duration_seconds",
documentation=(
"Wall time of the most recent update_weights_from_<source> call on "
"this scheduler rank (seconds). `source` label is one of: disk, "
"distributed, tensor, ipc. Event-detection via "
"changes(...[<range>]) > 0 — no separate counter needed."
),
labelnames=[*labels.keys(), "source"],
multiprocess_mode="mostrecent",
)
# ================================================================= # =================================================================
# Speculative decoding # Speculative decoding
# ================================================================= # =================================================================
@@ -1124,6 +1139,14 @@ class SchedulerMetricsCollector(_StatLoggerDIMixin):
def observe_queue_time(self, latency: float) -> None: def observe_queue_time(self, latency: float) -> None:
self._log_histogram(self.queue_time, latency) self._log_histogram(self.queue_time, latency)
def observe_weight_load(self, duration_seconds: float, source: str) -> None:
# Edge-triggered: engine is paused during the update, so log_stats
# won't fire — write the gauge inline at end of update_weights_from_*.
# `source` is "disk" | "distributed" | "tensor" | "ipc".
self.weight_load_duration_seconds.labels(**self.labels, source=source).set(
duration_seconds
)
def observe_prefill_delayer_outcome( def observe_prefill_delayer_outcome(
self, self,
forward_passes: int, forward_passes: int,