diff --git a/python/sglang/srt/observability/ray_wrappers.py b/python/sglang/srt/observability/ray_wrappers.py new file mode 100644 index 000000000..deb43ae74 --- /dev/null +++ b/python/sglang/srt/observability/ray_wrappers.py @@ -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 diff --git a/python/sglang/test/observability/__init__.py b/python/sglang/test/observability/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/sglang/test/observability/fake_ray.py b/python/sglang/test/observability/fake_ray.py new file mode 100644 index 000000000..125a5b573 --- /dev/null +++ b/python/sglang/test/observability/fake_ray.py @@ -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) diff --git a/test/registered/observability/test_metrics.py b/test/registered/observability/test_metrics.py index a793da567..34a765d55 100644 --- a/test/registered/observability/test_metrics.py +++ b/test/registered/observability/test_metrics.py @@ -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([]) diff --git a/test/registered/unit/observability/test_ray_wrappers.py b/test/registered/unit/observability/test_ray_wrappers.py new file mode 100644 index 000000000..757372ed2 --- /dev/null +++ b/test/registered/unit/observability/test_ray_wrappers.py @@ -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()) diff --git a/test/registered/unit/observability/test_stat_loggers_di.py b/test/registered/unit/observability/test_stat_loggers_di.py index 271a4cbed..ba00b16f5 100644 --- a/test/registered/unit/observability/test_stat_loggers_di.py +++ b/test/registered/unit/observability/test_stat_loggers_di.py @@ -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