mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
feat(proxy): serve Prometheus /metrics from a separate process via --prometheus_metrics_port (#39889)
* feat(proxy): serve Prometheus /metrics from a separate process via --prometheus_metrics_port Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(proxy): ruff format prometheus_metrics_server Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): fail fast when the separate metrics server cannot start and force the multiproc dir whenever it is enabled - wait for the child's /health before starting uvicorn; raise a ClickException if it exits first (port in use) - create PROMETHEUS_MULTIPROC_DIR whenever --prometheus_metrics_port is set, so DB-configured prometheus callbacks work - honour lowercase prometheus_multiproc_dir; validate the port before spawning - cover main() entry point, readiness, bind failure and wildcard-host probing in tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): pin metrics-server readiness to the child pid so another service on the port cannot pass the health check Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): probe metrics-server readiness through the shared HTTPHandler instead of bare httpx.get Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): serve only /metrics on the prometheus metrics port Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): validate metrics server CLI args with pydantic instead of typing.cast Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): satisfy metrics server lint gate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin <yassin@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
5df0e12e0f
commit
80839bb33c
7 changed files with 649 additions and 33 deletions
|
|
@ -105,13 +105,13 @@
|
|||
"limit": 109
|
||||
},
|
||||
"reportUnknownMemberType": {
|
||||
"limit": 38283
|
||||
"limit": 38271
|
||||
},
|
||||
"reportUnknownParameterType": {
|
||||
"limit": 19584
|
||||
},
|
||||
"reportUnknownVariableType": {
|
||||
"limit": 29829
|
||||
"limit": 29814
|
||||
},
|
||||
"reportUnnecessaryCast": {
|
||||
"limit": 110
|
||||
|
|
|
|||
167
litellm/proxy/prometheus_metrics_server.py
Normal file
167
litellm/proxy/prometheus_metrics_server.py
Normal file
|
|
@ -0,0 +1,167 @@
|
|||
"""Serve Prometheus `/metrics` from its own process so a scrape never runs on an inference worker.
|
||||
|
||||
Workers write their samples to `PROMETHEUS_MULTIPROC_DIR`; this process reads them back with a
|
||||
``MultiProcessCollector`` and serves the aggregated output on a separate port. The proxy CLI starts
|
||||
it with ``--prometheus_metrics_port``. It can also run as a sidecar sharing the same directory:
|
||||
``python -m litellm.proxy.prometheus_metrics_server --host 0.0.0.0 --port 4001``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import atexit
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Sequence
|
||||
from contextlib import closing
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from fastapi import FastAPI
|
||||
from prometheus_client import CollectorRegistry, multiprocess
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
from starlette.types import ASGIApp, Message, Receive, Scope, Send
|
||||
|
||||
from litellm.integrations.prometheus_metrics_endpoint import make_metrics_asgi_app
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
|
||||
METRICS_PATH: Final = "/metrics"
|
||||
PID_HEADER: Final = "x-litellm-metrics-pid"
|
||||
_PARENT_POLL_INTERVAL_SECONDS: Final = 1.0
|
||||
_STARTUP_TIMEOUT_SECONDS: Final = 30.0
|
||||
_STARTUP_POLL_INTERVAL_SECONDS: Final = 0.1
|
||||
_STARTUP_PROBE_TIMEOUT_SECONDS: Final = 1.0
|
||||
_WILDCARD_TO_LOOPBACK: Final = MappingProxyType({"0.0.0.0": "127.0.0.1", "::": "::1"})
|
||||
|
||||
|
||||
class _CliArgs(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
host: str
|
||||
port: int
|
||||
multiproc_dir: str | None
|
||||
|
||||
|
||||
class MetricsServerStartupError(RuntimeError):
|
||||
"""The metrics process died or never answered on its port before the proxy started serving."""
|
||||
|
||||
|
||||
def _add_pid_header(app: ASGIApp) -> ASGIApp:
|
||||
async def app_with_pid(scope: Scope, receive: Receive, send: Send) -> None:
|
||||
async def send_with_pid(message: Message) -> None:
|
||||
if message["type"] == "http.response.start":
|
||||
await send(
|
||||
{
|
||||
**message,
|
||||
"headers": [
|
||||
*message["headers"],
|
||||
(PID_HEADER.encode(), str(os.getpid()).encode()),
|
||||
],
|
||||
}
|
||||
)
|
||||
return
|
||||
await send(message)
|
||||
|
||||
await app(scope, receive, send_with_pid)
|
||||
|
||||
return app_with_pid
|
||||
|
||||
|
||||
def build_metrics_app(multiproc_dir: str) -> FastAPI:
|
||||
registry: Final = CollectorRegistry()
|
||||
multiprocess.MultiProcessCollector(registry, path=multiproc_dir)
|
||||
app: Final = FastAPI(title="LiteLLM Prometheus metrics", docs_url=None, redoc_url=None, openapi_url=None)
|
||||
app.mount(METRICS_PATH, _add_pid_header(make_metrics_asgi_app(registry)))
|
||||
|
||||
return app
|
||||
|
||||
|
||||
def _exit_when_parent_dies(parent_pid: int) -> None:
|
||||
def watch() -> None:
|
||||
while os.getppid() == parent_pid:
|
||||
time.sleep(_PARENT_POLL_INTERVAL_SECONDS)
|
||||
os._exit(0)
|
||||
|
||||
threading.Thread(target=watch, name="litellm-metrics-parent-watchdog", daemon=True).start()
|
||||
|
||||
|
||||
def run_metrics_server(host: str, port: int, multiproc_dir: str) -> None:
|
||||
import uvicorn
|
||||
|
||||
_exit_when_parent_dies(os.getppid())
|
||||
uvicorn.run(build_metrics_app(multiproc_dir), host=host, port=port, log_level="warning", access_log=False)
|
||||
|
||||
|
||||
def metrics_url(host: str, port: int) -> str:
|
||||
probe_host: Final = _WILDCARD_TO_LOOPBACK.get(host, host)
|
||||
netloc: Final = f"[{probe_host}]" if ":" in probe_host else probe_host
|
||||
return f"http://{netloc}:{port}{METRICS_PATH}"
|
||||
|
||||
|
||||
def _answered_by(http: HTTPHandler, url: str, pid: int) -> bool:
|
||||
"""True only when the metrics response comes from our child, not from whatever else holds the port."""
|
||||
try:
|
||||
response: Final = http.get(url) # pyright: ignore[reportUnknownMemberType] # HTTPHandler.get exposes untyped optional mappings
|
||||
return response.status_code == 200 and response.headers.get(PID_HEADER) == str(pid)
|
||||
except httpx.TransportError:
|
||||
return False
|
||||
|
||||
|
||||
def _wait_until_serving(process: subprocess.Popen[bytes], host: str, port: int) -> None:
|
||||
url: Final = metrics_url(host, port)
|
||||
deadline: Final = time.monotonic() + _STARTUP_TIMEOUT_SECONDS
|
||||
with closing(HTTPHandler(timeout=_STARTUP_PROBE_TIMEOUT_SECONDS)) as http:
|
||||
while time.monotonic() < deadline:
|
||||
if (returncode := process.poll()) is not None:
|
||||
raise MetricsServerStartupError(
|
||||
f"Prometheus metrics server exited with code {returncode} before serving {host}:{port}; "
|
||||
"is the port already in use?"
|
||||
)
|
||||
if _answered_by(http, url, process.pid):
|
||||
return
|
||||
time.sleep(_STARTUP_POLL_INTERVAL_SECONDS)
|
||||
process.terminate()
|
||||
raise MetricsServerStartupError(
|
||||
f"Prometheus metrics server did not answer {url} within {_STARTUP_TIMEOUT_SECONDS:.0f}s"
|
||||
)
|
||||
|
||||
|
||||
def start_metrics_server_process(host: str, port: int, multiproc_dir: str) -> subprocess.Popen[bytes]:
|
||||
"""Spawn the metrics server next to the proxy and block until it answers on its port."""
|
||||
process: Final = subprocess.Popen(
|
||||
(
|
||||
sys.executable,
|
||||
"-m",
|
||||
"litellm.proxy.prometheus_metrics_server",
|
||||
"--host",
|
||||
host,
|
||||
"--port",
|
||||
str(port),
|
||||
"--multiproc_dir",
|
||||
multiproc_dir,
|
||||
)
|
||||
)
|
||||
atexit.register(process.terminate)
|
||||
_wait_until_serving(process, host, port)
|
||||
return process
|
||||
|
||||
|
||||
def main(argv: Sequence[str] | None = None) -> None:
|
||||
parser: Final = argparse.ArgumentParser(
|
||||
description="Serve LiteLLM Prometheus metrics from PROMETHEUS_MULTIPROC_DIR"
|
||||
)
|
||||
parser.add_argument("--host", default="0.0.0.0")
|
||||
parser.add_argument("--port", type=int, required=True)
|
||||
parser.add_argument("--multiproc_dir", default=os.environ.get("PROMETHEUS_MULTIPROC_DIR"))
|
||||
args: Final = _CliArgs.model_validate(vars(parser.parse_args(argv)))
|
||||
if not args.multiproc_dir:
|
||||
parser.error("--multiproc_dir or PROMETHEUS_MULTIPROC_DIR is required")
|
||||
run_metrics_server(host=args.host, port=args.port, multiproc_dir=args.multiproc_dir)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -7,7 +7,7 @@ import re
|
|||
import subprocess
|
||||
import sys
|
||||
import urllib.parse as urlparse
|
||||
from collections.abc import Iterable
|
||||
from collections.abc import Iterable, Mapping, Sequence
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
|
|
@ -610,48 +610,49 @@ class ProxyInitializationHelpers:
|
|||
return None # Let uvicorn choose the default loop on Windows
|
||||
return "uvloop"
|
||||
|
||||
@staticmethod
|
||||
def _prometheus_callback_configured(litellm_settings: Mapping[str, object] | None) -> bool:
|
||||
if litellm_settings is None:
|
||||
return False
|
||||
configured: Final = tuple(
|
||||
litellm_settings.get(key) for key in ("callbacks", "success_callback", "failure_callback")
|
||||
)
|
||||
return any(
|
||||
setting == "prometheus"
|
||||
if isinstance(setting, str)
|
||||
else isinstance(setting, Sequence) and "prometheus" in setting
|
||||
for setting in configured
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _maybe_setup_prometheus_multiproc_dir(
|
||||
num_workers: int,
|
||||
litellm_settings: dict | None,
|
||||
) -> None:
|
||||
prometheus_metrics_port: int | None = None,
|
||||
) -> str | None:
|
||||
"""
|
||||
Auto-create PROMETHEUS_MULTIPROC_DIR when running with multiple workers
|
||||
and prometheus is configured as a callback.
|
||||
Auto-create PROMETHEUS_MULTIPROC_DIR when another process needs to read the samples: extra workers
|
||||
with prometheus configured as a callback in config.yaml, or the separate metrics server (always, since
|
||||
callbacks may also be enabled from the DB after startup).
|
||||
"""
|
||||
import tempfile
|
||||
|
||||
if num_workers <= 1 or litellm_settings is None:
|
||||
return
|
||||
|
||||
# Check if prometheus is in any callback list
|
||||
# Each setting can be a list or a single string; normalize to list
|
||||
callbacks = litellm_settings.get("callbacks") or []
|
||||
success_callbacks = litellm_settings.get("success_callback") or []
|
||||
failure_callbacks = litellm_settings.get("failure_callback") or []
|
||||
if isinstance(callbacks, str):
|
||||
callbacks = [callbacks]
|
||||
if isinstance(success_callbacks, str):
|
||||
success_callbacks = [success_callbacks]
|
||||
if isinstance(failure_callbacks, str):
|
||||
failure_callbacks = [failure_callbacks]
|
||||
all_callbacks: Final = callbacks + success_callbacks + failure_callbacks
|
||||
if "prometheus" not in all_callbacks:
|
||||
return
|
||||
if prometheus_metrics_port is None and (
|
||||
num_workers <= 1 or not ProxyInitializationHelpers._prometheus_callback_configured(litellm_settings)
|
||||
):
|
||||
return None
|
||||
|
||||
from litellm.proxy.prometheus_cleanup import wipe_directory
|
||||
|
||||
multiproc_dir = os.environ.get("PROMETHEUS_MULTIPROC_DIR") or os.environ.get("prometheus_multiproc_dir")
|
||||
|
||||
auto_created: Final = not multiproc_dir
|
||||
if not multiproc_dir:
|
||||
multiproc_dir = os.path.join(tempfile.gettempdir(), "litellm_prometheus_multiproc")
|
||||
os.environ["PROMETHEUS_MULTIPROC_DIR"] = multiproc_dir
|
||||
configured_dir: Final = os.environ.get("PROMETHEUS_MULTIPROC_DIR") or os.environ.get("prometheus_multiproc_dir")
|
||||
multiproc_dir: Final = configured_dir or os.path.join(tempfile.gettempdir(), "litellm_prometheus_multiproc")
|
||||
os.environ["PROMETHEUS_MULTIPROC_DIR"] = multiproc_dir
|
||||
|
||||
os.makedirs(multiproc_dir, exist_ok=True)
|
||||
wipe_directory(multiproc_dir)
|
||||
action: Final = "Auto-created" if auto_created else "Using existing"
|
||||
action: Final = "Using existing" if configured_dir else "Auto-created"
|
||||
print(f"LiteLLM: {action} PROMETHEUS_MULTIPROC_DIR={multiproc_dir}")
|
||||
return multiproc_dir
|
||||
|
||||
|
||||
@click.command()
|
||||
|
|
@ -930,6 +931,19 @@ class ProxyInitializationHelpers:
|
|||
default=False,
|
||||
help="Enable uvicorn hot reload (dev only). Also reloads when the --config YAML file changes. Incompatible with --num_workers>1, --run_gunicorn, and --run_hypercorn.",
|
||||
)
|
||||
@click.option(
|
||||
"--prometheus_metrics_port",
|
||||
default=None,
|
||||
type=click.IntRange(min=1, max=65535),
|
||||
help=(
|
||||
"Serve Prometheus /metrics from a separate process on this port (bound to --host) so scraping and "
|
||||
"multi-worker aggregation never run on an inference worker's event loop. Samples appear once the "
|
||||
"`prometheus` callback is enabled (config.yaml or DB). /metrics stays mounted on the main port as well; "
|
||||
"the separate port has no virtual-key auth, so keep it off public ingress. Startup fails if the metrics "
|
||||
"server cannot bind."
|
||||
),
|
||||
envvar="PROMETHEUS_METRICS_PORT",
|
||||
)
|
||||
def run_server(
|
||||
cli_args,
|
||||
host,
|
||||
|
|
@ -980,6 +994,7 @@ def run_server(
|
|||
enforce_prisma_migration_check: bool,
|
||||
use_v2_migration_resolver: bool,
|
||||
reload: bool,
|
||||
prometheus_metrics_port: int | None,
|
||||
):
|
||||
if cli_args:
|
||||
if cli_args == ("xai-oauth", "login"):
|
||||
|
|
@ -1364,6 +1379,8 @@ def run_server(
|
|||
)
|
||||
if port == 4000 and ProxyInitializationHelpers._is_port_in_use(port):
|
||||
port = random.randint(1024, 49152)
|
||||
if prometheus_metrics_port == port:
|
||||
raise click.UsageError("--prometheus_metrics_port must differ from --port")
|
||||
|
||||
import litellm
|
||||
|
||||
|
|
@ -1374,9 +1391,10 @@ def run_server(
|
|||
from litellm.proxy.proxy_server import app
|
||||
|
||||
# Auto-create PROMETHEUS_MULTIPROC_DIR for multi-worker setups
|
||||
ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir(
|
||||
prometheus_multiproc_dir: Final = ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir(
|
||||
num_workers=num_workers,
|
||||
litellm_settings=litellm_settings if config else None,
|
||||
prometheus_metrics_port=prometheus_metrics_port,
|
||||
)
|
||||
|
||||
# Skip server startup if requested (after all setup is done)
|
||||
|
|
@ -1384,6 +1402,20 @@ def run_server(
|
|||
print("LiteLLM: Setup complete. Skipping server startup as requested.")
|
||||
return
|
||||
|
||||
if prometheus_metrics_port is not None and prometheus_multiproc_dir is not None:
|
||||
from litellm.proxy.prometheus_metrics_server import MetricsServerStartupError, start_metrics_server_process
|
||||
|
||||
try:
|
||||
metrics_process: Final = start_metrics_server_process(
|
||||
host=host, port=prometheus_metrics_port, multiproc_dir=prometheus_multiproc_dir
|
||||
)
|
||||
except MetricsServerStartupError as error:
|
||||
raise click.ClickException(str(error)) from error
|
||||
print(
|
||||
f"\033[1;32mLiteLLM: Serving Prometheus metrics on {host}:{prometheus_metrics_port}/metrics "
|
||||
f"(pid {metrics_process.pid})\033[0m"
|
||||
)
|
||||
|
||||
running_uvicorn: Final = run_gunicorn is False and run_hypercorn is False
|
||||
uvicorn_args: Final = ProxyInitializationHelpers._get_default_unvicorn_init_args(
|
||||
host=host,
|
||||
|
|
|
|||
|
|
@ -131,3 +131,43 @@ class TestMaybeSetupPrometheusMultiprocDir:
|
|||
|
||||
# Cleanup
|
||||
os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"litellm_settings",
|
||||
[
|
||||
{"callbacks": ["prometheus"]},
|
||||
{"callbacks": ["langfuse"]},
|
||||
None,
|
||||
],
|
||||
)
|
||||
def test_separate_metrics_port_forces_dir_for_single_worker(self, litellm_settings):
|
||||
"""The separate metrics process reads the samples, so one worker still needs the shared dir, even when
|
||||
prometheus is not in config.yaml (callbacks can be turned on from the DB after startup)."""
|
||||
with patch.dict(os.environ, {}, clear=False):
|
||||
os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None)
|
||||
os.environ.pop("prometheus_multiproc_dir", None)
|
||||
|
||||
result_dir = ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir(
|
||||
num_workers=1,
|
||||
litellm_settings=litellm_settings,
|
||||
prometheus_metrics_port=4001,
|
||||
)
|
||||
|
||||
assert result_dir is not None
|
||||
assert os.environ.get("PROMETHEUS_MULTIPROC_DIR") == result_dir
|
||||
assert os.path.isdir(result_dir)
|
||||
|
||||
os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None)
|
||||
|
||||
def test_lowercase_env_var_is_reused_and_exported_uppercase(self, tmp_path):
|
||||
"""prometheus_client honours both spellings; the metrics server only reads the uppercase one."""
|
||||
with patch.dict(os.environ, {"prometheus_multiproc_dir": str(tmp_path)}, clear=False):
|
||||
os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None)
|
||||
|
||||
result_dir = ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir(
|
||||
num_workers=4,
|
||||
litellm_settings={"callbacks": "prometheus"},
|
||||
)
|
||||
|
||||
assert result_dir == str(tmp_path)
|
||||
assert os.environ["PROMETHEUS_MULTIPROC_DIR"] == str(tmp_path)
|
||||
|
|
|
|||
259
tests/test_litellm/proxy/test_prometheus_metrics_server.py
Normal file
259
tests/test_litellm/proxy/test_prometheus_metrics_server.py
Normal file
|
|
@ -0,0 +1,259 @@
|
|||
"""The separate metrics server must aggregate PROMETHEUS_MULTIPROC_DIR, expose only /metrics, and follow its
|
||||
parent's lifetime.
|
||||
|
||||
Everything here runs on loopback against a child of this test process; no LLM keys or external network.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import socket
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
from unittest.mock import patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from fastapi.testclient import TestClient
|
||||
from prometheus_client import values
|
||||
|
||||
from litellm.proxy.prometheus_metrics_server import (
|
||||
PID_HEADER,
|
||||
MetricsServerStartupError,
|
||||
build_metrics_app,
|
||||
main,
|
||||
metrics_url,
|
||||
start_metrics_server_process,
|
||||
)
|
||||
|
||||
_STARTUP_TIMEOUT_SECONDS: Final = 60.0
|
||||
_SHUTDOWN_TIMEOUT_SECONDS: Final = 15.0
|
||||
|
||||
|
||||
def _write_worker_sample(pid: int, value: float) -> None:
|
||||
"""Write one counter sample into PROMETHEUS_MULTIPROC_DIR the way a proxy worker would."""
|
||||
counter: Final = values.MultiProcessValue(process_identifier=lambda: pid)(
|
||||
"counter",
|
||||
"litellm_requests_metric_total",
|
||||
"litellm_requests_metric_total",
|
||||
("model",),
|
||||
("gpt-5",),
|
||||
"Total number of LLM calls",
|
||||
)
|
||||
counter.inc(value)
|
||||
|
||||
|
||||
def _free_port() -> int:
|
||||
with socket.socket() as sock:
|
||||
sock.bind(("127.0.0.1", 0))
|
||||
return sock.getsockname()[1]
|
||||
|
||||
|
||||
def _wait_for_metrics(port: int, pid: int) -> httpx.Response:
|
||||
deadline: Final = time.monotonic() + _STARTUP_TIMEOUT_SECONDS
|
||||
while time.monotonic() < deadline:
|
||||
try:
|
||||
response: Final = httpx.get(f"http://127.0.0.1:{port}/metrics", follow_redirects=True, timeout=1.0)
|
||||
if response.status_code == 200 and response.headers.get(PID_HEADER) == str(pid):
|
||||
return response
|
||||
except httpx.TransportError:
|
||||
pass
|
||||
time.sleep(0.2)
|
||||
raise AssertionError(f"metrics server on port {port} never served metrics")
|
||||
|
||||
|
||||
def _wait_until_down(port: int) -> None:
|
||||
deadline: Final = time.monotonic() + _SHUTDOWN_TIMEOUT_SECONDS
|
||||
while time.monotonic() < deadline:
|
||||
try:
|
||||
httpx.get(f"http://127.0.0.1:{port}/metrics", timeout=1.0)
|
||||
except httpx.TransportError:
|
||||
return
|
||||
time.sleep(0.2)
|
||||
raise AssertionError(f"metrics server on port {port} kept running after its parent died")
|
||||
|
||||
|
||||
def test_metrics_app_aggregates_multiproc_dir_and_reports_pid(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setenv("PROMETHEUS_MULTIPROC_DIR", str(tmp_path))
|
||||
_write_worker_sample(pid=1001, value=2)
|
||||
_write_worker_sample(pid=1002, value=3)
|
||||
other_dir: Final = tmp_path / "other"
|
||||
other_dir.mkdir()
|
||||
|
||||
client: Final = TestClient(build_metrics_app(str(tmp_path)))
|
||||
metrics: Final = client.get("/metrics")
|
||||
assert metrics.status_code == 200
|
||||
assert metrics.headers[PID_HEADER] == str(os.getpid())
|
||||
assert 'litellm_requests_metric_total{model="gpt-5"} 5.0' in metrics.text
|
||||
|
||||
assert client.get("/health").status_code == 404
|
||||
|
||||
empty: Final = TestClient(build_metrics_app(str(other_dir))).get("/metrics")
|
||||
assert empty.status_code == 200
|
||||
assert "litellm_requests_metric_total" not in empty.text
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("host", "expected"),
|
||||
(
|
||||
("0.0.0.0", "http://127.0.0.1:4001/metrics"),
|
||||
("::", "http://[::1]:4001/metrics"),
|
||||
("10.1.2.3", "http://10.1.2.3:4001/metrics"),
|
||||
("metrics.internal", "http://metrics.internal:4001/metrics"),
|
||||
),
|
||||
)
|
||||
def test_metrics_url_probes_loopback_for_wildcard_binds(host: str, expected: str):
|
||||
assert metrics_url(host, 4001) == expected
|
||||
|
||||
|
||||
def test_main_serves_the_app_for_the_given_dir_with_uvicorn(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setenv("PROMETHEUS_MULTIPROC_DIR", str(tmp_path))
|
||||
_write_worker_sample(pid=2001, value=6)
|
||||
with patch("uvicorn.run") as run:
|
||||
main(["--host", "10.1.2.3", "--port", "4001", "--multiproc_dir", str(tmp_path)])
|
||||
|
||||
run.assert_called_once()
|
||||
assert run.call_args.kwargs["host"] == "10.1.2.3"
|
||||
assert run.call_args.kwargs["port"] == 4001
|
||||
client: Final = TestClient(run.call_args.args[0])
|
||||
metrics: Final = client.get("/metrics")
|
||||
assert metrics.status_code == 200
|
||||
assert metrics.headers[PID_HEADER] == str(os.getpid())
|
||||
assert 'litellm_requests_metric_total{model="gpt-5"} 6.0' in client.get("/metrics").text
|
||||
|
||||
|
||||
def test_main_falls_back_to_env_multiproc_dir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setenv("PROMETHEUS_MULTIPROC_DIR", str(tmp_path))
|
||||
with patch("uvicorn.run") as run:
|
||||
main(["--port", "4001"])
|
||||
|
||||
(app,), served_on = run.call_args
|
||||
assert served_on["host"] == "0.0.0.0"
|
||||
metrics: Final = TestClient(app).get("/metrics")
|
||||
assert metrics.status_code == 200
|
||||
assert metrics.headers[PID_HEADER] == str(os.getpid())
|
||||
|
||||
|
||||
def test_main_rejects_missing_multiproc_dir(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.delenv("PROMETHEUS_MULTIPROC_DIR", raising=False)
|
||||
with patch("uvicorn.run") as run, pytest.raises(SystemExit) as exit_info:
|
||||
main(["--port", "4001"])
|
||||
|
||||
assert exit_info.value.code == 2
|
||||
run.assert_not_called()
|
||||
|
||||
|
||||
def test_start_metrics_server_process_returns_only_once_child_serves(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setenv("PROMETHEUS_MULTIPROC_DIR", str(tmp_path))
|
||||
_write_worker_sample(pid=3001, value=4)
|
||||
port: Final = _free_port()
|
||||
with patch("atexit.register") as register:
|
||||
process: Final = start_metrics_server_process(host="127.0.0.1", port=port, multiproc_dir=str(tmp_path))
|
||||
try:
|
||||
register.assert_called_once_with(process.terminate)
|
||||
assert process.poll() is None
|
||||
startup_metrics: Final = httpx.get(f"http://127.0.0.1:{port}/metrics", follow_redirects=True, timeout=5.0)
|
||||
assert startup_metrics.status_code == 200
|
||||
assert startup_metrics.headers[PID_HEADER] == str(process.pid)
|
||||
metrics: Final = httpx.get(f"http://127.0.0.1:{port}/metrics", follow_redirects=True, timeout=10.0)
|
||||
assert 'litellm_requests_metric_total{model="gpt-5"} 4.0' in metrics.text
|
||||
finally:
|
||||
process.kill()
|
||||
process.wait(timeout=10)
|
||||
|
||||
|
||||
def test_start_metrics_server_process_fails_when_port_is_taken(tmp_path: Path):
|
||||
with socket.socket() as occupied:
|
||||
occupied.bind(("127.0.0.1", 0))
|
||||
occupied.listen()
|
||||
port: Final = occupied.getsockname()[1]
|
||||
with (
|
||||
patch("atexit.register"),
|
||||
pytest.raises(
|
||||
MetricsServerStartupError, match=rf"exited with code [1-9]\d* before serving 127.0.0.1:{port}"
|
||||
),
|
||||
):
|
||||
start_metrics_server_process(host="127.0.0.1", port=port, multiproc_dir=str(tmp_path))
|
||||
|
||||
|
||||
class _ImpostorMetrics(BaseHTTPRequestHandler):
|
||||
"""An unrelated service already on the port that answers /metrics with 200 and plausible metrics."""
|
||||
|
||||
def do_GET(self) -> None:
|
||||
body: Final = b"# HELP impostor_metric A plausible metric\n# TYPE impostor_metric counter\nimpostor_metric 1\n"
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "text/plain")
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
|
||||
def log_message(self, format: str, *args: object) -> None:
|
||||
return
|
||||
|
||||
|
||||
def test_start_metrics_server_process_rejects_metrics_from_another_service_on_the_port(tmp_path: Path):
|
||||
with ThreadingHTTPServer(("127.0.0.1", 0), _ImpostorMetrics) as impostor:
|
||||
threading.Thread(target=impostor.serve_forever, daemon=True).start()
|
||||
port: Final = impostor.server_address[1]
|
||||
impostor_response: Final = httpx.get(f"http://127.0.0.1:{port}/metrics")
|
||||
assert impostor_response.status_code == 200
|
||||
assert "# HELP impostor_metric" in impostor_response.text
|
||||
assert PID_HEADER not in impostor_response.headers
|
||||
with (
|
||||
patch("atexit.register"),
|
||||
pytest.raises(MetricsServerStartupError, match=rf"exited with code [1-9]\d* before serving 127.0.0.1:{port}"),
|
||||
):
|
||||
start_metrics_server_process(host="127.0.0.1", port=port, multiproc_dir=str(tmp_path))
|
||||
impostor.shutdown()
|
||||
|
||||
|
||||
def test_metrics_server_process_serves_and_exits_with_parent(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.setenv("PROMETHEUS_MULTIPROC_DIR", str(tmp_path))
|
||||
_write_worker_sample(pid=2001, value=7)
|
||||
port: Final = _free_port()
|
||||
server_argv: Final = (
|
||||
sys.executable,
|
||||
"-m",
|
||||
"litellm.proxy.prometheus_metrics_server",
|
||||
"--host",
|
||||
"127.0.0.1",
|
||||
"--port",
|
||||
str(port),
|
||||
"--multiproc_dir",
|
||||
str(tmp_path),
|
||||
)
|
||||
parent: Final = subprocess.Popen(
|
||||
(
|
||||
sys.executable,
|
||||
"-c",
|
||||
"import subprocess, sys, time; p = subprocess.Popen(sys.argv[1:]); print(p.pid, flush=True); time.sleep(600)",
|
||||
*server_argv,
|
||||
),
|
||||
stdout=subprocess.PIPE,
|
||||
text=True,
|
||||
)
|
||||
assert parent.stdout is not None
|
||||
server_pid: Final = int(parent.stdout.readline())
|
||||
try:
|
||||
metrics: Final = _wait_for_metrics(port, server_pid)
|
||||
assert metrics.status_code == 200
|
||||
assert metrics.headers[PID_HEADER] == str(server_pid)
|
||||
|
||||
scrape: Final = httpx.get(f"http://127.0.0.1:{port}/metrics", follow_redirects=True, timeout=10.0)
|
||||
assert scrape.status_code == 200
|
||||
assert scrape.headers[PID_HEADER] == str(server_pid)
|
||||
assert 'litellm_requests_metric_total{model="gpt-5"} 7.0' in scrape.text
|
||||
|
||||
parent.kill()
|
||||
parent.wait(timeout=10)
|
||||
_wait_until_down(port)
|
||||
finally:
|
||||
parent.kill()
|
||||
try:
|
||||
os.kill(server_pid, 9)
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
|
|
@ -662,6 +662,124 @@ class TestProxyInitializationHelpers:
|
|||
assert "Invalid value for '--limit_concurrency'" in result.output
|
||||
mock_uvicorn_run.assert_not_called()
|
||||
|
||||
@patch("uvicorn.run")
|
||||
@patch("httpx.HTTPTransport.handle_request")
|
||||
@patch("atexit.register")
|
||||
@patch("subprocess.Popen")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above
|
||||
@patch( # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above
|
||||
"litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False
|
||||
)
|
||||
def test_prometheus_metrics_port_starts_separate_metrics_process(
|
||||
self,
|
||||
mock_should_update,
|
||||
mock_setup_db,
|
||||
mock_popen,
|
||||
mock_atexit_register,
|
||||
mock_handle_request,
|
||||
mock_uvicorn_run,
|
||||
tmp_path,
|
||||
):
|
||||
"""--prometheus_metrics_port must spawn `python -m litellm.proxy.prometheus_metrics_server` on --host
|
||||
with the shared multiproc dir, wait for its /metrics response, and only then start uvicorn. It must stay off by
|
||||
default, refuse to share --port, and abort the proxy when the child dies before serving."""
|
||||
import httpx
|
||||
from click.testing import CliRunner
|
||||
|
||||
from litellm.proxy.proxy_cli import run_server
|
||||
|
||||
runner = CliRunner()
|
||||
mock_popen.return_value = MagicMock(pid=4242, **{"poll.return_value": None})
|
||||
probed_urls: list[str] = []
|
||||
|
||||
def child_metrics(request: httpx.Request) -> httpx.Response:
|
||||
probed_urls.append(str(request.url))
|
||||
return httpx.Response(200, headers={"x-litellm-metrics-pid": "4242"}, content=b"")
|
||||
|
||||
mock_handle_request.side_effect = child_metrics
|
||||
mock_proxy_module = MagicMock(
|
||||
app=MagicMock(),
|
||||
ProxyConfig=MagicMock(),
|
||||
KeyManagementSettings=MagicMock(),
|
||||
save_worker_config=MagicMock(),
|
||||
)
|
||||
clean_env = {
|
||||
k: v
|
||||
for k, v in os.environ.items()
|
||||
if k not in ("DATABASE_URL", "DIRECT_URL", "PROMETHEUS_METRICS_PORT")
|
||||
}
|
||||
clean_env["PROMETHEUS_MULTIPROC_DIR"] = str(tmp_path)
|
||||
with (
|
||||
patch.dict(os.environ, clean_env, clear=True),
|
||||
patch.dict(
|
||||
"sys.modules",
|
||||
{
|
||||
"proxy_server": mock_proxy_module,
|
||||
"litellm.proxy.proxy_server": mock_proxy_module,
|
||||
},
|
||||
),
|
||||
patch( # test-quality-ok: same isolation as the sibling CLI tests above
|
||||
"litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args"
|
||||
) as mock_get_args,
|
||||
):
|
||||
mock_get_args.side_effect = lambda *a, **k: {
|
||||
"app": "litellm.proxy.proxy_server:app",
|
||||
"host": "localhost",
|
||||
"port": 8000,
|
||||
}
|
||||
|
||||
result = runner.invoke(
|
||||
run_server,
|
||||
["--local", "--host", "127.0.0.1", "--port", "4000", "--prometheus_metrics_port", "4001"],
|
||||
)
|
||||
assert (
|
||||
result.exit_code == 0
|
||||
), f"exit_code={result.exit_code}, output={result.output}"
|
||||
mock_popen.assert_called_once()
|
||||
spawned = list(mock_popen.call_args.args[0])
|
||||
assert spawned[1:3] == ["-m", "litellm.proxy.prometheus_metrics_server"]
|
||||
assert spawned[3:] == ["--host", "127.0.0.1", "--port", "4001", "--multiproc_dir", str(tmp_path)]
|
||||
assert probed_urls == ["http://127.0.0.1:4001/metrics"]
|
||||
assert "Serving Prometheus metrics on 127.0.0.1:4001/metrics (pid 4242)" in result.output
|
||||
mock_uvicorn_run.assert_called_once()
|
||||
|
||||
mock_popen.reset_mock()
|
||||
mock_uvicorn_run.reset_mock()
|
||||
mock_popen.return_value = MagicMock(pid=4243, **{"poll.return_value": 1})
|
||||
result = runner.invoke(
|
||||
run_server,
|
||||
["--local", "--port", "4000", "--prometheus_metrics_port", "4001"],
|
||||
)
|
||||
assert result.exit_code == 1, f"exit_code={result.exit_code}, output={result.output}"
|
||||
assert "Prometheus metrics server exited with code 1 before serving 0.0.0.0:4001" in result.output
|
||||
mock_uvicorn_run.assert_not_called()
|
||||
|
||||
mock_popen.reset_mock()
|
||||
mock_uvicorn_run.reset_mock()
|
||||
result = runner.invoke(run_server, ["--local"])
|
||||
assert (
|
||||
result.exit_code == 0
|
||||
), f"exit_code={result.exit_code}, output={result.output}"
|
||||
mock_popen.assert_not_called()
|
||||
mock_uvicorn_run.assert_called_once()
|
||||
|
||||
mock_uvicorn_run.reset_mock()
|
||||
result = runner.invoke(
|
||||
run_server,
|
||||
["--local", "--port", "4000", "--prometheus_metrics_port", "4000"],
|
||||
)
|
||||
assert result.exit_code == 2
|
||||
assert "--prometheus_metrics_port must differ from --port" in result.output
|
||||
mock_popen.assert_not_called()
|
||||
mock_uvicorn_run.assert_not_called()
|
||||
|
||||
result = runner.invoke(
|
||||
run_server, ["--local", "--prometheus_metrics_port", "0"]
|
||||
)
|
||||
assert result.exit_code == 2
|
||||
assert "Invalid value for '--prometheus_metrics_port'" in result.output
|
||||
mock_popen.assert_not_called()
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"timeout_config,expected_timeout",
|
||||
[
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@
|
|||
"limit": 22180
|
||||
},
|
||||
"LIT002": {
|
||||
"limit": 26745
|
||||
"limit": 26729
|
||||
},
|
||||
"LIT003": {
|
||||
"limit": 261
|
||||
|
|
@ -27,7 +27,7 @@
|
|||
"limit": 0
|
||||
},
|
||||
"LIT010": {
|
||||
"limit": 16462
|
||||
"limit": 16430
|
||||
},
|
||||
"LIT011": {
|
||||
"limit": 5506
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue