[ray] Support Ray metric backend for engine metrics (#31415)
Signed-off-by: Jeffrey Wang <jeffreywang@anyscale.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Qiaolin Yu <liin1211@outlook.com>
This commit is contained in:
co-authored by
Claude Opus 4.8
Qiaolin Yu
parent
8392c36bce
commit
bf68369a18
@@ -1847,6 +1847,10 @@ class TokenizerMetricsCollector(_StatLoggerDIMixin):
|
|||||||
).observe(value)
|
).observe(value)
|
||||||
|
|
||||||
def check_time_to_first_token_straggler(self, value: float) -> bool:
|
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(
|
his = self.histogram_time_to_first_token.labels(
|
||||||
**self.labels, is_streaming="true"
|
**self.labels, is_streaming="true"
|
||||||
)
|
)
|
||||||
@@ -1865,10 +1869,17 @@ class TokenizerMetricsCollector(_StatLoggerDIMixin):
|
|||||||
self, labels: Dict[str, str], internval: float, num_new_tokens: int
|
self, labels: Dict[str, str], internval: float, num_new_tokens: int
|
||||||
):
|
):
|
||||||
adjusted_interval = internval / num_new_tokens
|
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.
|
# 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
|
# 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)
|
his._sum.inc(internval)
|
||||||
|
|
||||||
for i, bound in enumerate(his._upper_bounds):
|
for i, bound in enumerate(his._upper_bounds):
|
||||||
|
|||||||
@@ -142,6 +142,18 @@ class RayPrometheusMetric:
|
|||||||
"""
|
"""
|
||||||
return re.sub(r"[^a-zA-Z0-9_]", "_", name)
|
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):
|
class RayCounterWrapper(RayPrometheusMetric):
|
||||||
"""``prometheus_client.Counter`` compatible wrapper."""
|
"""``prometheus_client.Counter`` compatible wrapper."""
|
||||||
@@ -157,7 +169,7 @@ class RayCounterWrapper(RayPrometheusMetric):
|
|||||||
name = self._get_sanitized_opentelemetry_name(name)
|
name = self._get_sanitized_opentelemetry_name(name)
|
||||||
self.metric = ray_metrics.Counter(
|
self.metric = ray_metrics.Counter(
|
||||||
name=name,
|
name=name,
|
||||||
description=documentation,
|
description=self._get_ascii_documentation(documentation),
|
||||||
tag_keys=tag_keys,
|
tag_keys=tag_keys,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -186,7 +198,7 @@ class RayGaugeWrapper(RayPrometheusMetric):
|
|||||||
name = self._get_sanitized_opentelemetry_name(name)
|
name = self._get_sanitized_opentelemetry_name(name)
|
||||||
self.metric = ray_metrics.Gauge(
|
self.metric = ray_metrics.Gauge(
|
||||||
name=name,
|
name=name,
|
||||||
description=documentation,
|
description=self._get_ascii_documentation(documentation),
|
||||||
tag_keys=tag_keys,
|
tag_keys=tag_keys,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -212,7 +224,7 @@ class RayHistogramWrapper(RayPrometheusMetric):
|
|||||||
name = self._get_sanitized_opentelemetry_name(name)
|
name = self._get_sanitized_opentelemetry_name(name)
|
||||||
self.metric = ray_metrics.Histogram(
|
self.metric = ray_metrics.Histogram(
|
||||||
name=name,
|
name=name,
|
||||||
description=documentation,
|
description=self._get_ascii_documentation(documentation),
|
||||||
tag_keys=tag_keys,
|
tag_keys=tag_keys,
|
||||||
boundaries=self._coerce_positive_boundaries(buckets),
|
boundaries=self._coerce_positive_boundaries(buckets),
|
||||||
)
|
)
|
||||||
@@ -254,7 +266,7 @@ class RaySummaryWrapper(RayPrometheusMetric):
|
|||||||
name = self._get_sanitized_opentelemetry_name(name)
|
name = self._get_sanitized_opentelemetry_name(name)
|
||||||
self.metric = ray_metrics.Histogram(
|
self.metric = ray_metrics.Histogram(
|
||||||
name=name,
|
name=name,
|
||||||
description=documentation,
|
description=self._get_ascii_documentation(documentation),
|
||||||
tag_keys=tag_keys,
|
tag_keys=tag_keys,
|
||||||
boundaries=self._coerce_positive_boundaries(self.DEFAULT_BOUNDARIES),
|
boundaries=self._coerce_positive_boundaries(self.DEFAULT_BOUNDARIES),
|
||||||
)
|
)
|
||||||
@@ -306,3 +318,22 @@ class RayExpertDispatchCollector(ExpertDispatchCollector):
|
|||||||
"""``ExpertDispatchCollector`` that emits via Ray's metric system."""
|
"""``ExpertDispatchCollector`` that emits via Ray's metric system."""
|
||||||
|
|
||||||
_histogram_cls = RayHistogramWrapper
|
_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,
|
||||||
|
}
|
||||||
|
|||||||
@@ -240,6 +240,12 @@ class RayEngine(Engine):
|
|||||||
self._placement_group = kwargs.pop("placement_group", None)
|
self._placement_group = kwargs.pop("placement_group", None)
|
||||||
if "log_level" not in kwargs:
|
if "log_level" not in kwargs:
|
||||||
kwargs["log_level"] = "error"
|
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))
|
super().__init__(server_args=ServerArgs(**kwargs))
|
||||||
|
|
||||||
def shutdown(self):
|
def shutdown(self):
|
||||||
|
|||||||
@@ -11,13 +11,19 @@ from __future__ import annotations
|
|||||||
|
|
||||||
import sys
|
import sys
|
||||||
import unittest
|
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.ci.ci_register import register_cpu_ci
|
||||||
from sglang.test.observability.fake_ray import (
|
from sglang.test.observability.fake_ray import (
|
||||||
clear_fake_ray_modules,
|
clear_fake_ray_modules,
|
||||||
load_ray_wrappers_with_fake_ray,
|
load_ray_wrappers_with_fake_ray,
|
||||||
load_ray_wrappers_without_ray,
|
load_ray_wrappers_without_ray,
|
||||||
)
|
)
|
||||||
|
from sglang.test.test_utils import CustomTestCase
|
||||||
|
|
||||||
register_cpu_ci(est_time=11, suite="base-a-test-cpu")
|
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())
|
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__":
|
if __name__ == "__main__":
|
||||||
sys.exit(unittest.main())
|
sys.exit(unittest.main())
|
||||||
|
|||||||
Reference in New Issue
Block a user