/test/registered/observability/test_metrics.py

import json
import os
import tempfile
import unittest
from typing import Dict, List

import requests
from prometheus_client.parser import text_string_to_metric_families
from prometheus_client.samples import Sample

from sglang.srt.environ import envs
from sglang.srt.observability.metrics_collector import (
    ROUTING_KEY_REQ_COUNT_BUCKET_BOUNDS,
    STAT_LOGGER_ROLE_SCHEDULER,
    SchedulerMetricsCollector,
    compute_routing_key_stats,
)
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.test_utils import (
    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
    DEFAULT_URL_FOR_TEST,
    CustomTestCase,
    is_in_ci,
    popen_launch_server,
)

register_cuda_ci(est_time=74, stage="base-b", runner_config="1-gpu-small")
register_amd_ci(est_time=32, suite="stage-b-test-1-gpu-small-amd")

_MODEL_NAME = "Qwen/Qwen3-0.6B"


class TestEnableMetrics(CustomTestCase):
    def test_metrics_1gpu(self):
        """Test that metrics endpoint returns data when enabled"""
        self._execute_core(
            other_args=[],
            verify_metrics_extra=None,
            expect_mfu_metrics=True,
            enable_mfu_metrics=True,
        )

    def test_mfu_metrics_gate_disabled(self):
        """MFU metrics should not be emitted when the gate is disabled."""
        self._execute_core(
            other_args=[],
            verify_metrics_extra=None,
            expect_mfu_metrics=False,
            enable_mfu_metrics=False,
        )

    def test_metrics_2gpu(self):
        # TODO enable when we have 2-gpu runner in nightly CI
        if is_in_ci():
            print("Skip test_metrics_2gpu since in 1-gpu CI")
            return

        def _verify_metrics_extra(metrics):
            metrics_to_check = [
                (
                    "sglang:dp_cooperation_realtime_tokens_total",
                    {"mode": "prefill_compute"},
                ),
                (
                    "sglang:dp_cooperation_realtime_tokens_total",
                    {"mode": "decode"},
                ),
                (
                    "sglang:dp_cooperation_forward_execution_seconds_total",
                    {"category": "extend"},
                ),
                (
                    "sglang:dp_cooperation_forward_execution_seconds_total",
                    {"category": "decode"},
                ),
            ]
            _check_metrics_positive(self, metrics, metrics_to_check)

            num_prefill_ranks_values = {
                s.labels["num_prefill_ranks"]
                for s in metrics["sglang:dp_cooperation_realtime_tokens_total"]
            }
            self.assertIn("0", num_prefill_ranks_values)
            self.assertIn("1", num_prefill_ranks_values)

        self._execute_core(
            other_args=["--tp", "2", "--dp", "2", "--enable-dp-attention"],
            verify_metrics_extra=_verify_metrics_extra,
            expect_mfu_metrics=True,
            enable_mfu_metrics=True,
        )

    def _execute_core(
        self,
        other_args,
        verify_metrics_extra,
        expect_mfu_metrics: bool,
        enable_mfu_metrics: bool,
    ):
        with (
            envs.SGLANG_ENABLE_METRICS_DP_ATTENTION.override(True),
            envs.SGLANG_ENABLE_METRICS_DEVICE_TIMER.override(True),
            envs.SGLANG_TEST_RETRACT.override(True),
        ):
            launch_args = [
                "--enable-metrics",
                "--cuda-graph-max-bs-decode",
                2,
                *other_args,
            ]
            if enable_mfu_metrics:
                launch_args.insert(1, "--enable-mfu-metrics")
            process = popen_launch_server(
                _MODEL_NAME,
                DEFAULT_URL_FOR_TEST,
                timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
                other_args=launch_args,
            )

        try:
            # Make some requests to generate some metrics
            response = requests.get(f"{DEFAULT_URL_FOR_TEST}/health_generate")
            self.assertEqual(response.status_code, 200)

            response = requests.post(
                f"{DEFAULT_URL_FOR_TEST}/generate",
                json={
                    "text": ["The capital of France is"] * 20,
                    "sampling_params": {
                        "temperature": 0,
                        "max_new_tokens": 50,
                    },
                    "stream": True,
                    "ignore_eos": True,
                },
                stream=True,
            )
            for _ in response.iter_lines(decode_unicode=False):
                pass

            for i in range(2):
                # Send the request twice to trigger cached token metrics
                response = requests.post(
                    f"{DEFAULT_URL_FOR_TEST}/generate",
                    json={
                        "text": "Hello, " * 100,
                        "sampling_params": {"temperature": 0, "max_new_tokens": 5},
                    },
                    headers={"x-smg-routing-key": "test-key"},
                )
                self.assertEqual(response.status_code, 200)

            # Get metrics
            metrics_response = requests.get(f"{DEFAULT_URL_FOR_TEST}/metrics")
            self.assertEqual(metrics_response.status_code, 200)
            metrics_text = metrics_response.text

            print(f"metrics_text=\n{metrics_text}")

            metrics = _parse_prometheus_metrics(metrics_text)
            self._verify_metrics_common(metrics_text, metrics, expect_mfu_metrics)
            if verify_metrics_extra is not None:
                verify_metrics_extra(metrics)
        finally:
            kill_process_tree(process.pid)

    def _verify_metrics_common(self, metrics_text, metrics, expect_mfu_metrics: bool):
        essential_metrics = [
            "sglang:num_running_reqs",
            "sglang:num_used_tokens",
            "sglang:token_usage",
            "sglang:gen_throughput",
            "sglang:num_queue_reqs",
            "sglang:num_grammar_queue_reqs",
            "sglang:cache_hit_rate",
            "sglang:spec_accept_length",
            "sglang:prompt_tokens_total",
            "sglang:generation_tokens_total",
            "sglang:cached_tokens_total",
            "sglang:num_requests_total",
            "sglang:time_to_first_token_seconds",
            "sglang:inter_token_latency_seconds",
            "sglang:e2e_request_latency_seconds",
            "sglang:http_requests_active",
            "sglang:routing_keys_active",
            "sglang:num_unique_running_routing_keys",
            "sglang:routing_key_running_req_count",
            "sglang:routing_key_all_req_count",
        ]
        mfu_metrics = [
            "sglang:estimated_flops_per_gpu_total",
            "sglang:estimated_read_bytes_per_gpu_total",
            "sglang:estimated_write_bytes_per_gpu_total",
        ]
        if expect_mfu_metrics:
            essential_metrics.extend(mfu_metrics)
        for metric in essential_metrics:
            self.assertIn(metric, metrics_text, f"Missing metric: {metric}")

        # Verify routing key GaugeHistogram buckets
        expected_buckets = len(ROUTING_KEY_REQ_COUNT_BUCKET_BOUNDS) + 1
        for metric_name in [
            "sglang:routing_key_running_req_count",
            "sglang:routing_key_all_req_count",
        ]:
            gt_le_pairs = set()
            for sample in metrics.get(metric_name, []):
                gt_le_pairs.add((sample.labels.get("gt"), sample.labels.get("le")))
            self.assertEqual(
                len(gt_le_pairs),
                expected_buckets,
                f"{metric_name}: Expected {expected_buckets} buckets, got {len(gt_le_pairs)}",
            )

        self.assertIn(f'model_name="{_MODEL_NAME}"', metrics_text)
        self.assertIn("_sum{", metrics_text)
        self.assertIn("_count{", metrics_text)
        self.assertIn("_bucket{", metrics_text)

        metrics_to_check = [
            ("sglang:realtime_tokens_total", {"mode": "prefill_compute"}),
            ("sglang:realtime_tokens_total", {"mode": "decode"}),
            ("sglang:forward_execution_seconds_total", {"category": "extend"}),
            ("sglang:forward_execution_seconds_total", {"category": "decode"}),
            ("sglang:process_cpu_seconds_total", {"component": "tokenizer"}),
        ]
        _check_metrics_positive(self, metrics, metrics_to_check)

        if expect_mfu_metrics:
            # Estimated perf metrics may have multiple series (e.g., by rank). Ensure
            # that at least one series for this model has a positive accumulated value.
            for metric_name in mfu_metrics:
                values = [
                    sample.value
                    for sample in metrics.get(metric_name, [])
                    if sample.labels.get("model_name") == _MODEL_NAME
                ]
                self.assertTrue(
                    values, f"{metric_name}: no samples for model {_MODEL_NAME}"
                )
                self.assertGreater(
                    sum(values),
                    0,
                    f"{metric_name}: expected positive total for model {_MODEL_NAME}",
                )
        else:
            # With only --enable-metrics (without --enable-mfu-metrics), MFU
            # counters should not emit positive values.
            for metric_name in mfu_metrics:
                values = [
                    sample.value
                    for sample in metrics.get(metric_name, [])
                    if sample.labels.get("model_name") == _MODEL_NAME
                ]
                if values:
                    self.assertEqual(
                        sum(values),
                        0,
                        f"{metric_name}: expected no positive samples with MFU metrics gate disabled",
                    )


