diff --git a/python/sglang/srt/eplb/expert_distribution.py b/python/sglang/srt/eplb/expert_distribution.py index 30a1f302c..17d20fbbd 100644 --- a/python/sglang/srt/eplb/expert_distribution.py +++ b/python/sglang/srt/eplb/expert_distribution.py @@ -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 diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index ec2b34926..cae9ed06f 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -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, diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 802cf8c66..6eaefc714 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -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: diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py index 76a5426aa..c3e4c7a80 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -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 diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 01ad1966c..c8cadf0b4 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -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: diff --git a/python/sglang/srt/observability/metrics_collector.py b/python/sglang/srt/observability/metrics_collector.py index d82b7eeb9..41c6496d3 100644 --- a/python/sglang/srt/observability/metrics_collector.py +++ b/python/sglang/srt/observability/metrics_collector.py @@ -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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 2ad0d6226..48b7bd163 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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): diff --git a/test/registered/observability/test_metrics.py b/test/registered/observability/test_metrics.py index 7e812b881..a793da567 100644 --- a/test/registered/observability/test_metrics.py +++ b/test/registered/observability/test_metrics.py @@ -1,3 +1,4 @@ +import os import unittest from typing import Dict, List @@ -8,6 +9,8 @@ from prometheus_client.samples import Sample from sglang.srt.environ import envs from sglang.srt.observability.metrics_collector import ( ROUTING_KEY_REQ_COUNT_BUCKET_BOUNDS, + STAT_LOGGER_ROLE_SCHEDULER, + SchedulerMetricsCollector, compute_routing_key_stats, ) 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}") +_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): def test_empty(self): num_unique, req_counts = compute_routing_key_stats([]) diff --git a/test/registered/unit/observability/test_stat_loggers_di.py b/test/registered/unit/observability/test_stat_loggers_di.py new file mode 100644 index 000000000..271a4cbed --- /dev/null +++ b/test/registered/unit/observability/test_stat_loggers_di.py @@ -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()