diff --git a/verifiers/v1/cli/serve.py b/verifiers/v1/cli/serve.py index 94038bc7b..7e7ac2674 100644 --- a/verifiers/v1/cli/serve.py +++ b/verifiers/v1/cli/serve.py @@ -63,6 +63,8 @@ def main(argv: list[str] | None = None) -> None: **pool_serve_kwargs(config.pool), legacy=config.is_legacy, address=config.address, + metrics_address=config.metrics_address, + metrics_port=config.metrics_port, log_setup=partial(setup_logging, level), **server_kwargs, ) diff --git a/verifiers/v1/configs/serve.py b/verifiers/v1/configs/serve.py index a71c2ecf8..6be0461ff 100644 --- a/verifiers/v1/configs/serve.py +++ b/verifiers/v1/configs/serve.py @@ -14,3 +14,7 @@ class ServeConfig(EnvServerConfig): """Log at debug level instead of info.""" dry_run: bool = False """Resolve + validate the config and dump it, then exit.""" + metrics_address: str = "127.0.0.1" + """Address for the optional Prometheus HTTP endpoint.""" + metrics_port: int | None = Field(None, ge=1, le=65535) + """Port for the optional Prometheus HTTP endpoint; unset disables metrics.""" diff --git a/verifiers/v1/serve/__init__.py b/verifiers/v1/serve/__init__.py index 01481ed08..1c6abd33d 100644 --- a/verifiers/v1/serve/__init__.py +++ b/verifiers/v1/serve/__init__.py @@ -1,6 +1,7 @@ """Serve V1 environments over ZMQ.""" from verifiers.v1.serve.client import EnvClient +from verifiers.v1.serve.metrics import MetricsServer, PoolMetrics from verifiers.v1.serve.pool import EnvServerPool, env_config_data, serve_env from verifiers.v1.serve.server import EnvServer from verifiers.v1.serve.types import ( @@ -20,6 +21,8 @@ "serve_env", "env_config_data", "EnvClient", + "MetricsServer", + "PoolMetrics", "HealthRequest", "HealthResponse", "InfoRequest", diff --git a/verifiers/v1/serve/metrics.py b/verifiers/v1/serve/metrics.py new file mode 100644 index 000000000..fe7747216 --- /dev/null +++ b/verifiers/v1/serve/metrics.py @@ -0,0 +1,150 @@ +"""Dependency-free Prometheus metrics for the v1 env-server pool.""" + +import asyncio +import contextlib +from collections.abc import Callable + +_REQUEST_TIMEOUT_SECONDS = 5 + + +def _metric_line( + name: str, value: int | float, labels: dict[str, str] | None = None +) -> str: + if labels: + escaped = { + key: value.replace("\\", "\\\\").replace('"', '\\"') + for key, value in labels.items() + } + label_text = ",".join(f'{key}="{value}"' for key, value in escaped.items()) + return f"{name}{{{label_text}}} {value}" + return f"{name} {value}" + + +class PoolMetrics: + """Mutable broker-side metrics state rendered on the broker event loop.""" + + def __init__(self, configured_workers: int | None) -> None: + self.configured_workers = configured_workers + self.request_total = 0 + self.request_latency_seconds_count = 0 + self.request_latency_seconds_sum = 0.0 + self.pending_depth = 0 + self.active_rollouts = 0 + self.workers: list[dict] = [] + + def render(self) -> str: + lines = [ + "# HELP verifiers_v1_env_active_rollouts Active rollout slots in the pool.", + "# TYPE verifiers_v1_env_active_rollouts gauge", + _metric_line("verifiers_v1_env_active_rollouts", self.active_rollouts), + "# HELP verifiers_v1_env_worker_active_rollouts Active rollout slots per worker.", + "# TYPE verifiers_v1_env_worker_active_rollouts gauge", + ] + lines.extend( + _metric_line( + "verifiers_v1_env_worker_active_rollouts", + worker["active"], + {"worker_index": str(worker["index"])}, + ) + for worker in self.workers + ) + lines.extend( + [ + "# HELP verifiers_v1_env_workers Current worker processes.", + "# TYPE verifiers_v1_env_workers gauge", + _metric_line("verifiers_v1_env_workers", len(self.workers)), + "# HELP verifiers_v1_env_configured_workers Configured worker limit; zero means unbounded.", + "# TYPE verifiers_v1_env_configured_workers gauge", + _metric_line( + "verifiers_v1_env_configured_workers", + self.configured_workers or 0, + ), + "# HELP verifiers_v1_env_pending_requests Requests currently awaiting a worker reply.", + "# TYPE verifiers_v1_env_pending_requests gauge", + _metric_line("verifiers_v1_env_pending_requests", self.pending_depth), + "# HELP verifiers_v1_env_requests_total Requests dispatched to workers.", + "# TYPE verifiers_v1_env_requests_total counter", + _metric_line("verifiers_v1_env_requests_total", self.request_total), + "# HELP verifiers_v1_env_request_latency_seconds Request latency from dispatch to worker reply.", + "# TYPE verifiers_v1_env_request_latency_seconds summary", + _metric_line( + "verifiers_v1_env_request_latency_seconds_count", + self.request_latency_seconds_count, + ), + _metric_line( + "verifiers_v1_env_request_latency_seconds_sum", + self.request_latency_seconds_sum, + ), + "", + ] + ) + return "\n".join(lines) + + +class MetricsServer: + """Minimal HTTP server exposing one Prometheus text endpoint.""" + + def __init__( + self, + address: str, + port: int, + render: Callable[[], str], + ) -> None: + self.address = address + self.port = port + self.render = render + self._server: asyncio.Server | None = None + + async def start(self) -> None: + self._server = await asyncio.start_server( + self._handle_connection, self.address, self.port + ) + sockets = self._server.sockets or [] + if sockets: + host, port, *_ = sockets[0].getsockname() + self.address = str(host) + self.port = int(port) + + async def _handle_connection( + self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter + ) -> None: + try: + request_line = await asyncio.wait_for( + reader.readline(), timeout=_REQUEST_TIMEOUT_SECONDS + ) + method, path, *_ = request_line.decode("ascii", "replace").split() + if method != "GET": + status = "405 Method Not Allowed" + body = b"method not allowed\n" + elif path.split("?", 1)[0] != "/metrics": + status = "404 Not Found" + body = b"not found\n" + else: + status = "200 OK" + body = self.render().encode() + headers = ( + f"HTTP/1.1 {status}\r\n" + "Content-Type: text/plain; version=0.0.4; charset=utf-8\r\n" + f"Content-Length: {len(body)}\r\n" + "Connection: close\r\n" + "\r\n" + ).encode() + writer.write(headers + body) + await writer.drain() + except asyncio.TimeoutError: + return + except (ValueError, UnicodeError): + writer.write(b"HTTP/1.1 400 Bad Request\r\nConnection: close\r\n\r\n") + with contextlib.suppress(Exception): + await writer.drain() + finally: + writer.close() + with contextlib.suppress(Exception): + await writer.wait_closed() + + async def close(self) -> None: + if self._server is None: + return + self._server.close() + await self._server.wait_closed() + self._server = None diff --git a/verifiers/v1/serve/pool.py b/verifiers/v1/serve/pool.py index faa8b5f50..67c4b3ce2 100644 --- a/verifiers/v1/serve/pool.py +++ b/verifiers/v1/serve/pool.py @@ -28,6 +28,7 @@ import os import signal import threading +import time import uuid from collections.abc import Callable @@ -37,6 +38,7 @@ from verifiers.v1.env import EnvConfig from verifiers.v1.serve.server import EnvServer +from verifiers.v1.serve.metrics import MetricsServer, PoolMetrics from verifiers.v1.serve.types import HealthResponse, RunGroupRequest logger = logging.getLogger(__name__) @@ -88,6 +90,8 @@ def __init__( log_setup: Callable[[], None] | None = None, multiplex: int = 128, elastic: bool = True, + metrics_address: str = "127.0.0.1", + metrics_port: int | None = None, ) -> None: self.server_kwargs = server_kwargs self.max_workers = max_workers @@ -99,6 +103,10 @@ def __init__( self.workers: list[dict] = [] self._mpctx = mp.get_context("spawn") self._poller: zmq.asyncio.Poller | None = None + self.metrics_address = metrics_address + self.metrics_port = metrics_port + self.metrics = PoolMetrics(max_workers) + self._metrics_server: MetricsServer | None = None self.ctx = zmq.asyncio.Context() self.frontend = self.ctx.socket(zmq.ROUTER) @@ -173,7 +181,7 @@ async def run(self) -> None: # (`max_workers` is a concrete count when elastic is off). for _ in range(1 if self.elastic else (self.max_workers or 1)): self._spawn_worker() - # request_id -> {client_id, worker, rollout_slots} + # request_id -> {client_id, worker, rollout_slots, dispatched_at} pending: dict[bytes, dict] = {} logger.info( "EnvServerPool up: address=%s workers=%d/%s multiplex=%d elastic=%s", @@ -185,6 +193,17 @@ async def run(self) -> None: ) try: in_flight = 0 + self.metrics.workers = self.workers + if self.metrics_port is not None: + self._metrics_server = MetricsServer( + self.metrics_address, self.metrics_port, self.metrics.render + ) + await self._metrics_server.start() + logger.info( + "EnvServerPool metrics: address=%s:%d", + self._metrics_server.address, + self._metrics_server.port, + ) while True: events = dict(await self._poller.poll()) if self.frontend in events: @@ -213,8 +232,13 @@ async def run(self) -> None: "client_id": client_id, "worker": worker, "rollout_slots": rollout_slots, + "dispatched_at": time.monotonic(), } in_flight += rollout_slots + self.metrics.request_total += 1 + self.metrics.pending_depth = len(pending) + self.metrics.active_rollouts = in_flight + self.metrics.workers = self.workers # forward without client_id — the DEALER identity is the worker's # `client_id`; we hold the real one in `pending`. await worker["dealer"].send_multipart( @@ -229,8 +253,15 @@ async def run(self) -> None: entry = pending.pop(request_id.bytes, None) if entry is None: continue + self.metrics.request_latency_seconds_count += 1 + self.metrics.request_latency_seconds_sum += ( + time.monotonic() - entry["dispatched_at"] + ) entry["worker"]["active"] -= entry["rollout_slots"] in_flight -= entry["rollout_slots"] + self.metrics.pending_depth = len(pending) + self.metrics.active_rollouts = in_flight + self.metrics.workers = self.workers with contextlib.suppress(zmq.ZMQError): await self.frontend.send_multipart( [entry["client_id"], request_id, data], copy=False @@ -238,6 +269,9 @@ async def run(self) -> None: except (asyncio.CancelledError, KeyboardInterrupt): pass finally: + if self._metrics_server is not None: + await self._metrics_server.close() + self._metrics_server = None self._shutdown() def _shutdown(self) -> None: @@ -279,6 +313,8 @@ def serve_env( log_setup: Callable[[], None] | None = None, multiplex: int = 128, elastic: bool = True, + metrics_address: str = "127.0.0.1", + metrics_port: int | None = None, **server_kwargs, ) -> None: """Serve one env over ZMQ: a single in-process `EnvServer` when `max_workers <= 1`, @@ -324,11 +360,18 @@ def serve_env( log_setup, multiplex, elastic, + metrics_address, + metrics_port, ) if address_queue is not None: address_queue.put(pool.address) asyncio.run(pool.run()) else: + if metrics_port is not None: + logger.warning( + "metrics endpoint is only available for EnvServerPool; " + "single-server mode is unchanged" + ) from verifiers.v1.legacy import LegacyEnvServer if (