[observability] add Ray metric backend wrappers (#26252)
Signed-off-by: Dongjun Na <kmu5544616@gmail.com>
This commit is contained in:
Dongjun Na
@@ -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
|
||||
Reference in New Issue
Block a user