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

[observability] add Ray metric backend wrappers (#26252)

Browse Source 
Signed-off-by: Dongjun Na <kmu5544616@gmail.com>
This commit is contained in:
Dongjun Na
2026-06-17 19:41:53 -07:00
committed by GitHub
parent d2b5488392
commit 462c01ea6b
6 changed files with 983 additions and 86 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/observability/ray_wrappers.py
+307
View File
@@ -0,0 +1,307 @@
# Copyright 2023-2024 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Ray-backed implementations of the prometheus_client API surface used by
sglang's ``*MetricsCollector`` classes.
The wrappers translate prometheus_client calls into ``ray.util.metrics`` so the
metrics emitted by an embedded sglang engine flow through Ray's metric agent and
appear on Ray's Prometheus endpoint / dashboard alongside other Ray metrics.
Mirrors ``vllm/v1/metrics/ray_wrappers.py`` with two sglang-specific additions:
* ``RaySummaryWrapper`` — Ray has no Summary primitive; we fall back to a
Histogram with conservative default boundaries. Quantile queries can be
approximated through ``histogram_quantile()`` in Prometheus.
* Five collector subclasses, one per ``*MetricsCollector`` defined in
:mod:`sglang.srt.observability.metrics_collector`, overriding only the
``_xxx_cls`` attributes that the corresponding collector actually uses.
Import is lazy: the module loads in environments without Ray installed, but
instantiating a wrapper without Ray raises a clear :class:`ImportError`.
"""
from __future__ import annotations
import copy
import re
import time
from typing import List, Optional
try:
from ray import serve as ray_serve
from ray.util import metrics as ray_metrics
from ray.util.metrics import Metric
except ImportError: # pragma: no cover - covered by a dedicated test
ray_metrics = None
ray_serve = None
Metric = None # type: ignore[assignment]
from sglang.srt.observability.metrics_collector import (
ExpertDispatchCollector,
RadixCacheMetricsCollector,
SchedulerMetricsCollector,
StorageMetricsCollector,
TokenizerMetricsCollector,
)
def _get_replica_id() -> Optional[str]:
"""Return the current Ray Serve replica ID, or ``None`` outside Serve."""
if ray_serve is None:
return None
try:
return ray_serve.get_replica_context().replica_id.unique_id
except ray_serve.exceptions.RayServeException:
return None
class RayPrometheusMetric:
"""Base wrapper that exposes the prometheus_client API on Ray metrics.
Subclasses populate ``self.metric`` with a ``ray.util.metrics`` instance in
their ``__init__``. Shared behaviour:
* A ``ReplicaId`` tag is appended to every metric and populated at
instantiation (and again on each ``labels()`` call) so Ray-Serve replicas
are distinguishable on dashboards.
* ``labels()`` returns a fresh copy of the wrapper with its tags bound,
mirroring the ``prometheus_client`` pattern and avoiding state sharing
between concurrent emits.
* Metric names are sanitised to satisfy Ray's OpenTelemetry naming rule
(no ``:``, no other punctuation).
"""
_is_labeled: bool = False
def __init__(self) -> None:
if ray_metrics is None:
raise ImportError(
"RayPrometheusMetric requires Ray to be installed. "
"Install with: pip install 'ray[serve]'"
)
self.metric: Optional[Metric] = None
self._tags: dict = {"ReplicaId": _get_replica_id() or ""}
@staticmethod
def _get_tag_keys(labelnames: Optional[List[str]]) -> tuple:
labels = list(labelnames) if labelnames else []
labels.append("ReplicaId")
return tuple(labels)
def _build_tags(self, *labels: str, **labelskwargs: str) -> dict:
if labels:
# The trailing entry of ``_tag_keys`` is always ``ReplicaId`` which we
# populate ourselves; positional args fill the preceding keys only.
expected = len(self.metric._tag_keys) - 1
if len(labels) != expected:
raise ValueError(
"Number of labels must match the number of tag keys. "
f"Expected {expected}, got {len(labels)}"
)
labelskwargs.update(zip(self.metric._tag_keys, labels))
labelskwargs["ReplicaId"] = _get_replica_id() or ""
return {k: v if isinstance(v, str) else str(v) for k, v in labelskwargs.items()}
def labels(self, *labels: str, **labelskwargs: str) -> RayPrometheusMetric:
if self._is_labeled:
raise ValueError("labels() cannot be called on an already-labeled metric.")
clone = copy.copy(self)
clone._tags = self._build_tags(*labels, **labelskwargs)
clone._is_labeled = True
return clone
@staticmethod
def _coerce_positive_boundaries(buckets):
# Ray (gRPC OpenCensus / OpenTelemetry export) rejects boundaries
# <= 0. sglang ships several histograms whose lowest bucket is 0.0
# (e.g. queue_time, e2e latency). Silently drop those so we never
# break engine startup when the metrics backend is Ray.
if not buckets:
return []
return [b for b in buckets if b > 0]
@staticmethod
def _get_sanitized_opentelemetry_name(name: str) -> str:
"""Replace characters Ray's OTel-backed metric name validator rejects.
Ray is migrating from OpenCensus to OpenTelemetry, whose instrument names
only allow ``a-zA-Z0-9_``. sglang's existing names use a ``sglang:foo``
prefix; converting ``:`` (and any other punctuation) to ``_`` keeps the
names valid without churn on the prometheus_client side.
"""
return re.sub(r"[^a-zA-Z0-9_]", "_", name)
class RayCounterWrapper(RayPrometheusMetric):
"""``prometheus_client.Counter`` compatible wrapper."""
def __init__(
self,
name: str,
documentation: Optional[str] = "",
labelnames: Optional[List[str]] = None,
) -> None:
super().__init__()
tag_keys = self._get_tag_keys(labelnames)
name = self._get_sanitized_opentelemetry_name(name)
self.metric = ray_metrics.Counter(
name=name,
description=documentation,
tag_keys=tag_keys,
)
def inc(self, value: float = 1.0) -> None:
if value == 0:
return
return self.metric.inc(value, tags=self._tags)
class RayGaugeWrapper(RayPrometheusMetric):
"""``prometheus_client.Gauge`` compatible wrapper."""
def __init__(
self,
name: str,
documentation: Optional[str] = "",
labelnames: Optional[List[str]] = None,
multiprocess_mode: Optional[str] = "",
) -> None:
# Ray aggregates per WorkerId/ReplicaId at the metric agent, so the
# prometheus_client multiproc modes ("mostrecent", "all", "sum") are not
# meaningful here. Accept and discard for API parity.
del multiprocess_mode
super().__init__()
tag_keys = self._get_tag_keys(labelnames)
name = self._get_sanitized_opentelemetry_name(name)
self.metric = ray_metrics.Gauge(
name=name,
description=documentation,
tag_keys=tag_keys,
)
def set(self, value: float) -> None:
return self.metric.set(value, tags=self._tags)
def set_to_current_time(self) -> None:
return self.set(time.time())
class RayHistogramWrapper(RayPrometheusMetric):
"""``prometheus_client.Histogram`` compatible wrapper."""
def __init__(
self,
name: str,
documentation: Optional[str] = "",
labelnames: Optional[List[str]] = None,
buckets: Optional[List[float]] = None,
) -> None:
super().__init__()
tag_keys = self._get_tag_keys(labelnames)
name = self._get_sanitized_opentelemetry_name(name)
self.metric = ray_metrics.Histogram(
name=name,
description=documentation,
tag_keys=tag_keys,
boundaries=self._coerce_positive_boundaries(buckets),
)
def observe(self, value: float) -> None:
return self.metric.observe(value, tags=self._tags)
class RaySummaryWrapper(RayPrometheusMetric):
"""``prometheus_client.Summary`` compatible wrapper.
``ray.util.metrics`` does not provide a Summary primitive. We approximate by
emitting through a Histogram with conservative default boundaries; quantile
queries can be approximated downstream via ``histogram_quantile()``.
"""
DEFAULT_BOUNDARIES: List[float] = [
0.005,
0.01,
0.025,
0.05,
0.1,
0.25,
0.5,
1.0,
2.5,
5.0,
10.0,
]
def __init__(
self,
name: str,
documentation: Optional[str] = "",
labelnames: Optional[List[str]] = None,
) -> None:
super().__init__()
tag_keys = self._get_tag_keys(labelnames)
name = self._get_sanitized_opentelemetry_name(name)
self.metric = ray_metrics.Histogram(
name=name,
description=documentation,
tag_keys=tag_keys,
boundaries=self._coerce_positive_boundaries(self.DEFAULT_BOUNDARIES),
)
def observe(self, value: float) -> None:
return self.metric.observe(value, tags=self._tags)
# ---------------------------------------------------------------------------
# Collector subclasses
#
# Each subclass only overrides the ``_xxx_cls`` attributes its parent actually
# uses; the parent's ``_StatLoggerDIMixin`` defaults handle the rest.
# ---------------------------------------------------------------------------
class RaySchedulerMetricsCollector(SchedulerMetricsCollector):
"""``SchedulerMetricsCollector`` that emits via Ray's metric system."""
_counter_cls = RayCounterWrapper
_gauge_cls = RayGaugeWrapper
_histogram_cls = RayHistogramWrapper
_summary_cls = RaySummaryWrapper
class RayTokenizerMetricsCollector(TokenizerMetricsCollector):
"""``TokenizerMetricsCollector`` that emits via Ray's metric system."""
_counter_cls = RayCounterWrapper
_histogram_cls = RayHistogramWrapper
class RayStorageMetricsCollector(StorageMetricsCollector):
"""``StorageMetricsCollector`` that emits via Ray's metric system."""
_counter_cls = RayCounterWrapper
_histogram_cls = RayHistogramWrapper
class RayRadixCacheMetricsCollector(RadixCacheMetricsCollector):
"""``RadixCacheMetricsCollector`` that emits via Ray's metric system."""
_counter_cls = RayCounterWrapper
_histogram_cls = RayHistogramWrapper
class RayExpertDispatchCollector(ExpertDispatchCollector):
"""``ExpertDispatchCollector`` that emits via Ray's metric system."""
_histogram_cls = RayHistogramWrapper
python/sglang/test/observability/__init__.py
View File
python/sglang/test/observability/fake_ray.py
+142
View File
@@ -0,0 +1,142 @@
# Copyright 2023-2024 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Shared test doubles for Ray-backed observability code paths.
These fakes let tests exercise :mod:`sglang.srt.observability.ray_wrappers`,
and any DI surface that wires through it, without requiring Ray to be
installed. The module is consumed by:
* ``test/registered/unit/observability/test_ray_wrappers.py`` — wrapper unit
tests.
* ``test/registered/unit/observability/test_stat_loggers_di.py`` — DI
integration tests that verify emissions flow through to the metric instance.
"""
from __future__ import annotations
import importlib
import sys
import types
RAY_FAKE_MODULE_NAMES = (
"ray",
"ray.util",
"ray.util.metrics",
"ray.serve",
"ray.serve.exceptions",
"sglang.srt.observability.ray_wrappers",
)
class FakeRayMetric:
"""Stand-in for ``ray.util.metrics.{Counter,Gauge,Histogram}``.
Records every ``inc`` / ``set`` / ``observe`` call so tests can assert on
the forwarded value and tags dict.
"""
def __init__(
self,
name: str = "",
description: str = "",
tag_keys: tuple = (),
boundaries=None,
):
self.name = name
self.description = description
self._tag_keys = tuple(tag_keys)
self.boundaries = list(boundaries) if boundaries is not None else None
self.calls = [] # list of (op, value, tags)
def inc(self, value, tags=None):
self.calls.append(("inc", value, dict(tags or {})))
def set(self, value, tags=None):
self.calls.append(("set", value, dict(tags or {})))
def observe(self, value, tags=None):
self.calls.append(("observe", value, dict(tags or {})))
class FakeRayServeException(Exception):
pass
def make_fake_ray_modules(replica_id: str = "test-replica") -> dict:
"""Build a dict of fake ray modules suitable for ``sys.modules.update``.
Each call returns fresh module objects so different tests can use
different ``replica_id`` values without cross-contamination.
"""
ray_pkg = types.ModuleType("ray")
ray_util = types.ModuleType("ray.util")
ray_util_metrics = types.ModuleType("ray.util.metrics")
ray_util_metrics.Counter = FakeRayMetric
ray_util_metrics.Gauge = FakeRayMetric
ray_util_metrics.Histogram = FakeRayMetric
ray_util_metrics.Metric = FakeRayMetric
ray_serve = types.ModuleType("ray.serve")
ray_serve_exc = types.ModuleType("ray.serve.exceptions")
ray_serve_exc.RayServeException = FakeRayServeException
ray_serve.exceptions = ray_serve_exc
class _ReplicaCtx:
class _Id:
unique_id = replica_id
replica_id = _Id()
ray_serve.get_replica_context = lambda: _ReplicaCtx()
return {
"ray": ray_pkg,
"ray.util": ray_util,
"ray.util.metrics": ray_util_metrics,
"ray.serve": ray_serve,
"ray.serve.exceptions": ray_serve_exc,
}
def load_ray_wrappers_with_fake_ray(replica_id: str = "test-replica"):
"""Inject fake ray modules into ``sys.modules`` and (re)import ``ray_wrappers``."""
fake = make_fake_ray_modules(replica_id=replica_id)
sys.modules.update(fake)
sys.modules.pop("sglang.srt.observability.ray_wrappers", None)
return importlib.import_module("sglang.srt.observability.ray_wrappers")
def load_ray_wrappers_without_ray():
"""Make ``import ray`` fail and (re)import ``ray_wrappers`` cleanly."""
for name in (
"ray",
"ray.util",
"ray.util.metrics",
"ray.serve",
"ray.serve.exceptions",
):
sys.modules.pop(name, None)
sys.modules[name] = None # type: ignore[assignment]
sys.modules.pop("sglang.srt.observability.ray_wrappers", None)
return importlib.import_module("sglang.srt.observability.ray_wrappers")
def clear_fake_ray_modules() -> None:
"""Remove the fake ray modules and the cached ray_wrappers import.
Tests call this from ``tearDown`` so module state does not leak between
tests.
"""
for name in RAY_FAKE_MODULE_NAMES:
sys.modules.pop(name, None)