minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[observability] add ServerArgs.stat_loggers for pluggable metrics backend (#24610)

Browse Source 
Signed-off-by: Dongjun Na <kmu5544616@gmail.com>
This commit is contained in:
Dongjun Na
2026-05-24 22:41:28 +08:00
committed by GitHub
parent b6f71d5850
commit 9d50cd9742
9 changed files with 411 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/eplb/expert_distribution.py
+11 -2
View File
@@ -30,7 +30,11 @@ import torch.distributed
from sglang.srt.environ import envs
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.observability.metrics_collector import ExpertDispatchCollector
from sglang.srt.observability.metrics_collector import (
STAT_LOGGER_ROLE_EXPERT_DISPATCH,
ExpertDispatchCollector,
resolve_collector_class,
)
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import Withable, get_device, get_int_env_var
@@ -672,7 +676,12 @@ class _UtilizationRateAccumulatorMixin(_Accumulator):
self.window_sizes = [10, 100, 1000]
self._history = _DequeCollection(maxlens=self.window_sizes)
self._rank = torch.distributed.get_rank()
self._expert_dispatch_collector = ExpertDispatchCollector(
expert_dispatch_cls = resolve_collector_class(
self._server_args,
STAT_LOGGER_ROLE_EXPERT_DISPATCH,
ExpertDispatchCollector,
)
self._expert_dispatch_collector = expert_dispatch_cls(
self._expert_location_metadata.ep_size
)
self._metric_heatmap_collection_counter = 0
python/sglang/srt/managers/tokenizer_manager.py
+11 -2
View File
@@ -85,7 +85,11 @@ from sglang.srt.managers.tokenizer_manager_score_mixin import (
)
from sglang.srt.managers.utils import is_health_check_generate_req
from sglang.srt.observability.cpu_monitor import start_cpu_monitor_thread
from sglang.srt.observability.metrics_collector import TokenizerMetricsCollector
from sglang.srt.observability.metrics_collector import (
STAT_LOGGER_ROLE_TOKENIZER,
TokenizerMetricsCollector,
resolve_collector_class,
)
from sglang.srt.observability.req_time_stats import (
APIServerReqTimeStats,
convert_time_to_realtime,
@@ -488,7 +492,12 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
labels[label] = ""
if self.server_args.extra_metric_labels:
labels.update(self.server_args.extra_metric_labels)
self.metrics_collector = TokenizerMetricsCollector(
tokenizer_collector_cls = resolve_collector_class(
self.server_args,
STAT_LOGGER_ROLE_TOKENIZER,
TokenizerMetricsCollector,
)
self.metrics_collector = tokenizer_collector_cls(
server_args=self.server_args,
labels=labels,
bucket_time_to_first_token=self.server_args.bucket_time_to_first_token,
python/sglang/srt/mem_cache/base_prefix_cache.py
+11 -2
View File
@@ -17,7 +17,11 @@ import torch
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.observability.metrics_collector import RadixCacheMetricsCollector
from sglang.srt.observability.metrics_collector import (
STAT_LOGGER_ROLE_RADIX_CACHE,
RadixCacheMetricsCollector,
resolve_collector_class,
)
if TYPE_CHECKING:
from sglang.srt.managers.schedule_batch import Req
@@ -207,7 +211,12 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
labels = {"cache_type": self.__class__.__name__}
if server_args.extra_metric_labels:
labels.update(server_args.extra_metric_labels)
self.metrics_collector = RadixCacheMetricsCollector(labels=labels)
radix_cache_cls = resolve_collector_class(
server_args,
STAT_LOGGER_ROLE_RADIX_CACHE,
RadixCacheMetricsCollector,
)
self.metrics_collector = radix_cache_cls(labels=labels)
def update_eviction_metrics(self, num_evicted: int, start_time: float):
if self.metrics_collector is not None and num_evicted > 0:
python/sglang/srt/mem_cache/hi_mamba_radix_cache.py
+13 -2
View File
@@ -46,7 +46,11 @@ from sglang.srt.mem_cache.radix_cache import (
RadixKey,
)
from sglang.srt.mem_cache.utils import compute_node_hash_values, split_node_hash_value
from sglang.srt.observability.metrics_collector import StorageMetricsCollector
from sglang.srt.observability.metrics_collector import (
STAT_LOGGER_ROLE_STORAGE,
StorageMetricsCollector,
resolve_collector_class,
)
if TYPE_CHECKING:
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
@@ -1252,7 +1256,14 @@ class HiMambaRadixCache(MambaRadixCache):
}
if extra_metric_labels:
labels.update(extra_metric_labels)
storage_metrics_collector = StorageMetricsCollector(labels=labels)
from sglang.srt.server_args import get_global_server_args
storage_cls = resolve_collector_class(
get_global_server_args(),
STAT_LOGGER_ROLE_STORAGE,
StorageMetricsCollector,
)
storage_metrics_collector = storage_cls(labels=labels)
self.enable_storage = enable_storage
self.prefetch_threshold = prefetch_threshold
python/sglang/srt/mem_cache/hiradix_cache.py
+13 -2
View File
@@ -56,7 +56,11 @@ from sglang.srt.mem_cache.utils import (
compute_node_hash_values,
split_node_hash_value,
)
from sglang.srt.observability.metrics_collector import StorageMetricsCollector
from sglang.srt.observability.metrics_collector import (
STAT_LOGGER_ROLE_STORAGE,
StorageMetricsCollector,
resolve_collector_class,
)
if TYPE_CHECKING:
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
@@ -248,7 +252,14 @@ class HiRadixCache(RadixCache):
labels.update(extra_metric_labels)
existing_collector = getattr(self, "storage_metrics_collector", None)
if existing_collector is None:
self.storage_metrics_collector = StorageMetricsCollector(labels=labels)
from sglang.srt.server_args import get_global_server_args
storage_cls = resolve_collector_class(
get_global_server_args(),
STAT_LOGGER_ROLE_STORAGE,
StorageMetricsCollector,
)
self.storage_metrics_collector = storage_cls(labels=labels)
elif set(existing_collector.labels.keys()) == set(labels.keys()):
existing_collector.labels = labels
else:
python/sglang/srt/observability/metrics_collector.py
+74 -11
View File
@@ -183,6 +183,44 @@ class DPCooperationInfo:
return dataclasses.asdict(self)
# Role keys used by ServerArgs.stat_loggers to look up collector overrides.
# Embedded-use callers (e.g. Ray Serve LLM) pass {"scheduler": MyClass, ...} on
# ServerArgs and the five collector instantiation sites pick the right class.
STAT_LOGGER_ROLE_SCHEDULER = "scheduler"
STAT_LOGGER_ROLE_TOKENIZER = "tokenizer"
STAT_LOGGER_ROLE_STORAGE = "storage"
STAT_LOGGER_ROLE_RADIX_CACHE = "radix_cache"
STAT_LOGGER_ROLE_EXPERT_DISPATCH = "expert_dispatch"
def resolve_collector_class(
server_args: Optional["ServerArgs"], role: str, default_cls: type
) -> type:
"""Return the subclass registered for `role` on `server_args.stat_loggers`,
or `default_cls` if none is registered. Tolerates `server_args=None` and
`stat_loggers=None`."""
if server_args is None:
return default_cls
stat_loggers = getattr(server_args, "stat_loggers", None)
if not stat_loggers:
return default_cls
return stat_loggers.get(role, default_cls)
class _StatLoggerDIMixin:
"""Shared DI override hooks for all *MetricsCollector classes.
Subclasses (e.g. a Ray-backed wrapper) replace these class attributes with
classes that mirror the prometheus_client API but emit through a different
backend. ``None`` keeps the prometheus_client default.
"""
_counter_cls = None
_gauge_cls = None
_histogram_cls = None
_summary_cls = None
@dataclass(kw_only=True, frozen=True, slots=True)
class SchedulerMetricsCollectorContext:
enable_metrics: bool
@@ -192,7 +230,7 @@ class SchedulerMetricsCollectorContext:
collector: Optional["SchedulerMetricsCollector"]
class SchedulerMetricsCollector:
class SchedulerMetricsCollector(_StatLoggerDIMixin):
def __init__(
self,
@@ -203,7 +241,15 @@ class SchedulerMetricsCollector:
server_args: Optional["ServerArgs"] = None,
) -> None:
# We need to import prometheus_client after setting the env variable `PROMETHEUS_MULTIPROC_DIR`
from prometheus_client import Counter, Gauge, Histogram, Summary
from prometheus_client import Counter as _PromCounter
from prometheus_client import Gauge as _PromGauge
from prometheus_client import Histogram as _PromHistogram
from prometheus_client import Summary as _PromSummary
Counter = self._counter_cls or _PromCounter
Gauge = self._gauge_cls or _PromGauge
Histogram = self._histogram_cls or _PromHistogram
Summary = self._summary_cls or _PromSummary
self.labels = labels
self.enable_lora = enable_lora
@@ -989,7 +1035,10 @@ class SchedulerMetricsCollector:
labels["dp_rank"] = dp_rank
if server_args.extra_metric_labels:
labels.update(server_args.extra_metric_labels)
collector = cls(
scheduler_collector_cls = resolve_collector_class(
server_args, STAT_LOGGER_ROLE_SCHEDULER, cls
)
collector = scheduler_collector_cls(
labels=labels,
enable_lora=enable_lora,
enable_hierarchical_cache=enable_hierarchical_cache,
@@ -1318,7 +1367,7 @@ class SchedulerMetricsCollector:
)
class TokenizerMetricsCollector:
class TokenizerMetricsCollector(_StatLoggerDIMixin):
def __init__(
self,
server_args: Optional[ServerArgs] = None,
@@ -1328,7 +1377,11 @@ class TokenizerMetricsCollector:
bucket_e2e_request_latency: Optional[List[float]] = None,
) -> None:
# We need to import prometheus_client after setting the env variable `PROMETHEUS_MULTIPROC_DIR`
from prometheus_client import Counter, Histogram
from prometheus_client import Counter as _PromCounter
from prometheus_client import Histogram as _PromHistogram
Counter = self._counter_cls or _PromCounter
Histogram = self._histogram_cls or _PromHistogram
self.labels = labels or {}
@@ -1634,12 +1687,16 @@ class StorageMetrics:
backup_bandwidth: List[float] = field(default_factory=list)
class StorageMetricsCollector:
class StorageMetricsCollector(_StatLoggerDIMixin):
def __init__(
self,
labels: Dict[str, str],
):
from prometheus_client import Counter, Histogram
from prometheus_client import Counter as _PromCounter
from prometheus_client import Histogram as _PromHistogram
Counter = self._counter_cls or _PromCounter
Histogram = self._histogram_cls or _PromHistogram
self.labels = labels
@@ -1728,9 +1785,11 @@ class StorageMetricsCollector:
self._log_histogram(self.histogram_backup_bandwidth, v)
class ExpertDispatchCollector:
class ExpertDispatchCollector(_StatLoggerDIMixin):
def __init__(self, ep_size: int) -> None:
from prometheus_client import Histogram
from prometheus_client import Histogram as _PromHistogram
Histogram = self._histogram_cls or _PromHistogram
ep_size_buckets = [i for i in range(ep_size)]
self.eplb_gpu_physical_count = Histogram(
@@ -1741,13 +1800,17 @@ class ExpertDispatchCollector:
)
class RadixCacheMetricsCollector:
class RadixCacheMetricsCollector(_StatLoggerDIMixin):
def __init__(
self,
labels: Dict[str, str],
) -> None:
# We need to import prometheus_client after setting the env variable `PROMETHEUS_MULTIPROC_DIR`
from prometheus_client import Counter, Histogram
from prometheus_client import Counter as _PromCounter
from prometheus_client import Histogram as _PromHistogram
Counter = self._counter_cls or _PromCounter
Histogram = self._histogram_cls or _PromHistogram
self.labels = labels
python/sglang/srt/server_args.py
+14 -1
View File
@@ -479,6 +479,14 @@ class ServerArgs:
export_metrics_to_file: bool = False
export_metrics_to_file_dir: Optional[str] = None
# Class-level DI for the five *MetricsCollector classes. Maps collector role
# (one of: "scheduler", "tokenizer", "storage", "radix_cache", "expert_dispatch")
# to a subclass of the matching base collector. The five instantiation sites
# read from this map and fall back to the base class. Class-object only (no
# CLI surface) since this exists for embedded use cases that pass a Python
# class directly. Default None preserves existing behavior.
stat_loggers: Optional[Dict[str, type]] = None
# API related
api_key: Optional[str] = None
admin_api_key: Optional[str] = None
@@ -6905,7 +6913,12 @@ class ServerArgs:
args.dp_size = args.data_parallel_size
args.ep_size = args.expert_parallel_size
attrs = [attr.name for attr in dataclasses.fields(cls)]
# Some dataclass fields (e.g. stat_loggers) intentionally have no CLI
# surface and won't appear on the argparse Namespace. Skip them so the
# dataclass default applies.
attrs = [
attr.name for attr in dataclasses.fields(cls) if hasattr(args, attr.name)
]
return cls(**{attr: getattr(args, attr) for attr in attrs})
def url(self, port: Optional[int] = None):