gatewaysdk 0.3.2__py3-none-any.whl

gatewaysdk/cli/prometheus.py

@@ -0,0 +1,115 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Serve Prometheus metrics from the GatewaySDK multiprocess registry."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import logging
+import os
+from pathlib import Path
+from typing import Iterable
+
+from fastapi import FastAPI
+from prometheus_client import make_asgi_app  # pyright: ignore[reportUnknownVariableType]
+
+from gatewaysdk.logging import setup as setup_logging
+from gatewaysdk.utils.metrics import get_prometheus_registry
+from gatewaysdk.utils.server_launcher import PythonServerLauncher, PythonServerLauncherArgs
+
+logger = logging.getLogger(__name__)
+
+
+def ensure_prometheus_dir() -> str:
+    """Ensure PROMETHEUS_MULTIPROC_DIR is set and the directory exists."""
+
+    directory = os.getenv("PROMETHEUS_MULTIPROC_DIR")
+    if directory is None:
+        raise ValueError("PROMETHEUS_MULTIPROC_DIR is not set.")
+
+    Path(directory).mkdir(parents=True, exist_ok=True)
+    logger.info("Serving Prometheus multiprocess metrics from %s", directory)
+    return directory
+
+
+def create_prometheus_app(metrics_path: str = "/v1/prometheus") -> FastAPI:
+    """Create a FastAPI app that exposes Prometheus metrics and a health endpoint.
+
+    Args:
+        metrics_path: URL path to expose the Prometheus metrics endpoint on.
+
+    Returns:
+        A FastAPI application ready to serve metrics.
+    """
+
+    if not metrics_path.startswith("/"):
+        raise ValueError("metrics_path must start with '/'.")
+
+    normalized_path = metrics_path.rstrip("/")
+    if normalized_path in ("", "/"):
+        raise ValueError("metrics_path must not be '/'. Choose a sub-path such as /v1/prometheus.")
+
+    app = FastAPI(title="GatewaySDK Prometheus exporter", docs_url=None, redoc_url=None)
+    metrics_app = make_asgi_app(registry=get_prometheus_registry())  # pyright: ignore[reportUnknownVariableType]
+    app.mount(normalized_path, metrics_app)  # pyright: ignore[reportUnknownArgumentType]
+
+    @app.get("/health")
+    async def healthcheck() -> dict[str, str]:  # pyright: ignore[reportUnusedFunction]
+        return {"status": "ok"}
+
+    return app
+
+
+def main(argv: Iterable[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description="Serve Prometheus metrics outside the LightningStore server.")
+    parser.add_argument("--host", default="0.0.0.0", help="Host to bind the metrics server to.")
+    parser.add_argument("--port", type=int, default=4748, help="Port to expose the Prometheus metrics on.")
+    parser.add_argument(
+        "--metrics-path",
+        default="/v1/prometheus",
+        help="HTTP path used to expose metrics. Must start with '/' and not be the root path.",
+    )
+    parser.add_argument(
+        "--log-level",
+        default="INFO",
+        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
+        help="Configure the logging level for the metrics server.",
+    )
+    parser.add_argument(
+        "--access-log",
+        action="store_true",
+        help="Enable uvicorn access logs. Disabled by default to reduce noise.",
+    )
+    args = parser.parse_args(list(argv) if argv is not None else None)
+
+    setup_logging(args.log_level)
+    ensure_prometheus_dir()
+
+    try:
+        app = create_prometheus_app(args.metrics_path)
+    except ValueError as exc:
+        logger.error("Failed to configure prometheus app: %s", exc)
+        return 1
+
+    launcher_args = PythonServerLauncherArgs(
+        host=args.host,
+        port=args.port,
+        log_level=getattr(logging, args.log_level),
+        access_log=args.access_log,
+        healthcheck_url="/health",
+    )
+    launcher = PythonServerLauncher(app, launcher_args)
+
+    try:
+        asyncio.run(launcher.run_forever())
+    except KeyboardInterrupt:
+        logger.info("Received shutdown signal. Stopping Prometheus server.")
+    except RuntimeError as exc:
+        logger.error("Prometheus server failed to start: %s", exc, exc_info=True)
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

gatewaysdk/cli/release_gate.py

@@ -0,0 +1,215 @@
+# Copyright (c) Gateway. All rights reserved.
+
+"""`gateway release gate` — ask the release gate about a build, from CI.
+
+The platform applies one rule before a build is promoted (pass rate over
+graded rollouts, plus a floor on how many were graded). That rule lives on
+the server; this command only asks for the verdict and turns it into an exit
+code a pipeline can branch on:
+
+    0  pass      the gate is open on the build's latest run
+    1  blocked   the latest run failed the rule; the reasons say how
+    2  untested  nothing to judge yet: no run recorded this build, or its
+                 latest run is still going
+
+`--wait` polls until every run attributed to the build is terminal (or the
+timeout lapses), so a CI step can start a run and then block on the verdict.
+The verdict itself is never computed here: a second copy of the rule in the
+CLI is exactly the drift the server endpoint exists to prevent.
+
+Kept apart from ``platform.py`` so tests can drive it with a fake HTTP
+client and no typer app.
+"""
+
+from __future__ import annotations
+
+import json
+import time
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+EXIT_PASS = 0
+EXIT_BLOCKED = 1
+EXIT_UNTESTED = 2
+
+EXIT_CODES: Dict[str, int] = {
+    "pass": EXIT_PASS,
+    "blocked": EXIT_BLOCKED,
+    "untested": EXIT_UNTESTED,
+}
+
+# Statuses after which a run grades nothing further. Mirrors the server's
+# `isTerminalRunStatus`; `--wait` stops when every listed run is one of these.
+TERMINAL_STATES = frozenset({"COMPLETED", "FAILED", "ERROR", "TIMEOUT", "CANCELLED"})
+
+GATE_PATH = "/api/public/agent-builds/gate"
+
+
+class GateLookupError(ValueError):
+    """The caller named the build zero ways or more than one way."""
+
+
+def build_query(
+    *,
+    build: Optional[str],
+    agent: Optional[str],
+    version: Optional[str],
+    commit: Optional[str],
+) -> Dict[str, str]:
+    """Exactly one way of naming the build, as the endpoint's query string."""
+    ways = [build is not None, agent is not None or version is not None, commit is not None]
+    if sum(ways) != 1:
+        raise GateLookupError("name the build one way: --build, --agent + --version, or --commit")
+    if build is not None:
+        return {"buildId": build}
+    if commit is not None:
+        return {"commitSha": commit}
+    if not agent or not version:
+        raise GateLookupError("--agent and --version go together; supply both")
+    return {"agentId": agent, "version": version}
+
+
+def fetch_gate(client: Any, query: Dict[str, str]) -> Dict[str, Any]:
+    """One GET against the gate endpoint. Raises on any non-2xx status."""
+    response = client.get(GATE_PATH, params=query)
+    if response.status_code == 404:
+        detail = _error_detail(response)
+        raise LookupError(detail or "no build matches that lookup")
+    response.raise_for_status()
+    return response.json()
+
+
+def _error_detail(response: Any) -> Optional[str]:
+    try:
+        body = response.json()
+    except Exception:  # noqa: BLE001 - a non-JSON error body is still an error
+        return None
+    if isinstance(body, dict):
+        for key in ("message", "error"):
+            if isinstance(body.get(key), str):
+                return body[key]
+    return None
+
+
+def runs_settled(result: Dict[str, Any]) -> bool:
+    """True once the build has runs and none of them can still change."""
+    runs = result.get("evaluatedRuns") or []
+    if not runs:
+        return False
+    return all(str(run.get("state", "")).upper() in TERMINAL_STATES for run in runs)
+
+
+def wait_for_gate(
+    client: Any,
+    query: Dict[str, str],
+    *,
+    timeout_s: float,
+    interval_s: float,
+    clock: Callable[[], float] = time.monotonic,
+    sleep: Callable[[float], None] = time.sleep,
+    on_poll: Optional[Callable[[Dict[str, Any]], None]] = None,
+) -> Tuple[Dict[str, Any], bool]:
+    """Poll until the attributed runs are terminal or the timeout lapses.
+
+    Returns the last result and whether the runs settled. A timeout is not an
+    error: the caller gets the latest verdict (usually ``untested``) and the
+    honest exit code for it.
+    """
+    deadline = clock() + timeout_s
+    while True:
+        result = fetch_gate(client, query)
+        if on_poll is not None:
+            on_poll(result)
+        if runs_settled(result):
+            return result, True
+        if clock() >= deadline:
+            return result, False
+        sleep(interval_s)
+
+
+def exit_code_for(result: Dict[str, Any]) -> int:
+    verdict = str(result.get("gate", {}).get("verdict", "")).lower()
+    if verdict not in EXIT_CODES:
+        raise ValueError(f"unknown gate verdict {verdict!r}")
+    return EXIT_CODES[verdict]
+
+
+def _fmt_rate(rate: Optional[float]) -> str:
+    return "n/a" if rate is None else f"{rate:.2f}"
+
+
+def render_table(result: Dict[str, Any]) -> str:
+    """The verdict as a short table, then the reasons and the judged runs."""
+    build = result.get("build", {})
+    gate = result.get("gate", {})
+    compared = result.get("comparedTo")
+
+    rows: List[Tuple[str, str]] = [
+        ("verdict", str(gate.get("verdict", "?")).upper()),
+        ("agent", str(build.get("agentId", "?"))),
+        ("version", str(build.get("version") or "-")),
+        ("build", str(build.get("id", "?"))),
+        ("commit", str(build.get("commitSha") or "-")),
+        ("pass rate", _fmt_rate(gate.get("passRate"))),
+        ("graded", str(gate.get("gradedRollouts", 0))),
+        (
+            "rule",
+            f"needs {_fmt_rate(gate.get('threshold'))} pass rate "
+            f"and at least {gate.get('minGradedRollouts', '?')} graded",
+        ),
+    ]
+    if compared:
+        stable = compared.get("stableVersion") or compared.get("stableBuildId")
+        rows.append(("production", str(stable)))
+
+    width = max(len(k) for k, _ in rows)
+    lines = [f"{k.ljust(width)}  {v}" for k, v in rows]
+
+    reasons = gate.get("reasons") or []
+    if reasons:
+        lines.append("")
+        lines.extend(f"- {reason}" for reason in reasons)
+
+    runs = result.get("evaluatedRuns") or []
+    if runs:
+        lines.append("")
+        lines.append("runs (newest first):")
+        for run in runs:
+            lines.append(
+                f"  {run.get('runId')}  {run.get('state')}  "
+                f"world={run.get('world') or '-'}  model={run.get('model') or '-'}  "
+                f"pass={_fmt_rate(run.get('passRate'))}  graded={run.get('graded', 0)}"
+            )
+    return "\n".join(lines)
+
+
+def render_json(result: Dict[str, Any]) -> str:
+    return json.dumps(result, indent=2, sort_keys=True)
+
+
+def write_step_summary(result: Dict[str, Any], path: Optional[str]) -> None:
+    """Append the verdict to a GitHub Actions step summary, when one exists.
+
+    A nicety, never the contract: the exit code is what a pipeline acts on,
+    so a summary that cannot be written is ignored rather than fatal.
+    """
+    if not path:
+        return
+    build = result.get("build", {})
+    gate = result.get("gate", {})
+    lines = [
+        f"### Release gate: {build.get('agentId', '?')} {build.get('version') or ''}".rstrip(),
+        "",
+        "| | |",
+        "|---|---|",
+        f"| verdict | {str(gate.get('verdict', '?')).upper()} |",
+        f"| pass rate | {_fmt_rate(gate.get('passRate'))} |",
+        f"| graded | {gate.get('gradedRollouts', 0)} |",
+        f"| rule | {_fmt_rate(gate.get('threshold'))} / {gate.get('minGradedRollouts', '?')} graded |",
+    ]
+    for reason in gate.get("reasons") or []:
+        lines.append(f"| reason | {reason} |")
+    try:
+        with open(path, "a", encoding="utf-8") as fh:
+            fh.write("\n".join(lines) + "\n\n")
+    except OSError:
+        pass

gatewaysdk/cli/store.py

@@ -0,0 +1,131 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Run a LightningStore server for persistent access from multiple processes."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import logging
+from typing import Iterable, List
+
+from gatewaysdk import setup_logging
+from gatewaysdk.store.client_server import LightningStoreServer
+from gatewaysdk.store.memory import InMemoryLightningStore
+from gatewaysdk.utils.metrics import (
+    ConsoleMetricsBackend,
+    MetricsBackend,
+    MultiMetricsBackend,
+    PrometheusMetricsBackend,
+    setup_multiprocess_prometheus,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def main(argv: Iterable[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description="Run a LightningStore server")
+    parser.add_argument("--host", default="0.0.0.0", help="Host to bind the server to")
+    parser.add_argument("--port", type=int, default=4747, help="Port to run the server on")
+    parser.add_argument(
+        "--cors-origin",
+        dest="cors_origins",
+        action="append",
+        help="Allowed CORS origin. Repeat for multiple origins. Use '*' to allow all origins.",
+    )
+    parser.add_argument(
+        "--log-level",
+        default="INFO",
+        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
+        help="Configure the logging level for the store.",
+    )
+    parser.add_argument(
+        "--tracker",
+        nargs="+",
+        choices=["prometheus", "console"],
+        help="Enable metrics tracking. Repeat for multiple trackers.",
+    )
+    parser.add_argument(
+        "--n-workers",
+        default=1,
+        type=int,
+        help=(
+            "Number of workers to run in the server. When it's greater than 1, the server will be run using `mp` launch mode. "
+            "Only applicable for zero-copy stores such as MongoDB backend."
+        ),
+    )
+
+    parser.add_argument(
+        "--backend",
+        choices=["memory", "mongo"],
+        default="memory",
+        help="Backend to use for the store.",
+    )
+    parser.add_argument(
+        "--mongo-uri",
+        default="mongodb://localhost:27017/?replicaSet=rs0",
+        help="MongoDB URI to use for the store. Applicable only if --backend is 'mongo'.",
+    )
+
+    args = parser.parse_args(list(argv) if argv is not None else None)
+
+    setup_logging(args.log_level)
+
+    trackers: List[MetricsBackend] = []
+    if args.tracker:
+        if "prometheus" in args.tracker:
+            logger.info("Enabling Prometheus metrics tracking.")
+            if args.n_workers > 1:
+                # This has to be done before prometheus_client is imported
+                setup_multiprocess_prometheus()
+                logger.info("Setting up Prometheus multiprocess directory for metrics tracking.")
+            trackers.append(PrometheusMetricsBackend())
+
+        if "console" in args.tracker:
+            logger.info("Enabling console metrics tracking.")
+            trackers.append(ConsoleMetricsBackend())
+
+    if len(trackers) == 0:
+        tracker: MetricsBackend | None = None
+    elif len(trackers) == 1:
+        tracker = trackers[0]
+    else:
+        tracker = MultiMetricsBackend(trackers)
+
+    if args.backend == "memory":
+        store = InMemoryLightningStore(
+            thread_safe=True,  # Using thread_safe store for server
+            tracker=tracker,
+        )
+    elif args.backend == "mongo":
+        from gatewaysdk.store.mongo import MongoLightningStore
+
+        store = MongoLightningStore(mongo_uri=args.mongo_uri, tracker=tracker)
+    else:
+        raise ValueError(f"Invalid backend: {args.backend}")
+
+    if args.n_workers > 1:
+        logger.info(f"Running the server using `mp` launch mode with {args.n_workers} workers.")
+        launch_mode = "mp"
+    else:
+        logger.info("Running the server using `asyncio` launch mode.")
+        launch_mode = "asyncio"
+    server = LightningStoreServer(
+        store,
+        host=args.host,
+        port=args.port,
+        cors_allow_origins=args.cors_origins,
+        launch_mode=launch_mode,
+        tracker=tracker,
+        n_workers=args.n_workers,
+    )
+    try:
+        asyncio.run(server.run_forever())
+    except RuntimeError as exc:
+        logger.error("LightningStore server failed to start: %s", exc, exc_info=True)
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

gatewaysdk/cli/vllm.py

@@ -0,0 +1,29 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from typing import Iterable
+
+
+def main(argv: Iterable[str] | None = None) -> int:
+    import sys
+
+    from vllm.entrypoints.cli.main import main as vllm_main
+
+    from gatewaysdk.instrumentation.vllm import instrument_vllm
+
+    instrument_vllm()
+    if argv is not None:
+        original_argv = sys.argv
+        sys.argv = [original_argv[0], *list(argv)]
+        try:
+            vllm_main()
+        finally:
+            sys.argv = original_argv
+    else:
+        vllm_main()
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())