From bf68369a181c05dc19e46323c12e6aeaa4c08a82 Mon Sep 17 00:00:00 2001 From: Jeffrey Wang Date: Mon, 7 Sep 2026 12:43:19 -0700 Subject: [PATCH] [ray] Support Ray metric backend for engine metrics (#31415) Signed-off-by: Jeffrey Wang Co-authored-by: Claude Opus 4.8 (1M context) Co-authored-by: Qiaolin Yu --- .../srt/observability/metrics_collector.py | 13 ++- .../sglang/srt/observability/ray_wrappers.py | 39 ++++++++- python/sglang/srt/ray/engine.py | 6 ++ .../unit/observability/test_ray_wrappers.py | 79 +++++++++++++++++++ 4 files changed, 132 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/observability/metrics_collector.py b/python/sglang/srt/observability/metrics_collector.py index b8129a4a8..8f7d580e4 100644 --- a/python/sglang/srt/observability/metrics_collector.py +++ b/python/sglang/srt/observability/metrics_collector.py @@ -1847,6 +1847,10 @@ class TokenizerMetricsCollector(_StatLoggerDIMixin): ).observe(value) def check_time_to_first_token_straggler(self, value: float) -> bool: + # Injected backends (e.g. Ray) route metrics out of process and can't + # introspect prometheus_client buckets here. + if self._histogram_cls is not None: + return False his = self.histogram_time_to_first_token.labels( **self.labels, is_streaming="true" ) @@ -1865,10 +1869,17 @@ class TokenizerMetricsCollector(_StatLoggerDIMixin): self, labels: Dict[str, str], internval: float, num_new_tokens: int ): adjusted_interval = internval / num_new_tokens + his = self.histogram_inter_token_latency.labels(**labels) + + if self._histogram_cls is not None: + # Injected backend (e.g. Ray): the bucket internals below don't + # exist, so use the public observe() API. + for _ in range(num_new_tokens): + his.observe(adjusted_interval) + return # A faster version of the Histogram::observe which observes multiple values at the same time. # reference: https://github.com/prometheus/client_python/blob/v0.21.1/prometheus_client/metrics.py#L639 - his = self.histogram_inter_token_latency.labels(**labels) his._sum.inc(internval) for i, bound in enumerate(his._upper_bounds): diff --git a/python/sglang/srt/observability/ray_wrappers.py b/python/sglang/srt/observability/ray_wrappers.py index 62faff0b8..828950df2 100644 --- a/python/sglang/srt/observability/ray_wrappers.py +++ b/python/sglang/srt/observability/ray_wrappers.py @@ -142,6 +142,18 @@ class RayPrometheusMetric: """ return re.sub(r"[^a-zA-Z0-9_]", "_", name) + @staticmethod + def _get_ascii_documentation(documentation: Optional[str]) -> str: + """ASCII-coerce a description; Ray's metric backend rejects non-ASCII.""" + if not documentation: + return documentation or "" + return ( + documentation.replace("—", "-") + .replace("–", "-") + .encode("ascii", "ignore") + .decode("ascii") + ) + class RayCounterWrapper(RayPrometheusMetric): """``prometheus_client.Counter`` compatible wrapper.""" @@ -157,7 +169,7 @@ class RayCounterWrapper(RayPrometheusMetric): name = self._get_sanitized_opentelemetry_name(name) self.metric = ray_metrics.Counter( name=name, - description=documentation, + description=self._get_ascii_documentation(documentation), tag_keys=tag_keys, ) @@ -186,7 +198,7 @@ class RayGaugeWrapper(RayPrometheusMetric): name = self._get_sanitized_opentelemetry_name(name) self.metric = ray_metrics.Gauge( name=name, - description=documentation, + description=self._get_ascii_documentation(documentation), tag_keys=tag_keys, ) @@ -212,7 +224,7 @@ class RayHistogramWrapper(RayPrometheusMetric): name = self._get_sanitized_opentelemetry_name(name) self.metric = ray_metrics.Histogram( name=name, - description=documentation, + description=self._get_ascii_documentation(documentation), tag_keys=tag_keys, boundaries=self._coerce_positive_boundaries(buckets), ) @@ -254,7 +266,7 @@ class RaySummaryWrapper(RayPrometheusMetric): name = self._get_sanitized_opentelemetry_name(name) self.metric = ray_metrics.Histogram( name=name, - description=documentation, + description=self._get_ascii_documentation(documentation), tag_keys=tag_keys, boundaries=self._coerce_positive_boundaries(self.DEFAULT_BOUNDARIES), ) @@ -306,3 +318,22 @@ class RayExpertDispatchCollector(ExpertDispatchCollector): """``ExpertDispatchCollector`` that emits via Ray's metric system.""" _histogram_cls = RayHistogramWrapper + + +def build_ray_stat_loggers() -> dict: + """Build the ``ServerArgs.stat_loggers`` map of role -> Ray-backed collector.""" + 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, + ) + + return { + STAT_LOGGER_ROLE_SCHEDULER: RaySchedulerMetricsCollector, + STAT_LOGGER_ROLE_TOKENIZER: RayTokenizerMetricsCollector, + STAT_LOGGER_ROLE_STORAGE: RayStorageMetricsCollector, + STAT_LOGGER_ROLE_RADIX_CACHE: RayRadixCacheMetricsCollector, + STAT_LOGGER_ROLE_EXPERT_DISPATCH: RayExpertDispatchCollector, + } diff --git a/python/sglang/srt/ray/engine.py b/python/sglang/srt/ray/engine.py index 545d5145c..35cb9c9c8 100644 --- a/python/sglang/srt/ray/engine.py +++ b/python/sglang/srt/ray/engine.py @@ -240,6 +240,12 @@ class RayEngine(Engine): self._placement_group = kwargs.pop("placement_group", None) if "log_level" not in kwargs: kwargs["log_level"] = "error" + # Schedulers are separate Ray actors; default to the Ray-backed + # collectors so enable_metrics reaches Ray's Prometheus endpoint. + if kwargs.get("enable_metrics") and kwargs.get("stat_loggers") is None: + from sglang.srt.observability.ray_wrappers import build_ray_stat_loggers + + kwargs["stat_loggers"] = build_ray_stat_loggers() super().__init__(server_args=ServerArgs(**kwargs)) def shutdown(self): diff --git a/test/registered/unit/observability/test_ray_wrappers.py b/test/registered/unit/observability/test_ray_wrappers.py index de44e262d..094be4529 100644 --- a/test/registered/unit/observability/test_ray_wrappers.py +++ b/test/registered/unit/observability/test_ray_wrappers.py @@ -11,13 +11,19 @@ from __future__ import annotations import sys import unittest +from functools import partial +from prometheus_client import CollectorRegistry, Counter, Gauge, Histogram + +from sglang.srt.observability.metrics_collector import TokenizerMetricsCollector +from sglang.srt.runtime_context import get_context from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.observability.fake_ray import ( clear_fake_ray_modules, load_ray_wrappers_with_fake_ray, load_ray_wrappers_without_ray, ) +from sglang.test.test_utils import CustomTestCase register_cpu_ci(est_time=11, suite="base-a-test-cpu") @@ -277,5 +283,78 @@ class TestRayMissingImportError(unittest.TestCase): self.assertIsNone(self.rw._get_replica_id()) +class TestAsciiDocumentation(TestRayWrapperBase): + """Ray's metric backend rejects non-ASCII, so a wrapper whose constructor + skips ``_get_ascii_documentation`` would only crash at deploy time.""" + + def test_all_wrappers_fold_non_ascii_description(self): + for cls in ( + self.rw.RayCounterWrapper, + self.rw.RayGaugeWrapper, + self.rw.RayHistogramWrapper, + self.rw.RaySummaryWrapper, + ): + with self.subTest(wrapper=cls.__name__): + metric = cls("sglang:x", documentation="load — seconds").metric + self.assertEqual(metric.description, "load - seconds") + + +class TestInterTokenLatencyEquivalence(CustomTestCase): + """``observe_inter_token_latency`` writes histogram internals directly for + the default backend but replays ``observe()`` for an injected one; both must + record identical sums and bucket counts, or ITL diverges between the default + and Ray backends.""" + + _BUCKETS = [0.05, 0.1, 0.5, 1.0] + _LABELS = {"model_name": "m"} + + def setUp(self): + super().setUp() + override = get_context().override_server_args( + prompt_tokens_buckets=None, + generation_tokens_buckets=None, + ) + self.server_args = override.install() + self.addCleanup(override.restore) + + def _build_collector(self, *, force_fallback: bool): + # A private registry per collector avoids duplicate ``sglang:`` names. + registry = CollectorRegistry() + + class _Collector(TokenizerMetricsCollector): + _counter_cls = partial(Counter, registry=registry) + _gauge_cls = partial(Gauge, registry=registry) + _histogram_cls = partial(Histogram, registry=registry) + + collector = _Collector( + server_args=self.server_args, + labels=self._LABELS, + bucket_time_to_first_token=[0.1, 1.0], + bucket_inter_token_latency=self._BUCKETS, + bucket_e2e_request_latency=[0.1, 1.0], + ) + if not force_fallback: + # _histogram_cls=None routes observe to the default-backend path. + collector._histogram_cls = None + return collector + + def test_fast_and_fallback_agree(self): + fast = self._build_collector(force_fallback=False) + fallback = self._build_collector(force_fallback=True) + + for collector in (fast, fallback): + collector.observe_inter_token_latency(self._LABELS, 0.24, 4) + collector.observe_inter_token_latency(self._LABELS, 6.0, 3) # +Inf bucket + collector.observe_inter_token_latency(self._LABELS, 0.3, 2) + + fast_h = fast.histogram_inter_token_latency.labels(**self._LABELS) + fb_h = fallback.histogram_inter_token_latency.labels(**self._LABELS) + self.assertEqual( + [b.get() for b in fast_h._buckets], + [b.get() for b in fb_h._buckets], + ) + self.assertAlmostEqual(fast_h._sum.get(), fb_h._sum.get()) + + if __name__ == "__main__": sys.exit(unittest.main())