[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
|
||||
@@ -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)
|
||||
@@ -1,4 +1,6 @@
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from typing import Dict, List
|
||||
|
||||
@@ -296,12 +298,125 @@ class _MarkingSchedulerCollector(SchedulerMetricsCollector):
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
|
||||
# Path to the cross-process marker file for the FakeRayMetric-style recording
|
||||
# variant below. Distinct from ``_DI_MARKER_PATH`` so the two scheduler
|
||||
# collector subclasses (instantiation-marker vs. emission-recording) cannot
|
||||
# stomp on each other when both tests run in the same CI shard.
|
||||
_DI_RECORDING_MARKER_PATH = os.path.join(
|
||||
tempfile.gettempdir(), "sglang_stat_loggers_di_marker.jsonl"
|
||||
)
|
||||
|
||||
|
||||
class _FileRecordingMetric:
|
||||
"""Module-level recording metric.
|
||||
|
||||
Mirrors the ``FakeRayMetric`` from
|
||||
``sglang.test.observability.fake_ray`` (records ``(op, value, tags)``
|
||||
triples) but exposes the prometheus_client ``.labels(...).inc/.set/
|
||||
.observe(...)`` shape that ``SchedulerMetricsCollector`` calls into.
|
||||
|
||||
Defined at module level so the scheduler subprocess can unpickle the
|
||||
``_RecordingSchedulerCollector`` reference. Recordings are appended as
|
||||
JSON lines to ``_DI_RECORDING_MARKER_PATH`` so the test runner can read
|
||||
them across the process boundary.
|
||||
"""
|
||||
|
||||
def __init__(self, name="", documentation="", labelnames=(), **kwargs):
|
||||
self.name = name
|
||||
self.documentation = documentation
|
||||
self._labelnames = tuple(labelnames or ())
|
||||
# Sink for in-process introspection. The subprocess uses the file
|
||||
# marker instead, since in-memory state is not visible to the test
|
||||
# runner.
|
||||
self.calls = []
|
||||
|
||||
def labels(self, **kwargs):
|
||||
return _FileRecordingMetricBound(self, dict(kwargs))
|
||||
|
||||
|
||||
class _FileRecordingMetricBound:
|
||||
"""The object returned by ``_FileRecordingMetric.labels(...)``.
|
||||
|
||||
All three terminal verbs append a JSON line to the marker file so the
|
||||
test runner can verify emissions made inside the scheduler subprocess.
|
||||
"""
|
||||
|
||||
def __init__(self, parent: "_FileRecordingMetric", tags: dict):
|
||||
self._parent = parent
|
||||
self._tags = tags
|
||||
|
||||
def _record(self, op: str, value):
|
||||
self._parent.calls.append((op, value, dict(self._tags)))
|
||||
try:
|
||||
with open(_DI_RECORDING_MARKER_PATH, "a") as f:
|
||||
f.write(
|
||||
json.dumps(
|
||||
{
|
||||
"name": self._parent.name,
|
||||
"op": op,
|
||||
"value": value,
|
||||
"tags": self._tags,
|
||||
}
|
||||
)
|
||||
+ "\n"
|
||||
)
|
||||
except OSError:
|
||||
# Marker file is best-effort. Never let a recording failure
|
||||
# disturb the scheduler's hot path.
|
||||
pass
|
||||
|
||||
def inc(self, amount=1):
|
||||
self._record("inc", amount)
|
||||
|
||||
def set(self, value):
|
||||
self._record("set", value)
|
||||
|
||||
def observe(self, value):
|
||||
self._record("observe", value)
|
||||
|
||||
|
||||
class _RecordingSchedulerCollector(SchedulerMetricsCollector):
|
||||
"""A custom ``SchedulerMetricsCollector`` that records every emission to
|
||||
a filesystem marker.
|
||||
|
||||
Achieves both halves of the reviewer's request:
|
||||
|
||||
1. Its mere instantiation proves that ``resolve_collector_class()``
|
||||
picked the injected subclass inside the scheduler subprocess
|
||||
(the marker file exists).
|
||||
2. Each emission lands on the ``_FileRecordingMetric`` double, which
|
||||
writes a JSON line. The test reads the file after shutdown and
|
||||
asserts that a few representative metrics received positive values.
|
||||
|
||||
Defined at module level so the scheduler subprocess can unpickle it.
|
||||
"""
|
||||
|
||||
_counter_cls = _FileRecordingMetric
|
||||
_gauge_cls = _FileRecordingMetric
|
||||
_histogram_cls = _FileRecordingMetric
|
||||
_summary_cls = _FileRecordingMetric
|
||||
|
||||
|
||||
def _clear_sglang_metrics_from_default_registry() -> None:
|
||||
"""Drop any ``sglang:`` metrics left in the process-global prometheus default
|
||||
REGISTRY by a prior in-process Engine boot. Without this, a second in-process
|
||||
``sgl.Engine(enable_metrics=True)`` in the same test process re-registers the
|
||||
same Counters and raises "Duplicated timeseries in CollectorRegistry"."""
|
||||
from prometheus_client import REGISTRY
|
||||
|
||||
for collector in list(getattr(REGISTRY, "_collector_to_names", {})):
|
||||
names = REGISTRY._collector_to_names.get(collector, set())
|
||||
if any(name.startswith("sglang:") for name in names):
|
||||
REGISTRY.unregister(collector)
|
||||
|
||||
|
||||
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:
|
||||
_clear_sglang_metrics_from_default_registry()
|
||||
try:
|
||||
os.unlink(_DI_MARKER_PATH)
|
||||
except FileNotFoundError:
|
||||
@@ -337,6 +452,115 @@ class TestStatLoggersDI(CustomTestCase):
|
||||
)
|
||||
|
||||
|
||||
class TestStatLoggersDIRecording(CustomTestCase):
|
||||
"""Boot a real ``sgl.Engine`` with a custom scheduler collector that
|
||||
swaps the four DI hook classes for a FakeRayMetric-style recording
|
||||
double and verify that emissions land on the double.
|
||||
|
||||
Combines the discriminating power of ``_RecordingSchedulerCollector``
|
||||
(proves the subclass was actually instantiated in the scheduler
|
||||
subprocess) with value recording (proves emissions flow through to the
|
||||
metric instance). Per the reviewer's framing, we pick a few
|
||||
representative metrics rather than enumerate all of them.
|
||||
"""
|
||||
|
||||
def setUp(self) -> None:
|
||||
# Avoid stale PROMETHEUS_MULTIPROC_DIR from prior in-process Engine boots.
|
||||
os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None)
|
||||
_clear_sglang_metrics_from_default_registry()
|
||||
try:
|
||||
os.unlink(_DI_RECORDING_MARKER_PATH)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
|
||||
def tearDown(self) -> None:
|
||||
try:
|
||||
os.unlink(_DI_RECORDING_MARKER_PATH)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
|
||||
def _read_marker(self):
|
||||
"""Return all recorded emissions as a list of dicts.
|
||||
|
||||
Each entry has keys ``name`` (str), ``op`` (one of ``inc``/``set``/
|
||||
``observe``), ``value`` (numeric) and ``tags`` (dict).
|
||||
"""
|
||||
entries = []
|
||||
with open(_DI_RECORDING_MARKER_PATH) as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
entries.append(json.loads(line))
|
||||
return entries
|
||||
|
||||
def test_engine_custom_scheduler_collector_emits_through_fake_metric(self):
|
||||
import sglang as sgl
|
||||
|
||||
engine = sgl.Engine(
|
||||
model_path=_MODEL_NAME,
|
||||
enable_metrics=True,
|
||||
stat_loggers={
|
||||
STAT_LOGGER_ROLE_SCHEDULER: _RecordingSchedulerCollector,
|
||||
},
|
||||
)
|
||||
try:
|
||||
# One small generation triggers scheduler init (which is where
|
||||
# resolve_collector_class picks the injected subclass) and is
|
||||
# enough to produce gauge ``.set()`` emissions on the basic
|
||||
# queue-state metrics.
|
||||
engine.generate("Hello", {"max_new_tokens": 4})
|
||||
finally:
|
||||
engine.shutdown()
|
||||
|
||||
# Discrimination: the marker file exists, proving the custom
|
||||
# subclass was instantiated inside the scheduler subprocess.
|
||||
self.assertTrue(
|
||||
os.path.exists(_DI_RECORDING_MARKER_PATH),
|
||||
"Custom SchedulerMetricsCollector was not instantiated; "
|
||||
"stat_loggers DI did not take effect.",
|
||||
)
|
||||
|
||||
entries = self._read_marker()
|
||||
self.assertGreater(
|
||||
len(entries),
|
||||
0,
|
||||
"Marker file exists but contains no emissions; "
|
||||
"the recording double was not wired through the DI hooks.",
|
||||
)
|
||||
|
||||
# Value verification: pick a few representative metrics and check
|
||||
# that they actually received emissions with sensible shapes. We do
|
||||
# not enumerate all metrics; the reviewer's framing was "just pick
|
||||
# a few".
|
||||
by_name = {}
|
||||
for e in entries:
|
||||
by_name.setdefault(e["name"], []).append(e)
|
||||
|
||||
# 1) num_running_reqs: a Gauge that the scheduler ``.set()``s every
|
||||
# stats tick. After one generation it should have at least one
|
||||
# emission.
|
||||
self.assertIn(
|
||||
"sglang:num_running_reqs",
|
||||
by_name,
|
||||
f"Expected num_running_reqs emissions, saw: {sorted(by_name)[:10]}",
|
||||
)
|
||||
running_ops = {e["op"] for e in by_name["sglang:num_running_reqs"]}
|
||||
self.assertIn("set", running_ops)
|
||||
|
||||
# 2) num_queue_reqs: same shape, different metric. Two metrics from
|
||||
# the same collector firing confirm the DI hook applied uniformly.
|
||||
self.assertIn("sglang:num_queue_reqs", by_name)
|
||||
queue_ops = {e["op"] for e in by_name["sglang:num_queue_reqs"]}
|
||||
self.assertIn("set", queue_ops)
|
||||
|
||||
# 3) Tag propagation: every recorded emission must carry the labels
|
||||
# keys the scheduler installed (model_name, engine_type, ...).
|
||||
any_running = by_name["sglang:num_running_reqs"][0]
|
||||
self.assertIn("model_name", any_running["tags"])
|
||||
self.assertEqual(any_running["tags"]["model_name"], _MODEL_NAME)
|
||||
|
||||
|
||||
class TestComputeRoutingKeyStats(unittest.TestCase):
|
||||
def test_empty(self):
|
||||
num_unique, req_counts = compute_routing_key_stats([])
|
||||
|
||||
@@ -0,0 +1,281 @@
|
||||
"""Unit tests for :mod:`sglang.srt.observability.ray_wrappers`.
|
||||
|
||||
The wrapper module is designed to import cleanly even without Ray installed; we
|
||||
inject a fake ``ray``/``ray.util.metrics``/``ray.serve`` triple into
|
||||
``sys.modules`` before importing the module so the tests run on the CPU CI
|
||||
runners that don't ship Ray. The fakes are shared with the DI integration
|
||||
tests in ``test_stat_loggers_di.py``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
import unittest
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
register_cpu_ci(est_time=3, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestRayWrapperBase(unittest.TestCase):
|
||||
def setUp(self) -> None:
|
||||
self.rw = load_ray_wrappers_with_fake_ray(replica_id="rep-001")
|
||||
|
||||
def tearDown(self) -> None:
|
||||
clear_fake_ray_modules()
|
||||
|
||||
|
||||
class TestNameSanitization(TestRayWrapperBase):
|
||||
def test_replaces_colons_with_underscores(self):
|
||||
sanitized = self.rw.RayPrometheusMetric._get_sanitized_opentelemetry_name(
|
||||
"sglang:num_running_reqs"
|
||||
)
|
||||
self.assertEqual(sanitized, "sglang_num_running_reqs")
|
||||
|
||||
def test_replaces_all_punctuation(self):
|
||||
sanitized = self.rw.RayPrometheusMetric._get_sanitized_opentelemetry_name(
|
||||
"sglang:foo.bar-baz/qux"
|
||||
)
|
||||
self.assertEqual(sanitized, "sglang_foo_bar_baz_qux")
|
||||
|
||||
def test_keeps_already_valid_names(self):
|
||||
sanitized = self.rw.RayPrometheusMetric._get_sanitized_opentelemetry_name(
|
||||
"sglang_foo_bar"
|
||||
)
|
||||
self.assertEqual(sanitized, "sglang_foo_bar")
|
||||
|
||||
|
||||
class TestReplicaIdInjection(TestRayWrapperBase):
|
||||
def test_tag_keys_include_replica_id(self):
|
||||
counter = self.rw.RayCounterWrapper(
|
||||
"sglang:requests", "doc", labelnames=["model_name"]
|
||||
)
|
||||
self.assertEqual(counter.metric._tag_keys, ("model_name", "ReplicaId"))
|
||||
|
||||
def test_emit_uses_replica_id_tag(self):
|
||||
counter = self.rw.RayCounterWrapper(
|
||||
"sglang:requests", "doc", labelnames=["model_name"]
|
||||
)
|
||||
counter.labels(model_name="m1").inc(1)
|
||||
op, value, tags = counter.metric.calls[-1]
|
||||
self.assertEqual(op, "inc")
|
||||
self.assertEqual(value, 1)
|
||||
self.assertEqual(tags, {"model_name": "m1", "ReplicaId": "rep-001"})
|
||||
|
||||
|
||||
class TestCounterWrapper(TestRayWrapperBase):
|
||||
def test_inc_forwards_value_and_tags(self):
|
||||
counter = self.rw.RayCounterWrapper("sglang:requests", "doc", labelnames=["m"])
|
||||
counter.labels(m="x").inc(5)
|
||||
self.assertEqual(
|
||||
counter.metric.calls[-1],
|
||||
("inc", 5, {"m": "x", "ReplicaId": "rep-001"}),
|
||||
)
|
||||
|
||||
def test_inc_zero_is_noop(self):
|
||||
counter = self.rw.RayCounterWrapper("sglang:requests", "doc", labelnames=["m"])
|
||||
counter.labels(m="x").inc(0)
|
||||
# No call recorded — inc(0) should short-circuit.
|
||||
self.assertEqual(counter.metric.calls, [])
|
||||
|
||||
|
||||
class TestGaugeWrapper(TestRayWrapperBase):
|
||||
def test_set_forwards_value_and_tags(self):
|
||||
gauge = self.rw.RayGaugeWrapper("sglang:running", "doc", labelnames=["m"])
|
||||
gauge.labels(m="x").set(12)
|
||||
self.assertEqual(
|
||||
gauge.metric.calls[-1],
|
||||
("set", 12, {"m": "x", "ReplicaId": "rep-001"}),
|
||||
)
|
||||
|
||||
def test_set_to_current_time_uses_set(self):
|
||||
gauge = self.rw.RayGaugeWrapper("sglang:start_time", "doc")
|
||||
gauge.set_to_current_time()
|
||||
op, value, _ = gauge.metric.calls[-1]
|
||||
self.assertEqual(op, "set")
|
||||
self.assertIsInstance(value, float)
|
||||
self.assertGreater(value, 0)
|
||||
|
||||
def test_accepts_multiprocess_mode_for_api_parity(self):
|
||||
# multiprocess_mode is irrelevant under Ray; the wrapper must still
|
||||
# accept the kwarg so existing call sites in metrics_collector.py work.
|
||||
gauge = self.rw.RayGaugeWrapper(
|
||||
"sglang:running", "doc", labelnames=["m"], multiprocess_mode="livesum"
|
||||
)
|
||||
gauge.labels(m="x").set(7)
|
||||
self.assertEqual(gauge.metric.calls[-1][0], "set")
|
||||
|
||||
|
||||
class TestHistogramWrapper(TestRayWrapperBase):
|
||||
def test_observe_forwards_value_and_tags(self):
|
||||
hist = self.rw.RayHistogramWrapper(
|
||||
"sglang:ttft_seconds", "doc", labelnames=["m"], buckets=[0.1, 1.0]
|
||||
)
|
||||
hist.labels(m="x").observe(0.3)
|
||||
self.assertEqual(
|
||||
hist.metric.calls[-1],
|
||||
("observe", 0.3, {"m": "x", "ReplicaId": "rep-001"}),
|
||||
)
|
||||
|
||||
def test_buckets_translate_to_boundaries(self):
|
||||
hist = self.rw.RayHistogramWrapper(
|
||||
"sglang:ttft_seconds", "doc", buckets=[0.1, 0.5, 1.0, 2.0]
|
||||
)
|
||||
self.assertEqual(hist.metric.boundaries, [0.1, 0.5, 1.0, 2.0])
|
||||
|
||||
def test_no_buckets_defaults_to_empty_list(self):
|
||||
hist = self.rw.RayHistogramWrapper("sglang:ttft_seconds", "doc")
|
||||
self.assertEqual(hist.metric.boundaries, [])
|
||||
|
||||
def test_non_positive_boundaries_dropped(self):
|
||||
# Ray.util.metrics rejects boundaries <= 0; sglang's queue_time and a
|
||||
# few other histograms include 0.0 as their lowest bucket. The wrapper
|
||||
# silently filters non-positive entries so engine startup never breaks
|
||||
# when the Ray backend is in use.
|
||||
hist = self.rw.RayHistogramWrapper(
|
||||
"sglang:queue_time_seconds", "doc", buckets=[0.0, 0.001, 1.0]
|
||||
)
|
||||
self.assertEqual(hist.metric.boundaries, [0.001, 1.0])
|
||||
|
||||
def test_negative_boundaries_dropped(self):
|
||||
hist = self.rw.RayHistogramWrapper(
|
||||
"sglang:demo", "doc", buckets=[-1.0, 0.0, 0.5]
|
||||
)
|
||||
self.assertEqual(hist.metric.boundaries, [0.5])
|
||||
|
||||
|
||||
class TestSummaryWrapperFallback(TestRayWrapperBase):
|
||||
def test_observe_uses_default_boundaries(self):
|
||||
summary = self.rw.RaySummaryWrapper("sglang:request_latency", "doc")
|
||||
self.assertEqual(
|
||||
summary.metric.boundaries,
|
||||
self.rw.RaySummaryWrapper.DEFAULT_BOUNDARIES,
|
||||
)
|
||||
|
||||
def test_observe_forwards_value_and_tags(self):
|
||||
summary = self.rw.RaySummaryWrapper(
|
||||
"sglang:request_latency", "doc", labelnames=["m"]
|
||||
)
|
||||
summary.labels(m="x").observe(0.42)
|
||||
self.assertEqual(
|
||||
summary.metric.calls[-1],
|
||||
("observe", 0.42, {"m": "x", "ReplicaId": "rep-001"}),
|
||||
)
|
||||
|
||||
|
||||
class TestLabelsCopyAndGuard(TestRayWrapperBase):
|
||||
def test_labels_returns_copy_not_self(self):
|
||||
counter = self.rw.RayCounterWrapper("sglang:requests", "doc", labelnames=["m"])
|
||||
labeled = counter.labels(m="x")
|
||||
self.assertIsNot(labeled, counter)
|
||||
|
||||
def test_original_remains_unlabeled_after_labels_call(self):
|
||||
counter = self.rw.RayCounterWrapper("sglang:requests", "doc", labelnames=["m"])
|
||||
counter.labels(m="x")
|
||||
self.assertFalse(counter._is_labeled)
|
||||
|
||||
def test_double_labels_raises(self):
|
||||
counter = self.rw.RayCounterWrapper("sglang:requests", "doc", labelnames=["m"])
|
||||
labeled = counter.labels(m="x")
|
||||
with self.assertRaises(ValueError):
|
||||
labeled.labels(m="y")
|
||||
|
||||
def test_concurrent_labels_have_isolated_tags(self):
|
||||
counter = self.rw.RayCounterWrapper("sglang:requests", "doc", labelnames=["m"])
|
||||
a = counter.labels(m="alpha")
|
||||
b = counter.labels(m="beta")
|
||||
self.assertEqual(a._tags["m"], "alpha")
|
||||
self.assertEqual(b._tags["m"], "beta")
|
||||
|
||||
def test_positional_label_args_supported(self):
|
||||
counter = self.rw.RayCounterWrapper(
|
||||
"sglang:requests", "doc", labelnames=["model", "engine"]
|
||||
)
|
||||
counter.labels("m1", "e1").inc(1)
|
||||
_, _, tags = counter.metric.calls[-1]
|
||||
self.assertEqual(tags["model"], "m1")
|
||||
self.assertEqual(tags["engine"], "e1")
|
||||
|
||||
def test_wrong_positional_arity_raises(self):
|
||||
counter = self.rw.RayCounterWrapper(
|
||||
"sglang:requests", "doc", labelnames=["model", "engine"]
|
||||
)
|
||||
with self.assertRaises(ValueError):
|
||||
counter.labels("only_one_arg")
|
||||
|
||||
|
||||
class TestCollectorSubclassWiring(TestRayWrapperBase):
|
||||
"""Each Ray collector subclass must override only the ``_xxx_cls`` attrs the
|
||||
underlying collector actually uses."""
|
||||
|
||||
def test_scheduler_overrides_all_four(self):
|
||||
cls = self.rw.RaySchedulerMetricsCollector
|
||||
self.assertIs(cls._counter_cls, self.rw.RayCounterWrapper)
|
||||
self.assertIs(cls._gauge_cls, self.rw.RayGaugeWrapper)
|
||||
self.assertIs(cls._histogram_cls, self.rw.RayHistogramWrapper)
|
||||
self.assertIs(cls._summary_cls, self.rw.RaySummaryWrapper)
|
||||
|
||||
def test_tokenizer_overrides_counter_histogram_only(self):
|
||||
cls = self.rw.RayTokenizerMetricsCollector
|
||||
self.assertIs(cls._counter_cls, self.rw.RayCounterWrapper)
|
||||
self.assertIs(cls._histogram_cls, self.rw.RayHistogramWrapper)
|
||||
self.assertIsNone(cls._gauge_cls)
|
||||
self.assertIsNone(cls._summary_cls)
|
||||
|
||||
def test_storage_overrides_counter_histogram_only(self):
|
||||
cls = self.rw.RayStorageMetricsCollector
|
||||
self.assertIs(cls._counter_cls, self.rw.RayCounterWrapper)
|
||||
self.assertIs(cls._histogram_cls, self.rw.RayHistogramWrapper)
|
||||
self.assertIsNone(cls._gauge_cls)
|
||||
self.assertIsNone(cls._summary_cls)
|
||||
|
||||
def test_radix_overrides_counter_histogram_only(self):
|
||||
cls = self.rw.RayRadixCacheMetricsCollector
|
||||
self.assertIs(cls._counter_cls, self.rw.RayCounterWrapper)
|
||||
self.assertIs(cls._histogram_cls, self.rw.RayHistogramWrapper)
|
||||
self.assertIsNone(cls._gauge_cls)
|
||||
self.assertIsNone(cls._summary_cls)
|
||||
|
||||
def test_expert_dispatch_overrides_histogram_only(self):
|
||||
cls = self.rw.RayExpertDispatchCollector
|
||||
self.assertIs(cls._histogram_cls, self.rw.RayHistogramWrapper)
|
||||
self.assertIsNone(cls._counter_cls)
|
||||
self.assertIsNone(cls._gauge_cls)
|
||||
self.assertIsNone(cls._summary_cls)
|
||||
|
||||
|
||||
class TestRayMissingImportError(unittest.TestCase):
|
||||
"""Importing the module must succeed even without Ray; instantiating a
|
||||
wrapper without Ray must raise ImportError with a clear message."""
|
||||
|
||||
def setUp(self) -> None:
|
||||
self.rw = load_ray_wrappers_without_ray()
|
||||
|
||||
def tearDown(self) -> None:
|
||||
clear_fake_ray_modules()
|
||||
|
||||
def test_module_imports_without_ray(self):
|
||||
# If we got here without exception, the import succeeded.
|
||||
self.assertTrue(hasattr(self.rw, "RayCounterWrapper"))
|
||||
|
||||
def test_instantiating_wrapper_without_ray_raises(self):
|
||||
with self.assertRaises(ImportError) as ctx:
|
||||
self.rw.RayCounterWrapper("sglang:foo", "doc")
|
||||
self.assertIn("Ray", str(ctx.exception))
|
||||
|
||||
def test_get_replica_id_returns_none_without_ray(self):
|
||||
self.assertIsNone(self.rw._get_replica_id())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(unittest.main())
|
||||
@@ -1,5 +1,22 @@
|
||||
"""Unit tests for class-level DI on the five *MetricsCollector classes via
|
||||
ServerArgs.stat_loggers — no server, no model loading."""
|
||||
"""Pure-CPU unit tests for ``ServerArgs.stat_loggers`` DI plumbing.
|
||||
|
||||
These tests cover the small, in-process pieces of the ``stat_loggers``
|
||||
dependency injection feature:
|
||||
|
||||
* The four DI hook class attributes (``_counter_cls``/``_gauge_cls``/
|
||||
``_histogram_cls``/``_summary_cls``) default to ``None`` on every
|
||||
collector, so the existing prometheus_client backend is used unchanged.
|
||||
* ``resolve_collector_class()`` returns the registered subclass when a role
|
||||
is present in ``stat_loggers`` and falls back to the default otherwise.
|
||||
* Without any subclass override, collectors instantiate the real
|
||||
prometheus_client classes.
|
||||
|
||||
The full Engine-level integration test (which boots ``sgl.Engine`` and
|
||||
verifies that emissions land on a FakeRayMetric-style recording double in
|
||||
the scheduler subprocess) lives in
|
||||
``test/registered/observability/test_metrics.py`` alongside the other
|
||||
GPU-backed metrics tests.
|
||||
"""
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
@@ -25,18 +42,19 @@ from sglang.srt.observability.metrics_collector import (
|
||||
|
||||
|
||||
class _StubArgs:
|
||||
"""Minimal ServerArgs stand-in. Avoids triggering heavy ServerArgs import chain."""
|
||||
"""Minimal ServerArgs stand-in.
|
||||
|
||||
Avoids triggering the heavy real ServerArgs import chain for unit-level
|
||||
``resolve_collector_class`` cases.
|
||||
"""
|
||||
|
||||
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."""
|
||||
"""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)
|
||||
@@ -60,9 +78,6 @@ class TestCollectorClassAttrs(unittest.TestCase):
|
||||
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)
|
||||
@@ -81,7 +96,6 @@ class TestResolveCollectorClass(unittest.TestCase):
|
||||
self.assertIs(cls, SchedulerMetricsCollector)
|
||||
|
||||
def test_returns_default_when_role_missing(self):
|
||||
# Different role registered. Default still wins for "scheduler".
|
||||
class MyTokenizer(TokenizerMetricsCollector):
|
||||
pass
|
||||
|
||||
@@ -113,84 +127,13 @@ class TestResolveCollectorClass(unittest.TestCase):
|
||||
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)
|
||||
class TestDefaultBackend(unittest.TestCase):
|
||||
"""Without any subclass override, collectors instantiate the real
|
||||
prometheus_client classes; the existing behavior is unchanged."""
|
||||
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user