def _parse_prometheus_metrics(metrics_text: str) -> Dict[str, List[Sample]]:
    result = {}
    for family in text_string_to_metric_families(metrics_text):
        for sample in family.samples:
            if sample.name not in result:
                result[sample.name] = []
            result[sample.name].append(sample)
    return result


def _get_sample_value_by_labels(samples: List[Sample], labels: Dict[str, str]) -> float:
    for sample in samples:
        if all(sample.labels.get(k) == v for k, v in labels.items()):
            return sample.value
    raise KeyError(f"No sample found with labels {labels}")


def _check_metrics_positive(test_case, metrics, metrics_to_check):
    for metric_name, labels in metrics_to_check:
        value = _get_sample_value_by_labels(metrics[metric_name], labels)
        test_case.assertGreater(value, 0, f"{metric_name} {labels}")


_DI_MARKER_PATH = "/tmp/sglang_di_test_marker"


class _MarkingSchedulerCollector(SchedulerMetricsCollector):
    """Records its own instantiation to a file so the test can verify the
    custom subclass was used in the scheduler subprocess.

    Defined at module level so it is picklable into the scheduler process.
    Cross-process signalling uses a filesystem marker because the scheduler
    runs in its own subprocess and cannot share in-memory state with the
    test runner.
    """

    def __init__(self, *args, **kwargs):
        with open(_DI_MARKER_PATH, "w") as f:
            f.write("scheduler_collector_initialized\n")
        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:
            pass

    def tearDown(self) -> None:
        try:
            os.unlink(_DI_MARKER_PATH)
        except FileNotFoundError:
            pass

    def test_engine_custom_scheduler_collector(self):
        import sglang as sgl

        engine = sgl.Engine(
            model_path=_MODEL_NAME,
            enable_metrics=True,
            stat_loggers={
                STAT_LOGGER_ROLE_SCHEDULER: _MarkingSchedulerCollector,
            },
        )
        try:
            # One small generation triggers scheduler init, which is where
            # resolve_collector_class() picks the injected subclass.
            engine.generate("Hello", {"max_new_tokens": 4})
        finally:
            engine.shutdown()

        self.assertTrue(
            os.path.exists(_DI_MARKER_PATH),
            "Custom SchedulerMetricsCollector was not instantiated; "
            "stat_loggers DI did not take effect.",
        )


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([])
        self.assertEqual(num_unique, 0)
        self.assertEqual(req_counts, [])

    def test_all_none(self):
        num_unique, req_counts = compute_routing_key_stats([None, None, None])
        self.assertEqual(num_unique, 0)
        self.assertEqual(req_counts, [])

    def test_with_none(self):
        num_unique, req_counts = compute_routing_key_stats([None, "key1", None])
        self.assertEqual(num_unique, 1)
        self.assertEqual(req_counts, [1])

    def test_single_key_multiple_reqs(self):
        num_unique, req_counts = compute_routing_key_stats(["key1"] * 5)
        self.assertEqual(num_unique, 1)
        self.assertEqual(req_counts, [5])

    def test_distribution(self):
        routing_keys = ["key1"] * 5 + ["key2"] * 1 + ["key3"] * 15 + ["key4"] * 250
        num_unique, req_counts = compute_routing_key_stats(routing_keys)
        self.assertEqual(num_unique, 4)
        self.assertEqual(sorted(req_counts), [1, 5, 15, 250])


if __name__ == "__main__":
    unittest.main()