diff --git a/grpc_servicer/pyproject.toml b/grpc_servicer/pyproject.toml index f07d646a9..c8b873484 100644 --- a/grpc_servicer/pyproject.toml +++ b/grpc_servicer/pyproject.toml @@ -12,6 +12,7 @@ dependencies = [ "grpcio>=1.81.1", "grpcio-reflection>=1.81.1", "grpcio-health-checking>=1.81.1", + "prometheus-client>=0.20.0", ] readme = "README.md" license = "Apache-2.0" diff --git a/grpc_servicer/smg_grpc_servicer/metrics.py b/grpc_servicer/smg_grpc_servicer/metrics.py new file mode 100644 index 000000000..55e83c22a --- /dev/null +++ b/grpc_servicer/smg_grpc_servicer/metrics.py @@ -0,0 +1,290 @@ +"""Shared best-effort Prometheus ``/metrics`` HTTP sidecar for the servicers. + +The gRPC servicers (sglang, tokenspeed) start :func:`start_metrics_sidecar` +alongside their ``serve_grpc`` so the Rust gateway can scrape per-worker +Prometheus metrics over plain HTTP. The sidecar is intentionally dependency-light +(stdlib ``asyncio`` + ``prometheus_client`` only — no aiohttp/uvicorn) and +*best-effort*: any bind/serve failure is logged and swallowed so the gRPC server +keeps serving. + +Metrics come from each engine's native ``prometheus_client`` registry. For +engines without one (tokenspeed), :class:`SchedulerLoadCollector` re-exposes the +same scheduler load snapshot that ``GetLoads`` already computes. +""" + +from __future__ import annotations + +import asyncio +import ipaddress +import logging +import os +from collections.abc import Callable + +from prometheus_client import CollectorRegistry +from prometheus_client.core import GaugeMetricFamily +from prometheus_client.exposition import CONTENT_TYPE_LATEST, generate_latest +from prometheus_client.registry import Collector + +logger = logging.getLogger(__name__) + +# Env fallback for the sidecar port when a servicer's launcher can't take a +# ``--metrics-port`` flag (e.g. sglang's entrypoint comes from upstream). +METRICS_PORT_ENV = "SMG_METRICS_PORT" + +# Wildcard bind hosts: the gateway reaches the worker only at its real address +# (which it discovered from registration), never at a wildcard. +_WILDCARD_HOSTS = frozenset({"0.0.0.0", "::", "[::]", ""}) + + +def _is_unroutable_host(host: str) -> bool: + """True for hosts the gateway can't reach across the network. + + Covers wildcard binds and loopback (``127.0.0.0/8``, ``::1``, ``localhost``): + for these we advertise the numeric ``metrics_port`` only and let the gateway + pair it with the worker address it discovered. Any other hostname is assumed + routable. + """ + if host in _WILDCARD_HOSTS or host.lower() == "localhost": + return True + candidate = host[1:-1] if host.startswith("[") and host.endswith("]") else host + try: + return ipaddress.ip_address(candidate).is_loopback + except ValueError: + return False + + +# One budget for the whole request-head read (request line + headers) so a +# Slowloris-style client dripping headers can't pin a handler open. +_READ_TIMEOUT = 5.0 + + +def _coerce_port(value: object, source: str) -> int | None: + """Coerce ``value`` to a valid ``0 < port < 65536`` int, else log + return None.""" + try: + port = int(value) + except (TypeError, ValueError): + logger.warning("%s=%r is not an int; metrics sidecar disabled", source, value) + return None + if not 0 < port < 65536: + logger.warning("%s=%d out of range; metrics sidecar disabled", source, port) + return None + return port + + +def resolve_metrics_port(explicit: int | None = None) -> int | None: + """Resolve the sidecar port from an explicit value or ``SMG_METRICS_PORT``. + + Returns ``None`` when neither is set (sidecar disabled) or the value is not a + usable port. Both the explicit argument and the env var are range-validated; + a malformed value is logged and treated as unset rather than aborting startup. + """ + if explicit is not None: + return _coerce_port(explicit, "metrics_port") + raw = os.getenv(METRICS_PORT_ENV) + if not raw: + return None + return _coerce_port(raw, METRICS_PORT_ENV) + + +def metrics_url(host: str, port: int) -> str | None: + """Return an ``http://host:port/metrics`` URL, or ``None`` for unroutable binds. + + A servicer bound to a wildcard (``0.0.0.0``) or loopback (``127.0.0.1``, + ``localhost``, ``::1``) address isn't reachable by the gateway across the + network — the gateway reaches it at the worker address it discovered — so + advertising such a URL would be misleading. The numeric ``metrics_port`` is + advertised regardless; the gateway combines it with that worker address. + """ + if _is_unroutable_host(host): + return None + # Bracket bare IPv6 literals so the URL parses (``http://[::1]:9100/...``). + if ":" in host and not host.startswith("["): + host = f"[{host}]" + return f"http://{host}:{port}/metrics" + + +def metrics_server_args(host: str, metrics_port: int | None) -> dict[str, object]: + """Build the ``server_args`` keys advertising the sidecar to the gateway. + + Returns ``{"metrics_port": int}`` (always, when enabled) plus + ``{"metrics_url": str}`` when ``host`` is routable. The Rust gateway picks + these out of ``GetServerInfo.server_args`` via its key allowlist, so the key + names here are a cross-component contract. Empty dict when disabled. + """ + if metrics_port is None: + return {} + port = int(metrics_port) + out: dict[str, object] = {"metrics_port": port} + url = metrics_url(host or "", port) + if url is not None: + out["metrics_url"] = url + return out + + +class SchedulerLoadCollector(Collector): + """Re-expose a scheduler load snapshot as Prometheus gauges. + + ``snapshot_fn`` returns the same dict shape ``GetLoads`` builds its response + from: ``num_running_reqs``, ``num_waiting_reqs``, ``num_total_reqs`` and a + ``token_usage`` ratio in ``[0, 1]``. It is called on every scrape; a raising + ``snapshot_fn`` is caught and exposed as an empty snapshot so one failing call + doesn't break the whole ``/metrics`` response. + """ + + def __init__(self, snapshot_fn: Callable[[], dict[str, float]]): + self._snapshot_fn = snapshot_fn + + def collect(self): + try: + snapshot = self._snapshot_fn() or {} + except Exception: # noqa: BLE001 — a failing snapshot must not break the scrape. + logger.warning("scheduler load snapshot failed; exposing empty", exc_info=True) + snapshot = {} + gauges = ( + ("smg_scheduler_running_requests", "Requests currently running", "num_running_reqs"), + ("smg_scheduler_waiting_requests", "Requests waiting in queue", "num_waiting_reqs"), + ("smg_scheduler_total_requests", "Running + waiting requests", "num_total_reqs"), + ("smg_scheduler_token_usage", "KV-cache token usage ratio [0,1]", "token_usage"), + ) + for name, doc, key in gauges: + metric = GaugeMetricFamily(name, doc) + metric.add_metric([], float(snapshot.get(key, 0.0) or 0.0)) + yield metric + + +class MetricsSidecar: + """A running ``/metrics`` HTTP server backed by ``asyncio``. + + Serves Prometheus exposition for ``GET /metrics`` (and ``/``); every other + path gets ``404``. Kept minimal on purpose — it exists only so the gateway + can scrape, not as a general HTTP frontend. + """ + + def __init__(self, registry: CollectorRegistry, host: str, port: int): + self._registry = registry + self.host = host + self.port = port + self._server: asyncio.AbstractServer | None = None + + async def start(self) -> None: + self._server = await asyncio.start_server(self._handle, self.host, self.port) + # When started with port 0 the OS assigns an ephemeral port; surface it + # so callers (and tests) can read the real bound port. + if self.port == 0 and self._server.sockets: + self.port = self._server.sockets[0].getsockname()[1] + + async def close(self) -> None: + if self._server is not None: + self._server.close() + try: + await self._server.wait_closed() + except Exception: # noqa: BLE001 — shutdown is best-effort. + logger.debug("metrics sidecar wait_closed raised", exc_info=True) + self._server = None + + @staticmethod + async def _read_head(reader: asyncio.StreamReader) -> bytes: + """Read the request line + drain headers, returning the request line. + + Reads until the blank line that ends the header block (or EOF). Draining + the headers keeps the client's write side from seeing a reset. The caller + wraps this in a single ``wait_for`` so the *whole* head-read phase shares + one deadline; a per-line timeout would let a Slowloris client reset the + clock by dripping one header every interval and hold the handler open. + """ + request_line = await reader.readline() + # Empty request line means EOF / immediate disconnect — nothing to serve. + if not request_line: + return b"" + while True: + line = await reader.readline() + if line in (b"\r\n", b"\n", b""): + break + return request_line + + async def _handle( + self, + reader: asyncio.StreamReader, + writer: asyncio.StreamWriter, + ) -> None: + try: + # One deadline for the entire head read (request line + all headers), + # not one per line — see _read_head. + request_line = await asyncio.wait_for(self._read_head(reader), timeout=_READ_TIMEOUT) + # Empty request line means EOF / immediate disconnect — nothing to serve. + if not request_line: + return + + parts = request_line.split() + method = parts[0] if parts else b"" + path = parts[1].split(b"?", 1)[0] if len(parts) > 1 else b"" + + if method != b"GET": + await self._write_response(writer, 405, b"text/plain", b"method not allowed") + elif path in (b"/metrics", b"/"): + payload = generate_latest(self._registry) + await self._write_response(writer, 200, CONTENT_TYPE_LATEST.encode(), payload) + else: + await self._write_response(writer, 404, b"text/plain", b"not found") + except Exception: # noqa: BLE001 — never let a scrape kill the loop. + logger.debug("metrics sidecar request failed", exc_info=True) + finally: + try: + writer.close() + await writer.wait_closed() + except Exception: # noqa: BLE001 — fully release the socket, best-effort. + pass + + @staticmethod + async def _write_response( + writer: asyncio.StreamWriter, + status: int, + content_type: bytes, + body: bytes, + ) -> None: + reason = {200: "OK", 404: "Not Found", 405: "Method Not Allowed"}.get(status, "OK") + head = ( + f"HTTP/1.1 {status} {reason}\r\n".encode() + + b"Content-Type: " + + content_type + + b"\r\n" + + f"Content-Length: {len(body)}\r\n".encode() + + b"Connection: close\r\n\r\n" + ) + writer.write(head + body) + await writer.drain() + + +async def start_metrics_sidecar( + host: str, + port: int, + *, + registry: CollectorRegistry | None = None, +) -> MetricsSidecar | None: + """Best-effort: start the ``/metrics`` sidecar; return it, or ``None`` on failure. + + A bind/serve failure is logged at WARNING and swallowed — the caller's gRPC + server must keep running regardless. ``registry`` defaults to the global + ``prometheus_client`` registry, which is what the engines populate. + """ + if registry is None: + # Import lazily so the module loads even if prometheus_client's default + # registry is unavailable for some reason. + from prometheus_client import REGISTRY + + registry = REGISTRY + + sidecar = MetricsSidecar(registry, host, port) + try: + await sidecar.start() + except Exception: # noqa: BLE001 — sidecar is non-fatal. + logger.warning( + "Failed to start Prometheus /metrics sidecar on %s:%d; continuing without it", + host, + port, + exc_info=True, + ) + return None + + logger.info("Prometheus /metrics sidecar listening on %s:%d", host, sidecar.port) + return sidecar diff --git a/grpc_servicer/smg_grpc_servicer/sglang/server.py b/grpc_servicer/smg_grpc_servicer/sglang/server.py index fa2624719..d233f91b9 100644 --- a/grpc_servicer/smg_grpc_servicer/sglang/server.py +++ b/grpc_servicer/smg_grpc_servicer/sglang/server.py @@ -18,6 +18,7 @@ import grpc from grpc_health.v1 import health_pb2_grpc from grpc_reflection.v1alpha import reflection +from prometheus_client import CollectorRegistry from sglang.srt.configs.model_config import ModelConfig from sglang.srt.disaggregation.utils import FAKE_BOOTSTRAP_HOST, DisaggregationMode from sglang.srt.managers.disagg_service import start_disagg_service @@ -26,6 +27,11 @@ from sglang.utils import get_exception_traceback from smg_grpc_proto import sglang_scheduler_pb2, sglang_scheduler_pb2_grpc +from smg_grpc_servicer.metrics import ( + SchedulerLoadCollector, + resolve_metrics_port, + start_metrics_sidecar, +) from smg_grpc_servicer.sglang.health_servicer import SGLangHealthServicer from smg_grpc_servicer.sglang.request_manager import GrpcRequestManager from smg_grpc_servicer.sglang.scheduler_launcher import launch_scheduler_process_only @@ -38,6 +44,7 @@ async def serve_grpc( server_args: ServerArgs, model_info: dict | None = None, on_request_manager_ready: Callable | None = None, + metrics_port: int | None = None, ): """Start the standalone gRPC server with integrated scheduler. @@ -50,6 +57,8 @@ async def serve_grpc( server starts accepting requests. sglang's HTTP sidecar uses this to wire its admin endpoints to the scheduler, so the callback signature is a public contract. + metrics_port: Optional port for the best-effort Prometheus ``/metrics`` + sidecar (falls back to ``SMG_METRICS_PORT``); ``None`` disables it. """ # Start bootstrap server BEFORE launching scheduler processes (only in PREFILL mode) @@ -155,6 +164,10 @@ async def serve_grpc( ) health_pb2_grpc.add_HealthServicer_to_server(health_servicer, server) + # Requested port only decides whether to start the sidecar; the servicer + # advertises the *bound* port set after a successful start (see below). + metrics_port = resolve_metrics_port(metrics_port) + # Add SGLang service servicer = SGLangSchedulerServicer( request_manager=request_manager, @@ -163,6 +176,8 @@ async def serve_grpc( model_info=model_info, scheduler_info=scheduler_info, health_servicer=health_servicer, + # Set to the bound port only after the sidecar starts (below). + metrics_port=None, ) sglang_scheduler_pb2_grpc.add_SglangSchedulerServicer_to_server(servicer, server) @@ -293,6 +308,20 @@ def _cert_config_fetcher(): await server.start() + # Best-effort Prometheus sidecar. In gRPC mode the scheduler runs in + # subprocesses, so a SchedulerLoadCollector re-exposes the request manager's + # in-process load rather than relying on a populated global registry. + metrics_sidecar = None + if metrics_port is not None: + registry = CollectorRegistry() + registry.register(SchedulerLoadCollector(servicer.load_snapshot)) + metrics_sidecar = await start_metrics_sidecar( + server_args.host, metrics_port, registry=registry + ) + # Advertise only the actually-bound port; leave it None if the bind failed. + if metrics_sidecar is not None: + servicer.metrics_port = metrics_sidecar.port + # Start warmup in a separate thread warmup_thread = threading.Thread( target=_wait_and_warmup_grpc, @@ -320,6 +349,9 @@ def signal_handler(): # requests before we drain. health_servicer.set_not_serving() + if metrics_sidecar is not None: + await metrics_sidecar.close() + # Drain in-flight RPCs with the request manager's ZMQ sockets still # open, then tear it down. Closing ZMQ before server.stop() drops the # backing channel out from under streams that are still draining, so diff --git a/grpc_servicer/smg_grpc_servicer/sglang/servicer.py b/grpc_servicer/smg_grpc_servicer/sglang/servicer.py index 7efd4cb43..37728e8e9 100644 --- a/grpc_servicer/smg_grpc_servicer/sglang/servicer.py +++ b/grpc_servicer/smg_grpc_servicer/sglang/servicer.py @@ -202,6 +202,7 @@ def __init__( model_info: dict, scheduler_info: dict, health_servicer: SGLangHealthServicer | None = None, + metrics_port: int | None = None, ): """Initialize the standalone gRPC service.""" self.request_manager = request_manager @@ -211,6 +212,9 @@ def __init__( self.scheduler_info = scheduler_info self.start_time = time.time() self.health_servicer = health_servicer + # Advertised in GetServerInfo.server_args so the gateway can scrape the + # Prometheus sidecar; ``None`` when the sidecar is disabled. + self.metrics_port = metrics_port self.mm_receiver = None if ( self.server_args.language_only @@ -525,6 +529,29 @@ async def GetModelInfo( num_labels=self.model_info.get("num_labels") or 0, ) + def _metrics_server_args(self) -> dict: + """Sidecar address advertised in ``server_args`` for gateway discovery.""" + from smg_grpc_servicer.metrics import metrics_server_args + + return metrics_server_args(getattr(self.server_args, "host", "") or "", self.metrics_port) + + def load_snapshot(self) -> dict[str, float]: + """Sync, non-blocking load snapshot for the Prometheus sidecar. + + Reads only the request manager's in-process counters (no scheduler + communicator round-trip), so it is safe from a synchronous Prometheus + collect(). ``token_usage`` needs a scheduler round-trip and is reported + via ``GetLoads`` instead, so it is left at 0 here. + """ + info = self.request_manager.get_server_info() + active = int(info.get("active_requests", 0) or 0) + return { + "num_running_reqs": float(active), + "num_waiting_reqs": 0.0, + "num_total_reqs": float(active), + "token_usage": 0.0, + } + async def GetServerInfo( self, _request: sglang_scheduler_pb2.GetServerInfoRequest, @@ -549,6 +576,7 @@ def make_serializable(obj): return str(obj) serializable_args = make_serializable(server_args_dict) + serializable_args.update(self._metrics_server_args()) server_args_struct.update(serializable_args) # Convert scheduler_info to Struct diff --git a/grpc_servicer/smg_grpc_servicer/tokenspeed/__main__.py b/grpc_servicer/smg_grpc_servicer/tokenspeed/__main__.py index 5f20b4484..179b26ba0 100644 --- a/grpc_servicer/smg_grpc_servicer/tokenspeed/__main__.py +++ b/grpc_servicer/smg_grpc_servicer/tokenspeed/__main__.py @@ -6,10 +6,13 @@ All :class:`ServerArgs` flags are accepted — argv is parsed by ``prepare_server_args`` so there is no flag drift vs the HTTP frontend. +``--metrics-port`` (handled here, not by ``ServerArgs``) starts the Prometheus +``/metrics`` sidecar. """ from __future__ import annotations +import argparse import asyncio import logging import sys @@ -33,10 +36,16 @@ def main(argv: list[str] | None = None) -> None: format="%(asctime)s [%(name)s] %(levelname)s %(message)s", ) - server_args = prepare_server_args(argv) + # Pull --metrics-port out before handing the rest to ServerArgs, which + # doesn't know the flag; SMG_METRICS_PORT still applies when it's omitted. + parser = argparse.ArgumentParser(add_help=False) + parser.add_argument("--metrics-port", type=int, default=None) + known, rest = parser.parse_known_args(argv) + + server_args = prepare_server_args(rest) if uvloop is not None: asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) - asyncio.run(serve_grpc(server_args)) + asyncio.run(serve_grpc(server_args, metrics_port=known.metrics_port)) if __name__ == "__main__": diff --git a/grpc_servicer/smg_grpc_servicer/tokenspeed/server.py b/grpc_servicer/smg_grpc_servicer/tokenspeed/server.py index b41bd6007..ab3ec96e7 100644 --- a/grpc_servicer/smg_grpc_servicer/tokenspeed/server.py +++ b/grpc_servicer/smg_grpc_servicer/tokenspeed/server.py @@ -13,10 +13,16 @@ import grpc from grpc_health.v1 import health_pb2_grpc from grpc_reflection.v1alpha import reflection +from prometheus_client import CollectorRegistry from smg_grpc_proto import tokenspeed_scheduler_pb2_grpc from smg_grpc_proto.generated import tokenspeed_scheduler_pb2 from tokenspeed.runtime.utils.server_args import ServerArgs +from smg_grpc_servicer.metrics import ( + SchedulerLoadCollector, + resolve_metrics_port, + start_metrics_sidecar, +) from smg_grpc_servicer.tokenspeed.health_servicer import TokenSpeedHealthServicer from smg_grpc_servicer.tokenspeed.scheduler_launcher import launch_engine from smg_grpc_servicer.tokenspeed.servicer import TokenSpeedSchedulerServicer @@ -75,12 +81,20 @@ def _grpc_server_options(max_message_bytes: int) -> list[tuple[str, int]]: ] -async def serve_grpc(server_args: ServerArgs) -> None: - """Run the TokenSpeed gRPC server until a shutdown signal is received.""" +async def serve_grpc(server_args: ServerArgs, metrics_port: int | None = None) -> None: + """Run the TokenSpeed gRPC server until a shutdown signal is received. + + ``metrics_port`` (falling back to ``SMG_METRICS_PORT``) enables a best-effort + Prometheus ``/metrics`` sidecar the gateway can scrape; ``None`` disables it. + """ logger.info("Launching TokenSpeed scheduler + AsyncLLM...") async_llm, scheduler_info = launch_engine(server_args) + # Requested port only decides whether to start the sidecar; the servicer + # advertises the *bound* port set after a successful start (see below). + metrics_port = resolve_metrics_port(metrics_port) + max_message_bytes = _grpc_max_message_bytes() server = grpc.aio.server( futures.ThreadPoolExecutor(max_workers=10), @@ -98,6 +112,8 @@ async def serve_grpc(server_args: ServerArgs) -> None: server_args=server_args, scheduler_info=scheduler_info, health_servicer=health_servicer, + # Set to the bound port only after the sidecar starts (below). + metrics_port=None, ) tokenspeed_scheduler_pb2_grpc.add_TokenSpeedSchedulerServicer_to_server(servicer, server) @@ -114,6 +130,19 @@ async def serve_grpc(server_args: ServerArgs) -> None: await server.start() + # Best-effort Prometheus sidecar. TokenSpeed has no native registry, so a + # SchedulerLoadCollector re-exposes the in-process load the servicer reports. + metrics_sidecar = None + if metrics_port is not None: + registry = CollectorRegistry() + registry.register(SchedulerLoadCollector(servicer.load_snapshot)) + metrics_sidecar = await start_metrics_sidecar( + server_args.host, metrics_port, registry=registry + ) + # Advertise only the actually-bound port; leave it None if the bind failed. + if metrics_sidecar is not None: + servicer.metrics_port = metrics_sidecar.port + # Warmup on a background thread so the async server can handle the probe. warmup_thread = threading.Thread( target=_wait_and_warmup, @@ -140,6 +169,8 @@ def _signal_handler() -> None: await stop_event.wait() finally: logger.info("Shutting down TokenSpeed gRPC server") + if metrics_sidecar is not None: + await metrics_sidecar.close() try: await servicer.shutdown() except Exception: # noqa: BLE001 diff --git a/grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py b/grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py index a2791c156..ba20e10f6 100644 --- a/grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py +++ b/grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py @@ -121,11 +121,15 @@ def __init__( server_args: ServerArgs, scheduler_info: dict, health_servicer: TokenSpeedHealthServicer | None = None, + metrics_port: int | None = None, ): self.async_llm = async_llm self.server_args = server_args self.scheduler_info = scheduler_info self.health_servicer = health_servicer + # Advertised in GetServerInfo.server_args so the gateway can scrape the + # Prometheus sidecar; ``None`` when the sidecar is disabled. + self.metrics_port = metrics_port self.start_time = time.time() # Resolved ZMQ KV-events endpoint, or None when the worker was not @@ -463,6 +467,7 @@ async def GetServerInfo( server_args_dict = dataclasses.asdict(self.server_args) else: server_args_dict = dict(getattr(self.server_args, "__dict__", {})) + server_args_dict.update(self._metrics_server_args()) server_args_struct = Struct() server_args_struct.update(_make_json_serializable(server_args_dict)) @@ -495,6 +500,37 @@ async def GetServerInfo( max_total_num_tokens=int(self.scheduler_info.get("max_total_num_tokens", 0)), ) + def _metrics_server_args(self) -> dict[str, Any]: + """Sidecar address advertised in ``server_args`` for gateway discovery.""" + from smg_grpc_servicer.metrics import metrics_server_args + + return metrics_server_args(getattr(self.server_args, "host", "") or "", self.metrics_port) + + def load_snapshot(self) -> dict[str, float]: + """Sync, non-blocking load snapshot for the Prometheus sidecar. + + Mirrors the running/waiting/token-usage shape ``GetLoads`` reports, but + reads only in-process ``AsyncLLM`` state (no scheduler ZMQ round-trip) so + it is safe to call from a synchronous Prometheus collect(). + + ``rid_to_state`` retains finished-but-not-yet-cleaned entries and can't + distinguish waiting from running without a scheduler round-trip, so + (like sglang's ``load_snapshot``) waiting is reported as 0 and total as + running only; the scheduler-side breakdown is reported via ``GetLoads``. + """ + rid_to_state = getattr(self.async_llm, "rid_to_state", {}) or {} + running = 0 + for state in rid_to_state.values(): + if getattr(state, "finished", False): + continue + running += 1 + return { + "num_running_reqs": float(running), + "num_waiting_reqs": 0.0, + "num_total_reqs": float(running), + "token_usage": 0.0, + } + # ------------------------------------------------------------------ # GetLoads (unary) — bridges to TokenSpeed's scheduler-side load metrics # ------------------------------------------------------------------ diff --git a/grpc_servicer/tests/test_metrics_sidecar.py b/grpc_servicer/tests/test_metrics_sidecar.py new file mode 100644 index 000000000..c08426b01 --- /dev/null +++ b/grpc_servicer/tests/test_metrics_sidecar.py @@ -0,0 +1,396 @@ +"""Tests for the shared Prometheus ``/metrics`` sidecar (engine-free). + +``smg_grpc_servicer.metrics`` depends only on stdlib + prometheus_client, so +these run without any inference engine installed. They drive the async sidecar +via ``asyncio.run`` so no pytest-asyncio config is required. + +Run with: pytest grpc_servicer/tests/test_metrics_sidecar.py +""" + +from __future__ import annotations + +import asyncio +import urllib.error +import urllib.request + +import pytest + +pytest.importorskip("prometheus_client") + +from prometheus_client import CollectorRegistry # noqa: E402 +from prometheus_client.parser import text_string_to_metric_families # noqa: E402 +from smg_grpc_servicer import metrics as metrics_mod # noqa: E402 +from smg_grpc_servicer.metrics import ( # noqa: E402 + METRICS_PORT_ENV, + SchedulerLoadCollector, + metrics_server_args, + metrics_url, + resolve_metrics_port, + start_metrics_sidecar, +) + + +def _get(url: str) -> tuple[int, str, str]: + with urllib.request.urlopen(url, timeout=5) as resp: # noqa: S310 — loopback only. + return resp.status, resp.headers.get("Content-Type", ""), resp.read().decode() + + +async def _fetch(loop, url: str) -> tuple[int, str, str]: + """Run the blocking urlopen off the event loop so the sidecar can serve it.""" + return await loop.run_in_executor(None, _get, url) + + +# --------------------------------------------------------------------------- +# Pure helpers +# --------------------------------------------------------------------------- + + +def test_metrics_url_skips_wildcard_hosts(): + assert metrics_url("0.0.0.0", 9100) is None + assert metrics_url("::", 9100) is None + assert metrics_url("", 9100) is None + + +def test_metrics_url_skips_loopback_hosts(): + # Loopback isn't reachable from a gateway in another network namespace, so we + # advertise the port only and let the gateway pair it with the worker address. + assert metrics_url("127.0.0.1", 9100) is None + assert metrics_url("127.0.0.5", 9100) is None + assert metrics_url("localhost", 9100) is None + assert metrics_url("LocalHost", 9100) is None + assert metrics_url("::1", 9100) is None + assert metrics_url("[::1]", 9100) is None + + +def test_metrics_url_for_routable_host(): + assert metrics_url("10.1.2.3", 9100) == "http://10.1.2.3:9100/metrics" + + +def test_metrics_url_brackets_ipv6_literal(): + assert metrics_url("2001:db8::1", 9100) == "http://[2001:db8::1]:9100/metrics" + + +def test_metrics_url_does_not_double_bracket_ipv6(): + assert metrics_url("[2001:db8::1]", 9100) == "http://[2001:db8::1]:9100/metrics" + + +def test_metrics_server_args_disabled_is_empty(): + assert metrics_server_args("10.1.2.3", None) == {} + + +def test_metrics_server_args_routable_host_advertises_port_and_url(): + args = metrics_server_args("10.1.2.3", 9100) + assert args == {"metrics_port": 9100, "metrics_url": "http://10.1.2.3:9100/metrics"} + + +def test_metrics_server_args_wildcard_host_advertises_port_only(): + # Gateway combines metrics_port with the worker address it discovered. + assert metrics_server_args("0.0.0.0", 9100) == {"metrics_port": 9100} + + +def test_metrics_server_args_loopback_host_advertises_port_only(): + # Same as wildcard: a loopback bind yields no routable URL, only the port. + assert metrics_server_args("127.0.0.1", 9100) == {"metrics_port": 9100} + assert metrics_server_args("localhost", 9100) == {"metrics_port": 9100} + + +def test_resolve_metrics_port_prefers_explicit(): + assert resolve_metrics_port(9100) == 9100 + + +@pytest.mark.parametrize("value", [0, 70000, -1]) +def test_resolve_metrics_port_rejects_bad_explicit(monkeypatch, value): + # Explicit port is range-validated like the env var (was previously bypassed). + # Set the env to a *valid* port to prove the bad explicit value isn't silently + # falling through to it. + monkeypatch.setenv(METRICS_PORT_ENV, "9300") + assert resolve_metrics_port(value) is None + + +def test_resolve_metrics_port_none_without_env(monkeypatch): + monkeypatch.delenv(METRICS_PORT_ENV, raising=False) + assert resolve_metrics_port(None) is None + + +def test_resolve_metrics_port_reads_env(monkeypatch): + monkeypatch.setenv(METRICS_PORT_ENV, "9200") + assert resolve_metrics_port(None) == 9200 + + +@pytest.mark.parametrize("value", ["abc", "0", "70000", "-1"]) +def test_resolve_metrics_port_rejects_bad_env(monkeypatch, value): + monkeypatch.setenv(METRICS_PORT_ENV, value) + assert resolve_metrics_port(None) is None + + +# --------------------------------------------------------------------------- +# SchedulerLoadCollector +# --------------------------------------------------------------------------- + + +def test_scheduler_load_collector_emits_gauges(): + snapshot = { + "num_running_reqs": 3, + "num_waiting_reqs": 2, + "num_total_reqs": 5, + "token_usage": 0.42, + } + registry = CollectorRegistry() + registry.register(SchedulerLoadCollector(lambda: snapshot)) + + assert registry.get_sample_value("smg_scheduler_running_requests") == 3.0 + assert registry.get_sample_value("smg_scheduler_waiting_requests") == 2.0 + assert registry.get_sample_value("smg_scheduler_total_requests") == 5.0 + assert registry.get_sample_value("smg_scheduler_token_usage") == 0.42 + + +def test_scheduler_load_collector_tolerates_missing_keys(): + registry = CollectorRegistry() + registry.register(SchedulerLoadCollector(dict)) # empty snapshot + assert registry.get_sample_value("smg_scheduler_running_requests") == 0.0 + assert registry.get_sample_value("smg_scheduler_token_usage") == 0.0 + + +def test_scheduler_load_collector_tolerates_raising_snapshot(): + def boom(): + raise RuntimeError("snapshot unavailable") + + registry = CollectorRegistry() + registry.register(SchedulerLoadCollector(boom)) + # A failing snapshot_fn must not break the scrape: gauges fall back to 0. + assert registry.get_sample_value("smg_scheduler_running_requests") == 0.0 + assert registry.get_sample_value("smg_scheduler_total_requests") == 0.0 + + +# --------------------------------------------------------------------------- +# HTTP sidecar +# --------------------------------------------------------------------------- + + +def test_metrics_endpoint_returns_valid_exposition(): + async def scenario(): + registry = CollectorRegistry() + registry.register( + SchedulerLoadCollector( + lambda: { + "num_running_reqs": 7, + "num_waiting_reqs": 1, + "num_total_reqs": 8, + "token_usage": 0.5, + } + ) + ) + sidecar = await start_metrics_sidecar("127.0.0.1", 0, registry=registry) + assert sidecar is not None + loop = asyncio.get_running_loop() + try: + status, content_type, body = await _fetch( + loop, f"http://127.0.0.1:{sidecar.port}/metrics" + ) + finally: + await sidecar.close() + return status, content_type, body + + status, content_type, body = asyncio.run(scenario()) + assert status == 200 + assert content_type.startswith("text/plain") + families = {f.name: f for f in text_string_to_metric_families(body)} + assert "smg_scheduler_running_requests" in families + assert families["smg_scheduler_running_requests"].samples[0].value == 7.0 + + +def test_root_path_also_serves_metrics(): + async def scenario(): + sidecar = await start_metrics_sidecar("127.0.0.1", 0, registry=CollectorRegistry()) + assert sidecar is not None + loop = asyncio.get_running_loop() + try: + return await _fetch(loop, f"http://127.0.0.1:{sidecar.port}/") + finally: + await sidecar.close() + + status, _, body = asyncio.run(scenario()) + assert status == 200 + # Empty registry still yields valid (possibly empty) exposition. + list(text_string_to_metric_families(body)) + + +def test_unknown_path_returns_404(): + async def scenario(): + sidecar = await start_metrics_sidecar("127.0.0.1", 0, registry=CollectorRegistry()) + assert sidecar is not None + loop = asyncio.get_running_loop() + try: + await _fetch(loop, f"http://127.0.0.1:{sidecar.port}/healthz") + finally: + await sidecar.close() + + with pytest.raises(urllib.error.HTTPError) as exc: + asyncio.run(scenario()) + assert exc.value.code == 404 + + +def test_query_string_is_ignored(): + async def scenario(): + sidecar = await start_metrics_sidecar("127.0.0.1", 0, registry=CollectorRegistry()) + assert sidecar is not None + loop = asyncio.get_running_loop() + try: + return await _fetch(loop, f"http://127.0.0.1:{sidecar.port}/metrics?foo=bar") + finally: + await sidecar.close() + + status, _, _ = asyncio.run(scenario()) + assert status == 200 + + +def test_start_is_best_effort_on_bind_failure(): + async def scenario(): + registry = CollectorRegistry() + first = await start_metrics_sidecar("127.0.0.1", 0, registry=registry) + assert first is not None + try: + # Second bind to the same port must fail-soft (return None). + second = await start_metrics_sidecar("127.0.0.1", first.port, registry=registry) + return second + finally: + await first.close() + + assert asyncio.run(scenario()) is None + + +def test_close_is_idempotent(): + async def scenario(): + sidecar = await start_metrics_sidecar("127.0.0.1", 0, registry=CollectorRegistry()) + assert sidecar is not None + await sidecar.close() + await sidecar.close() # second close must not raise + + asyncio.run(scenario()) + + +def test_advertises_actually_bound_port_not_zero(): + # Mirrors server.py: started on port 0, the servicer advertises the *bound* + # port (sidecar.port), never the requested 0. The advertised host is the + # configured server host, independent of the loopback bind used here. + async def scenario(): + sidecar = await start_metrics_sidecar("127.0.0.1", 0, registry=CollectorRegistry()) + assert sidecar is not None + try: + return metrics_server_args("10.0.0.5", sidecar.port) + finally: + await sidecar.close() + + args = asyncio.run(scenario()) + assert args["metrics_port"] != 0 + assert args["metrics_url"] == f"http://10.0.0.5:{args['metrics_port']}/metrics" + + +def test_failed_bind_advertises_nothing(): + # When start returns None, server.py leaves metrics_port None → no advertisement. + async def scenario(): + first = await start_metrics_sidecar("127.0.0.1", 0, registry=CollectorRegistry()) + assert first is not None + try: + second = await start_metrics_sidecar( + "127.0.0.1", first.port, registry=CollectorRegistry() + ) + advertised_port = second.port if second is not None else None + return metrics_server_args("127.0.0.1", advertised_port) + finally: + await first.close() + + assert asyncio.run(scenario()) == {} + + +def test_handle_returns_on_immediate_disconnect(): + # EOF before any request line must not hang the handler; the server keeps + # serving subsequent connections. + async def scenario(): + sidecar = await start_metrics_sidecar("127.0.0.1", 0, registry=CollectorRegistry()) + assert sidecar is not None + loop = asyncio.get_running_loop() + try: + # Open and immediately close without sending anything (EOF). + reader, writer = await asyncio.open_connection("127.0.0.1", sidecar.port) + writer.close() + await writer.wait_closed() + # A normal request still succeeds afterwards. + return await _fetch(loop, f"http://127.0.0.1:{sidecar.port}/metrics") + finally: + await sidecar.close() + + status, _, _ = asyncio.run(scenario()) + assert status == 200 + + +def test_handle_times_out_slow_request(monkeypatch): + # A client that opens a connection but never finishes the request headers + # must be dropped by the read timeout rather than pinning the handler open. + monkeypatch.setattr(metrics_mod, "_READ_TIMEOUT", 0.2) + + async def scenario(): + sidecar = await start_metrics_sidecar("127.0.0.1", 0, registry=CollectorRegistry()) + assert sidecar is not None + try: + reader, writer = await asyncio.open_connection("127.0.0.1", sidecar.port) + # Request line only — never send the blank line that ends headers. + writer.write(b"GET /metrics HTTP/1.1\r\n") + await writer.drain() + # The server's read timeout (0.2s) closes the connection: read() then + # returns EOF well within this wait_for, so the handler didn't hang. + data = await asyncio.wait_for(reader.read(), timeout=3.0) + writer.close() + await writer.wait_closed() + return data + finally: + await sidecar.close() + + assert asyncio.run(scenario()) == b"" + + +def test_handle_times_out_slow_header_drip(monkeypatch): + # Slowloris: a client that drips one header line per interval *shorter* than + # the timeout but never finishes the block. A per-line timeout would reset the + # clock on every drip and hold the handler open forever; the whole-head-phase + # deadline must cut it off regardless of the drip rate. + monkeypatch.setattr(metrics_mod, "_READ_TIMEOUT", 0.3) + + async def scenario(): + sidecar = await start_metrics_sidecar("127.0.0.1", 0, registry=CollectorRegistry()) + assert sidecar is not None + try: + reader, writer = await asyncio.open_connection("127.0.0.1", sidecar.port) + writer.write(b"GET /metrics HTTP/1.1\r\n") + await writer.drain() + + async def drip(): + # Each drip lands well inside _READ_TIMEOUT (0.3s) yet the total + # span (>1s) blows the whole-phase budget; never sends the blank + # line that would end the header block. + for i in range(10): + await asyncio.sleep(0.15) + writer.write(f"X-Drip-{i}: keep-alive\r\n".encode()) + await writer.drain() + + dripper = asyncio.ensure_future(drip()) + try: + # The whole-phase deadline closes the connection: read() returns + # EOF far sooner than the dripper would finish (10 * 0.15 = 1.5s). + data = await asyncio.wait_for(reader.read(), timeout=1.0) + finally: + dripper.cancel() + try: + await dripper + except (asyncio.CancelledError, ConnectionError): + pass + writer.close() + try: + await writer.wait_closed() + except ConnectionError: + pass + return data + finally: + await sidecar.close() + + assert asyncio.run(scenario()) == b"" diff --git a/grpc_servicer/tests/test_servicer_metrics_advertise.py b/grpc_servicer/tests/test_servicer_metrics_advertise.py new file mode 100644 index 000000000..70ef7ffe5 --- /dev/null +++ b/grpc_servicer/tests/test_servicer_metrics_advertise.py @@ -0,0 +1,106 @@ +"""GetServerInfo advertises the metrics sidecar address in ``server_args``. + +These exercise the real servicer ``GetServerInfo`` RPCs, so they require the +engine packages (``tokenspeed`` / ``sglang``) to be importable and are skipped +otherwise. The engine-free advertisement logic itself is covered directly in +``test_metrics_sidecar.py`` via ``metrics_server_args``. + +Run with: pytest grpc_servicer/tests/test_servicer_metrics_advertise.py +""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import pytest + +pytest.importorskip("smg_grpc_proto") + + +def _struct_to_dict(struct) -> dict: + return {k: struct[k] for k in struct.keys()} + + +def test_tokenspeed_get_server_info_advertises_metrics(): + pytest.importorskip("tokenspeed") + from smg_grpc_servicer.tokenspeed.servicer import TokenSpeedSchedulerServicer + + servicer = TokenSpeedSchedulerServicer.__new__(TokenSpeedSchedulerServicer) + # Bypass __init__ (which spins up AsyncLLM loops); set just what + # GetServerInfo touches. + servicer.async_llm = SimpleNamespace(rid_to_state={}) + servicer.server_args = SimpleNamespace(host="10.0.0.5", served_model_name="m") + servicer.scheduler_info = {"max_total_num_tokens": 0} + servicer.health_servicer = None + servicer.metrics_port = 9101 + servicer.start_time = 0.0 + + resp = asyncio.run(servicer.GetServerInfo(object(), object())) + args = _struct_to_dict(resp.server_args) + assert args["metrics_port"] == 9101 + assert args["metrics_url"] == "http://10.0.0.5:9101/metrics" + + +def test_tokenspeed_get_server_info_omits_metrics_when_disabled(): + pytest.importorskip("tokenspeed") + from smg_grpc_servicer.tokenspeed.servicer import TokenSpeedSchedulerServicer + + servicer = TokenSpeedSchedulerServicer.__new__(TokenSpeedSchedulerServicer) + servicer.async_llm = SimpleNamespace(rid_to_state={}) + servicer.server_args = SimpleNamespace(host="10.0.0.5", served_model_name="m") + servicer.scheduler_info = {"max_total_num_tokens": 0} + servicer.health_servicer = None + servicer.metrics_port = None + servicer.start_time = 0.0 + + resp = asyncio.run(servicer.GetServerInfo(object(), object())) + args = _struct_to_dict(resp.server_args) + assert "metrics_port" not in args + assert "metrics_url" not in args + + +def test_tokenspeed_load_snapshot_reports_running_only(): + pytest.importorskip("tokenspeed") + from smg_grpc_servicer.tokenspeed.servicer import TokenSpeedSchedulerServicer + + servicer = TokenSpeedSchedulerServicer.__new__(TokenSpeedSchedulerServicer) + # Two running + one finished-but-not-cleaned entry: waiting can't be told + # apart from running without a scheduler round-trip, so it's reported as 0 + # and total mirrors running (excludes the finished entry). + servicer.async_llm = SimpleNamespace( + rid_to_state={ + "a": SimpleNamespace(finished=False), + "b": SimpleNamespace(finished=False), + "c": SimpleNamespace(finished=True), + } + ) + snap = servicer.load_snapshot() + assert snap["num_running_reqs"] == 2.0 + assert snap["num_waiting_reqs"] == 0.0 + assert snap["num_total_reqs"] == 2.0 + assert snap["token_usage"] == 0.0 + + +def test_sglang_get_server_info_advertises_metrics(): + pytest.importorskip("sglang") + from smg_grpc_servicer.sglang.servicer import SGLangSchedulerServicer + + servicer = SGLangSchedulerServicer.__new__(SGLangSchedulerServicer) + servicer.request_manager = SimpleNamespace( + get_server_info=lambda: { + "active_requests": 0, + "paused": False, + "last_receive_time": 0.0, + } + ) + # Wildcard host: only the port is advertised (gateway knows the real address). + servicer.server_args = SimpleNamespace(host="0.0.0.0") + servicer.scheduler_info = {"max_total_num_tokens": 0} + servicer.metrics_port = 9102 + servicer.start_time = 0.0 + + resp = asyncio.run(servicer.GetServerInfo(object(), object())) + args = _struct_to_dict(resp.server_args) + assert args["metrics_port"] == 9102 + assert "metrics_url" not in args