sgl-project--sglang/python/sglang/test/kits/fwd_occupancy_kit.py

"""Single-batch decode GPU occupancy sanity kit.

Probes ``sglang:fwd_occupancy`` (a 0-100 percentage averaged over the
last ``decode_log_interval`` batches; resets to NaN at window
boundaries) under one long single-batch ``/generate`` request, and
asserts median above a threshold. Single-batch is where CPU overhead
dominates -- overlap scheduler / cuda graph regressions surface here
before batched throughput moves.

Prerequisites on the consuming server:
    env:          SGLANG_ENABLE_METRICS_DEVICE_TIMER=1
    server flag:  --enable-metrics

Mix into a ``CustomTestCase`` subclass exposing ``self.base_url``.
"""

import re
import statistics
import threading
import time

import requests
import tabulate

_FWD_OCCUPANCY_RE = re.compile(
    r"^sglang:fwd_occupancy(?:\{[^}]*\})?\s+(\S+)", re.MULTILINE
)
_GENERATE_REQUEST_TIMEOUT = 600
_METRICS_REQUEST_TIMEOUT = 10


class FwdOccupancyMixin:
    """Assert single-batch ``sglang:fwd_occupancy`` median > threshold."""

    fwd_occupancy_threshold: float = 95.0
    fwd_occupancy_min_samples: int = 5
    fwd_occupancy_scrape_interval: float = 0.5

    # Spec-decoding accept-length floor. Only enforced when the server
    # is running with a spec algorithm (avg_spec_accept_length present
    # in /server_info); silently skipped otherwise. EAGLE3 3/1/4 on
    # 5090 + Llama-3.1-8B measured ~2.0 in CI; 1.8 leaves a small
    # buffer while still catching silent fallback to vanilla (~1.0).
    fwd_occupancy_acc_length_threshold: float = 1.8

    # Warmup: one short request to fill cuda graphs + get the
    # device-timer past its first NaN window.
    fwd_occupancy_warmup_max_new_tokens: int = 64
    fwd_occupancy_warmup_settle_seconds: float = 1.0

    # Measurement: one long single-batch request -- max_new_tokens must
    # span several decode_log_interval windows for enough samples.
    fwd_occupancy_max_new_tokens: int = 2048
    fwd_occupancy_prompt: str = (
        "Human: Give me a fully functional FastAPI server. Show the python code.\n\nAssistant:"
    )

    def _scrape_fwd_occupancy(self):
        """Max non-NaN gauge value across exposed labels (e.g. dp ranks);
        None on transient scrape failure."""
        try:
            resp = requests.get(
                self.base_url + "/metrics", timeout=_METRICS_REQUEST_TIMEOUT
            )
        except requests.RequestException:
            return None
        if resp.status_code != 200:
            return None
        vals = []
        for raw in _FWD_OCCUPANCY_RE.findall(resp.text):
            try:
                v = float(raw)
            except ValueError:
                continue
            if v == v:  # NaN filter (gauge resets to NaN on window boundary)
                vals.append(v)
        return max(vals) if vals else None

    def _assert_metrics_device_timer_enabled(self):
        """Fail loudly on missing flag/env -- otherwise a NaN-only gauge
        looks like a real occupancy regression."""
        resp = requests.get(
            self.base_url + "/metrics", timeout=_METRICS_REQUEST_TIMEOUT
        )
        if resp.status_code != 200:
            raise AssertionError(
                f"/metrics returned {resp.status_code}; the test class's "
                "server must be launched with --enable-metrics"
            )
        if "sglang:fwd_occupancy" not in resp.text:
            raise AssertionError(
                "sglang:fwd_occupancy gauge not exposed; set "
                "SGLANG_ENABLE_METRICS_DEVICE_TIMER=1 in the server's env "
                "and pass --enable-metrics"
            )

    def _fwd_occupancy_fire(self, prompt: str, max_new_tokens: int):
        """Fire one /generate, return (meta_info, wall_time).
        Must not be called concurrently -- that would break the
        single-batch invariant."""
        t0 = time.perf_counter()
        try:
            resp = requests.post(
                self.base_url + "/generate",
                json={
                    "text": prompt,
                    "sampling_params": {
                        "temperature": 0.0,
                        "max_new_tokens": max_new_tokens,
                        "ignore_eos": True,
                    },
                },
                timeout=_GENERATE_REQUEST_TIMEOUT,
            )
        except requests.RequestException:
            # Final stats-vs-threshold is the signal; individual fire
            # failure isn't.
            return {}, 0.0
        elapsed = time.perf_counter() - t0
        try:
            return resp.json().get("meta_info", {}), elapsed
        except ValueError:  # non-JSON body
            return {}, elapsed

    def _fwd_occupancy_warmup(self):
        """Fill cuda graphs + step the device-timer past its first NaN
        window before measurement starts."""
        self._fwd_occupancy_fire(
            "warmup " + self.fwd_occupancy_prompt,
            self.fwd_occupancy_warmup_max_new_tokens,
        )
        time.sleep(self.fwd_occupancy_warmup_settle_seconds)

    def _fwd_occupancy_measure(self):
        """Background-fire one long single-batch request, scrape
        /metrics on the foreground; return (non-NaN samples, perf)."""
        samples = []
        request_done = threading.Event()
        result = {"meta_info": {}, "elapsed": 0.0}

        def fire_one():
            try:
                result["meta_info"], result["elapsed"] = self._fwd_occupancy_fire(
                    self.fwd_occupancy_prompt,
                    self.fwd_occupancy_max_new_tokens,
                )
            finally:
                request_done.set()

        firer = threading.Thread(target=fire_one, daemon=True)
        firer.start()

        while not request_done.is_set():
            v = self._scrape_fwd_occupancy()
            if v is not None:
                samples.append(v)
            time.sleep(self.fwd_occupancy_scrape_interval)

        firer.join(timeout=_GENERATE_REQUEST_TIMEOUT)
        return samples, self._fwd_occupancy_perf(result)

    def _fwd_occupancy_perf(self, result):
        """Aggregate per-request perf metrics from the fire result
        (input/output tokens, decode tps, mean inter-token latency,
        wall-clock tps)."""
        meta = result["meta_info"] or {}
        elapsed = result["elapsed"]
        out = meta.get("completion_tokens", 0) or 0
        decode_tps = meta.get("decode_throughput", 0.0) or 0.0
        return {
            "input_tokens": meta.get("prompt_tokens", 0) or 0,
            "output_tokens": out,
            "decode_tps": decode_tps,
            "mean_itl_ms": (1000.0 / decode_tps) if decode_tps > 0 else 0.0,
            "wall_tps": (out / elapsed) if elapsed > 0 else 0.0,
        }

    def test_fwd_occupancy(self):
        self._assert_metrics_device_timer_enabled()
        self._fwd_occupancy_warmup()
        samples, perf = self._fwd_occupancy_measure()

        # The 2048-token decode above populates the spec running average
        # if a spec algorithm is enabled; absent otherwise (vanilla
        # decode skips this check).
        try:
            info = requests.get(
                self.base_url + "/server_info", timeout=_METRICS_REQUEST_TIMEOUT
            ).json()
            avg_accept = info["internal_states"][0].get("avg_spec_accept_length")
        except (requests.RequestException, KeyError, IndexError):
            avg_accept = None

        # Median is the steady-state signal; peak / p10 included in the
        # assertion message for triage. Both tables print before any
        # assertion so the numbers surface even on assertion failure.
        samples_sorted = sorted(samples)
        if samples_sorted:
            median = statistics.median(samples_sorted)
            peak = samples_sorted[-1]
            p10_idx = min(len(samples_sorted) - 1, max(0, len(samples_sorted) // 10))
            p10 = samples_sorted[p10_idx]
        else:
            median = peak = p10 = float("nan")

        perf_rows = [
            ["input tokens", perf["input_tokens"]],
            ["output tokens", perf["output_tokens"]],
            ["decode tps", f"{perf['decode_tps']:.2f}"],
            ["mean itl (ms)", f"{perf['mean_itl_ms']:.2f}"],
            ["wall tps", f"{perf['wall_tps']:.2f}"],
        ]
        if avg_accept is not None:
            perf_rows.append(["avg spec accept", f"{avg_accept:.3f}"])
        # Lead each table with a text title line so the two tables stay
        # visually separated even when CI prefixes every line with a timestamp
        # (which turns blank separator lines into non-empty lines).
        print(
            "\n[perf metrics]\n"
            + tabulate.tabulate(
                perf_rows, headers=["perf metric", "value"], tablefmt="github"
            )
        )

        print(
            "\n[fwd_occupancy stats]\n"
            + tabulate.tabulate(
                [
                    ["samples (n)", len(samples)],
                    ["median", f"{median:.2f}"],
                    ["peak", f"{peak:.2f}"],
                    ["p10", f"{p10:.2f}"],
                    ["threshold", f"{self.fwd_occupancy_threshold:.2f}"],
                ],
                headers=["fwd_occupancy", "value"],
                tablefmt="github",
            )
        )

        self.assertGreaterEqual(
            len(samples),
            self.fwd_occupancy_min_samples,
            f"only {len(samples)} non-NaN occupancy samples collected "
            f"(need >= {self.fwd_occupancy_min_samples}); the measurement "
            "window may be too short or the gauge stuck at NaN",
        )
        self.assertGreater(
            median,
            self.fwd_occupancy_threshold,
            f"sglang:fwd_occupancy median={median:.2f} did not exceed "
            f"threshold {self.fwd_occupancy_threshold} "
            f"(peak={peak:.2f}, p10={p10:.2f}, n={len(samples)})",
        )
        if avg_accept is not None:
            self.assertGreater(
                avg_accept,
                self.fwd_occupancy_acc_length_threshold,
                f"avg_spec_accept_length={avg_accept:.3f} did not exceed "
                f"threshold {self.fwd_occupancy_acc_length_threshold} -- spec "
                "barely accepted, possibly degraded to vanilla decode",
            )