[observability] add ServerArgs.stat_loggers for pluggable metrics backend (#24610)
Signed-off-by: Dongjun Na <kmu5544616@gmail.com>
This commit is contained in:
@@ -30,7 +30,11 @@ import torch.distributed
|
|||||||
|
|
||||||
from sglang.srt.environ import envs
|
from sglang.srt.environ import envs
|
||||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
|
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.server_args import ServerArgs
|
||||||
from sglang.srt.utils import Withable, get_device, get_int_env_var
|
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.window_sizes = [10, 100, 1000]
|
||||||
self._history = _DequeCollection(maxlens=self.window_sizes)
|
self._history = _DequeCollection(maxlens=self.window_sizes)
|
||||||
self._rank = torch.distributed.get_rank()
|
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._expert_location_metadata.ep_size
|
||||||
)
|
)
|
||||||
self._metric_heatmap_collection_counter = 0
|
self._metric_heatmap_collection_counter = 0
|
||||||
|
|||||||
@@ -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.managers.utils import is_health_check_generate_req
|
||||||
from sglang.srt.observability.cpu_monitor import start_cpu_monitor_thread
|
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 (
|
from sglang.srt.observability.req_time_stats import (
|
||||||
APIServerReqTimeStats,
|
APIServerReqTimeStats,
|
||||||
convert_time_to_realtime,
|
convert_time_to_realtime,
|
||||||
@@ -488,7 +492,12 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
|
|||||||
labels[label] = ""
|
labels[label] = ""
|
||||||
if self.server_args.extra_metric_labels:
|
if self.server_args.extra_metric_labels:
|
||||||
labels.update(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,
|
server_args=self.server_args,
|
||||||
labels=labels,
|
labels=labels,
|
||||||
bucket_time_to_first_token=self.server_args.bucket_time_to_first_token,
|
bucket_time_to_first_token=self.server_args.bucket_time_to_first_token,
|
||||||
|
|||||||
@@ -17,7 +17,11 @@ import torch
|
|||||||
|
|
||||||
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
|
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
|
||||||
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
|
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:
|
if TYPE_CHECKING:
|
||||||
from sglang.srt.managers.schedule_batch import Req
|
from sglang.srt.managers.schedule_batch import Req
|
||||||
@@ -207,7 +211,12 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
|
|||||||
labels = {"cache_type": self.__class__.__name__}
|
labels = {"cache_type": self.__class__.__name__}
|
||||||
if server_args.extra_metric_labels:
|
if server_args.extra_metric_labels:
|
||||||
labels.update(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):
|
def update_eviction_metrics(self, num_evicted: int, start_time: float):
|
||||||
if self.metrics_collector is not None and num_evicted > 0:
|
if self.metrics_collector is not None and num_evicted > 0:
|
||||||
|
|||||||
@@ -46,7 +46,11 @@ from sglang.srt.mem_cache.radix_cache import (
|
|||||||
RadixKey,
|
RadixKey,
|
||||||
)
|
)
|
||||||
from sglang.srt.mem_cache.utils import compute_node_hash_values, split_node_hash_value
|
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:
|
if TYPE_CHECKING:
|
||||||
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
|
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
|
||||||
@@ -1252,7 +1256,14 @@ class HiMambaRadixCache(MambaRadixCache):
|
|||||||
}
|
}
|
||||||
if extra_metric_labels:
|
if extra_metric_labels:
|
||||||
labels.update(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.enable_storage = enable_storage
|
||||||
self.prefetch_threshold = prefetch_threshold
|
self.prefetch_threshold = prefetch_threshold
|
||||||
|
|||||||
@@ -56,7 +56,11 @@ from sglang.srt.mem_cache.utils import (
|
|||||||
compute_node_hash_values,
|
compute_node_hash_values,
|
||||||
split_node_hash_value,
|
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:
|
if TYPE_CHECKING:
|
||||||
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
|
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
|
||||||
@@ -248,7 +252,14 @@ class HiRadixCache(RadixCache):
|
|||||||
labels.update(extra_metric_labels)
|
labels.update(extra_metric_labels)
|
||||||
existing_collector = getattr(self, "storage_metrics_collector", None)
|
existing_collector = getattr(self, "storage_metrics_collector", None)
|
||||||
if existing_collector is 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()):
|
elif set(existing_collector.labels.keys()) == set(labels.keys()):
|
||||||
existing_collector.labels = labels
|
existing_collector.labels = labels
|
||||||
else:
|
else:
|
||||||
|
|||||||
@@ -183,6 +183,44 @@ class DPCooperationInfo:
|
|||||||
return dataclasses.asdict(self)
|
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)
|
@dataclass(kw_only=True, frozen=True, slots=True)
|
||||||
class SchedulerMetricsCollectorContext:
|
class SchedulerMetricsCollectorContext:
|
||||||
enable_metrics: bool
|
enable_metrics: bool
|
||||||
@@ -192,7 +230,7 @@ class SchedulerMetricsCollectorContext:
|
|||||||
collector: Optional["SchedulerMetricsCollector"]
|
collector: Optional["SchedulerMetricsCollector"]
|
||||||
|
|
||||||
|
|
||||||
class SchedulerMetricsCollector:
|
class SchedulerMetricsCollector(_StatLoggerDIMixin):
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
@@ -203,7 +241,15 @@ class SchedulerMetricsCollector:
|
|||||||
server_args: Optional["ServerArgs"] = None,
|
server_args: Optional["ServerArgs"] = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
# We need to import prometheus_client after setting the env variable `PROMETHEUS_MULTIPROC_DIR`
|
# 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.labels = labels
|
||||||
self.enable_lora = enable_lora
|
self.enable_lora = enable_lora
|
||||||
@@ -989,7 +1035,10 @@ class SchedulerMetricsCollector:
|
|||||||
labels["dp_rank"] = dp_rank
|
labels["dp_rank"] = dp_rank
|
||||||
if server_args.extra_metric_labels:
|
if server_args.extra_metric_labels:
|
||||||
labels.update(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,
|
labels=labels,
|
||||||
enable_lora=enable_lora,
|
enable_lora=enable_lora,
|
||||||
enable_hierarchical_cache=enable_hierarchical_cache,
|
enable_hierarchical_cache=enable_hierarchical_cache,
|
||||||
@@ -1318,7 +1367,7 @@ class SchedulerMetricsCollector:
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
class TokenizerMetricsCollector:
|
class TokenizerMetricsCollector(_StatLoggerDIMixin):
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
server_args: Optional[ServerArgs] = None,
|
server_args: Optional[ServerArgs] = None,
|
||||||
@@ -1328,7 +1377,11 @@ class TokenizerMetricsCollector:
|
|||||||
bucket_e2e_request_latency: Optional[List[float]] = None,
|
bucket_e2e_request_latency: Optional[List[float]] = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
# We need to import prometheus_client after setting the env variable `PROMETHEUS_MULTIPROC_DIR`
|
# 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 {}
|
self.labels = labels or {}
|
||||||
|
|
||||||
@@ -1634,12 +1687,16 @@ class StorageMetrics:
|
|||||||
backup_bandwidth: List[float] = field(default_factory=list)
|
backup_bandwidth: List[float] = field(default_factory=list)
|
||||||
|
|
||||||
|
|
||||||
class StorageMetricsCollector:
|
class StorageMetricsCollector(_StatLoggerDIMixin):
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
labels: Dict[str, str],
|
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
|
self.labels = labels
|
||||||
|
|
||||||
@@ -1728,9 +1785,11 @@ class StorageMetricsCollector:
|
|||||||
self._log_histogram(self.histogram_backup_bandwidth, v)
|
self._log_histogram(self.histogram_backup_bandwidth, v)
|
||||||
|
|
||||||
|
|
||||||
class ExpertDispatchCollector:
|
class ExpertDispatchCollector(_StatLoggerDIMixin):
|
||||||
def __init__(self, ep_size: int) -> None:
|
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)]
|
ep_size_buckets = [i for i in range(ep_size)]
|
||||||
self.eplb_gpu_physical_count = Histogram(
|
self.eplb_gpu_physical_count = Histogram(
|
||||||
@@ -1741,13 +1800,17 @@ class ExpertDispatchCollector:
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
class RadixCacheMetricsCollector:
|
class RadixCacheMetricsCollector(_StatLoggerDIMixin):
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
labels: Dict[str, str],
|
labels: Dict[str, str],
|
||||||
) -> None:
|
) -> None:
|
||||||
# We need to import prometheus_client after setting the env variable `PROMETHEUS_MULTIPROC_DIR`
|
# 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
|
self.labels = labels
|
||||||
|
|
||||||
|
|||||||
@@ -479,6 +479,14 @@ class ServerArgs:
|
|||||||
export_metrics_to_file: bool = False
|
export_metrics_to_file: bool = False
|
||||||
export_metrics_to_file_dir: Optional[str] = None
|
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 related
|
||||||
api_key: Optional[str] = None
|
api_key: Optional[str] = None
|
||||||
admin_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.dp_size = args.data_parallel_size
|
||||||
args.ep_size = args.expert_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})
|
return cls(**{attr: getattr(args, attr) for attr in attrs})
|
||||||
|
|
||||||
def url(self, port: Optional[int] = None):
|
def url(self, port: Optional[int] = None):
|
||||||
|
|||||||
@@ -1,3 +1,4 @@
|
|||||||
|
import os
|
||||||
import unittest
|
import unittest
|
||||||
from typing import Dict, List
|
from typing import Dict, List
|
||||||
|
|
||||||
@@ -8,6 +9,8 @@ from prometheus_client.samples import Sample
|
|||||||
from sglang.srt.environ import envs
|
from sglang.srt.environ import envs
|
||||||
from sglang.srt.observability.metrics_collector import (
|
from sglang.srt.observability.metrics_collector import (
|
||||||
ROUTING_KEY_REQ_COUNT_BUCKET_BOUNDS,
|
ROUTING_KEY_REQ_COUNT_BUCKET_BOUNDS,
|
||||||
|
STAT_LOGGER_ROLE_SCHEDULER,
|
||||||
|
SchedulerMetricsCollector,
|
||||||
compute_routing_key_stats,
|
compute_routing_key_stats,
|
||||||
)
|
)
|
||||||
from sglang.srt.utils import kill_process_tree
|
from sglang.srt.utils import kill_process_tree
|
||||||
@@ -274,6 +277,66 @@ def _check_metrics_positive(test_case, metrics, metrics_to_check):
|
|||||||
test_case.assertGreater(value, 0, f"{metric_name} {labels}")
|
test_case.assertGreater(value, 0, f"{metric_name} {labels}")
|
||||||
|
|
||||||
|
|
||||||
|
_DI_MARKER_PATH = "/tmp/sglang_di_test_marker"
|
||||||
|
|
||||||
|
|
||||||
|
class _MarkingSchedulerCollector(SchedulerMetricsCollector):
|
||||||
|
"""Records its own instantiation to a file so the test can verify the
|
||||||
|
custom subclass was used in the scheduler subprocess.
|
||||||
|
|
||||||
|
Defined at module level so it is picklable into the scheduler process.
|
||||||
|
Cross-process signalling uses a filesystem marker because the scheduler
|
||||||
|
runs in its own subprocess and cannot share in-memory state with the
|
||||||
|
test runner.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, *args, **kwargs):
|
||||||
|
with open(_DI_MARKER_PATH, "w") as f:
|
||||||
|
f.write("scheduler_collector_initialized\n")
|
||||||
|
super().__init__(*args, **kwargs)
|
||||||
|
|
||||||
|
|
||||||
|
class TestStatLoggersDI(CustomTestCase):
|
||||||
|
"""Verify that a custom MetricsCollector subclass passed through
|
||||||
|
``ServerArgs.stat_loggers`` is the one instantiated inside the
|
||||||
|
scheduler subprocess."""
|
||||||
|
|
||||||
|
def setUp(self) -> None:
|
||||||
|
try:
|
||||||
|
os.unlink(_DI_MARKER_PATH)
|
||||||
|
except FileNotFoundError:
|
||||||
|
pass
|
||||||
|
|
||||||
|
def tearDown(self) -> None:
|
||||||
|
try:
|
||||||
|
os.unlink(_DI_MARKER_PATH)
|
||||||
|
except FileNotFoundError:
|
||||||
|
pass
|
||||||
|
|
||||||
|
def test_engine_custom_scheduler_collector(self):
|
||||||
|
import sglang as sgl
|
||||||
|
|
||||||
|
engine = sgl.Engine(
|
||||||
|
model_path=_MODEL_NAME,
|
||||||
|
enable_metrics=True,
|
||||||
|
stat_loggers={
|
||||||
|
STAT_LOGGER_ROLE_SCHEDULER: _MarkingSchedulerCollector,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
# One small generation triggers scheduler init, which is where
|
||||||
|
# resolve_collector_class() picks the injected subclass.
|
||||||
|
engine.generate("Hello", {"max_new_tokens": 4})
|
||||||
|
finally:
|
||||||
|
engine.shutdown()
|
||||||
|
|
||||||
|
self.assertTrue(
|
||||||
|
os.path.exists(_DI_MARKER_PATH),
|
||||||
|
"Custom SchedulerMetricsCollector was not instantiated; "
|
||||||
|
"stat_loggers DI did not take effect.",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class TestComputeRoutingKeyStats(unittest.TestCase):
|
class TestComputeRoutingKeyStats(unittest.TestCase):
|
||||||
def test_empty(self):
|
def test_empty(self):
|
||||||
num_unique, req_counts = compute_routing_key_stats([])
|
num_unique, req_counts = compute_routing_key_stats([])
|
||||||
|
|||||||
@@ -0,0 +1,201 @@
|
|||||||
|
"""Unit tests for class-level DI on the five *MetricsCollector classes via
|
||||||
|
ServerArgs.stat_loggers — no server, no model loading."""
|
||||||
|
|
||||||
|
from sglang.test.ci.ci_register import register_cpu_ci
|
||||||
|
|
||||||
|
register_cpu_ci(est_time=2, suite="base-a-test-cpu")
|
||||||
|
|
||||||
|
import unittest
|
||||||
|
|
||||||
|
import prometheus_client
|
||||||
|
|
||||||
|
from sglang.srt.observability.metrics_collector import (
|
||||||
|
STAT_LOGGER_ROLE_EXPERT_DISPATCH,
|
||||||
|
STAT_LOGGER_ROLE_RADIX_CACHE,
|
||||||
|
STAT_LOGGER_ROLE_SCHEDULER,
|
||||||
|
STAT_LOGGER_ROLE_STORAGE,
|
||||||
|
STAT_LOGGER_ROLE_TOKENIZER,
|
||||||
|
ExpertDispatchCollector,
|
||||||
|
RadixCacheMetricsCollector,
|
||||||
|
SchedulerMetricsCollector,
|
||||||
|
StorageMetricsCollector,
|
||||||
|
TokenizerMetricsCollector,
|
||||||
|
resolve_collector_class,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class _StubArgs:
|
||||||
|
"""Minimal ServerArgs stand-in. Avoids triggering heavy ServerArgs import chain."""
|
||||||
|
|
||||||
|
def __init__(self, stat_loggers=None):
|
||||||
|
self.stat_loggers = stat_loggers
|
||||||
|
|
||||||
|
|
||||||
|
# ── _gauge_cls / _counter_cls / _histogram_cls / _summary_cls override surface ──
|
||||||
|
|
||||||
|
|
||||||
|
class TestCollectorClassAttrs(unittest.TestCase):
|
||||||
|
"""All five collectors expose four DI hook class attrs, all defaulting to None
|
||||||
|
so the existing prometheus_client backend is used unchanged."""
|
||||||
|
|
||||||
|
def test_scheduler_collector_attrs_default_none(self):
|
||||||
|
self.assertIsNone(SchedulerMetricsCollector._counter_cls)
|
||||||
|
self.assertIsNone(SchedulerMetricsCollector._gauge_cls)
|
||||||
|
self.assertIsNone(SchedulerMetricsCollector._histogram_cls)
|
||||||
|
self.assertIsNone(SchedulerMetricsCollector._summary_cls)
|
||||||
|
|
||||||
|
def test_tokenizer_collector_attrs_default_none(self):
|
||||||
|
self.assertIsNone(TokenizerMetricsCollector._counter_cls)
|
||||||
|
self.assertIsNone(TokenizerMetricsCollector._histogram_cls)
|
||||||
|
|
||||||
|
def test_storage_collector_attrs_default_none(self):
|
||||||
|
self.assertIsNone(StorageMetricsCollector._counter_cls)
|
||||||
|
self.assertIsNone(StorageMetricsCollector._histogram_cls)
|
||||||
|
|
||||||
|
def test_expert_dispatch_collector_attrs_default_none(self):
|
||||||
|
self.assertIsNone(ExpertDispatchCollector._histogram_cls)
|
||||||
|
|
||||||
|
def test_radix_cache_collector_attrs_default_none(self):
|
||||||
|
self.assertIsNone(RadixCacheMetricsCollector._counter_cls)
|
||||||
|
self.assertIsNone(RadixCacheMetricsCollector._histogram_cls)
|
||||||
|
|
||||||
|
|
||||||
|
# ── resolve_collector_class helper ──
|
||||||
|
|
||||||
|
|
||||||
|
class TestResolveCollectorClass(unittest.TestCase):
|
||||||
|
def test_returns_default_when_server_args_none(self):
|
||||||
|
cls = resolve_collector_class(None, "scheduler", SchedulerMetricsCollector)
|
||||||
|
self.assertIs(cls, SchedulerMetricsCollector)
|
||||||
|
|
||||||
|
def test_returns_default_when_stat_loggers_none(self):
|
||||||
|
cls = resolve_collector_class(
|
||||||
|
_StubArgs(stat_loggers=None), "scheduler", SchedulerMetricsCollector
|
||||||
|
)
|
||||||
|
self.assertIs(cls, SchedulerMetricsCollector)
|
||||||
|
|
||||||
|
def test_returns_default_when_stat_loggers_empty(self):
|
||||||
|
cls = resolve_collector_class(
|
||||||
|
_StubArgs(stat_loggers={}), "scheduler", SchedulerMetricsCollector
|
||||||
|
)
|
||||||
|
self.assertIs(cls, SchedulerMetricsCollector)
|
||||||
|
|
||||||
|
def test_returns_default_when_role_missing(self):
|
||||||
|
# Different role registered. Default still wins for "scheduler".
|
||||||
|
class MyTokenizer(TokenizerMetricsCollector):
|
||||||
|
pass
|
||||||
|
|
||||||
|
cls = resolve_collector_class(
|
||||||
|
_StubArgs(stat_loggers={"tokenizer": MyTokenizer}),
|
||||||
|
"scheduler",
|
||||||
|
SchedulerMetricsCollector,
|
||||||
|
)
|
||||||
|
self.assertIs(cls, SchedulerMetricsCollector)
|
||||||
|
|
||||||
|
def test_returns_subclass_when_role_registered(self):
|
||||||
|
class MyScheduler(SchedulerMetricsCollector):
|
||||||
|
pass
|
||||||
|
|
||||||
|
cls = resolve_collector_class(
|
||||||
|
_StubArgs(stat_loggers={"scheduler": MyScheduler}),
|
||||||
|
"scheduler",
|
||||||
|
SchedulerMetricsCollector,
|
||||||
|
)
|
||||||
|
self.assertIs(cls, MyScheduler)
|
||||||
|
|
||||||
|
def test_role_constants_match_collector_keys(self):
|
||||||
|
"""The exported role constants must be the exact strings the
|
||||||
|
instantiation sites use to look up subclasses."""
|
||||||
|
self.assertEqual(STAT_LOGGER_ROLE_SCHEDULER, "scheduler")
|
||||||
|
self.assertEqual(STAT_LOGGER_ROLE_TOKENIZER, "tokenizer")
|
||||||
|
self.assertEqual(STAT_LOGGER_ROLE_STORAGE, "storage")
|
||||||
|
self.assertEqual(STAT_LOGGER_ROLE_RADIX_CACHE, "radix_cache")
|
||||||
|
self.assertEqual(STAT_LOGGER_ROLE_EXPERT_DISPATCH, "expert_dispatch")
|
||||||
|
|
||||||
|
|
||||||
|
# ── DI swap behavior — actually instantiate with a custom backend ──
|
||||||
|
|
||||||
|
|
||||||
|
class _RecordingGauge:
|
||||||
|
"""Test double that mirrors prometheus_client.Gauge constructor signature.
|
||||||
|
Records every instantiation so the test can assert the override took effect."""
|
||||||
|
|
||||||
|
instances = []
|
||||||
|
|
||||||
|
def __init__(self, *args, **kwargs):
|
||||||
|
type(self).instances.append((args, kwargs))
|
||||||
|
|
||||||
|
def labels(self, **kwargs):
|
||||||
|
return self
|
||||||
|
|
||||||
|
def set(self, value):
|
||||||
|
pass
|
||||||
|
|
||||||
|
def inc(self, amount=1):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class _RecordingCounter(_RecordingGauge):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class _RecordingHistogram(_RecordingGauge):
|
||||||
|
def observe(self, value):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class _RecordingSummary(_RecordingGauge):
|
||||||
|
def observe(self, value):
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class TestDISwap(unittest.TestCase):
|
||||||
|
"""Subclasses that set the DI hooks at class level cause the collector to
|
||||||
|
instantiate the test doubles instead of prometheus_client classes."""
|
||||||
|
|
||||||
|
def setUp(self):
|
||||||
|
_RecordingGauge.instances = []
|
||||||
|
_RecordingCounter.instances = []
|
||||||
|
_RecordingHistogram.instances = []
|
||||||
|
_RecordingSummary.instances = []
|
||||||
|
|
||||||
|
def test_radix_cache_di_swap(self):
|
||||||
|
"""Smallest collector (4 metrics, Counter + Histogram) — verifies the
|
||||||
|
DI shim flows through both class types."""
|
||||||
|
|
||||||
|
class RaySwapRadixCache(RadixCacheMetricsCollector):
|
||||||
|
_counter_cls = _RecordingCounter
|
||||||
|
_histogram_cls = _RecordingHistogram
|
||||||
|
|
||||||
|
labels = {"cache_type": "test"}
|
||||||
|
RaySwapRadixCache(labels=labels)
|
||||||
|
|
||||||
|
# 4 instruments total in RadixCacheMetricsCollector:
|
||||||
|
# eviction_duration_seconds (H), eviction_num_tokens (C),
|
||||||
|
# load_back_duration_seconds (H), load_back_num_tokens (C).
|
||||||
|
self.assertEqual(len(_RecordingCounter.instances), 2)
|
||||||
|
self.assertEqual(len(_RecordingHistogram.instances), 2)
|
||||||
|
|
||||||
|
def test_expert_dispatch_di_swap(self):
|
||||||
|
"""Smallest collector (1 Histogram metric)."""
|
||||||
|
|
||||||
|
class RaySwapExpert(ExpertDispatchCollector):
|
||||||
|
_histogram_cls = _RecordingHistogram
|
||||||
|
|
||||||
|
RaySwapExpert(ep_size=4)
|
||||||
|
self.assertEqual(len(_RecordingHistogram.instances), 1)
|
||||||
|
|
||||||
|
def test_default_path_uses_prometheus_client(self):
|
||||||
|
"""Without any subclass override, the collector instantiates the real
|
||||||
|
prometheus_client classes — the existing behavior is unchanged."""
|
||||||
|
labels = {"cache_type": "test_default"}
|
||||||
|
collector = RadixCacheMetricsCollector(labels=labels)
|
||||||
|
# The instruments must be real prometheus_client objects, not test doubles.
|
||||||
|
self.assertIsInstance(collector.eviction_num_tokens, prometheus_client.Counter)
|
||||||
|
self.assertIsInstance(
|
||||||
|
collector.eviction_duration_seconds, prometheus_client.Histogram
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
Reference in New Issue
Block a user