feat(control_plane): --control-plane flag serves the control plane API from the existing image
Some checks are pending
Cost map pins recheck / rerun-cost-map-guard (push) Waiting to run
LiteLLM Rust DB / db-queries (push) Waiting to run
LiteLLM Rust / rust-lint (push) Waiting to run
LiteLLM Rust / rust-test (push) Waiting to run
LiteLLM Rust / rust-wheel (push) Waiting to run
Terraform Modules / fmt, validate, test (aws) (push) Waiting to run
Terraform Modules / fmt, validate, test (gcp) (push) Waiting to run

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-10-13 21:30:40 +00:00
parent 93d1c74c24
commit 537919ab1d
5 changed files with 720 additions and 73 deletions

View file

@ -1127,7 +1127,7 @@ def _initialize_loggers_with_handler(handler: logging.Handler):
logging.getLogger().setLevel(resolve_log_level(configured_level))
def _get_uvicorn_json_log_config():
def _get_uvicorn_json_log_config() -> Mapping[str, object]:
"""
Generate a uvicorn log_config dictionary that applies JSON formatting to all loggers.

View file

@ -0,0 +1,142 @@
import os
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from typing import Final, Protocol
from fastapi import FastAPI
from typing_extensions import ReadOnly, TypedDict
from litellm.proxy.auth.litellm_license import LicenseCheck
from litellm.proxy.control_plane.app import License, create_app
from litellm.proxy.control_plane.container import build_container
from litellm.proxy.control_plane.settings import ControlPlaneSettings, load_settings
@dataclass(frozen=True, slots=True)
class Tls:
keyfile: str
certfile: str
class _LogConfigKwargs(TypedDict, total=False):
log_config: ReadOnly[str | dict[str, object]] # mutable-ok: Uvicorn accepts dict log configs
log_level: ReadOnly[int]
class _UvicornRunner(Protocol):
def __call__(
self,
*,
app: FastAPI,
host: str = ...,
port: int = ...,
server_header: bool = ...,
timeout_graceful_shutdown: int | None = ...,
timeout_keep_alive: int = ...,
h11_max_incomplete_event_size: int | None = ...,
ssl_keyfile: str | None = ...,
ssl_certfile: str | None = ...,
limit_concurrency: int | None = ...,
log_config: str | dict[str, object] | None = ..., # mutable-ok: Uvicorn accepts dict log configs
log_level: int | None = ...,
) -> None: ...
class Runner(Protocol):
def __call__(
self,
app: FastAPI,
*,
host: str,
port: int,
tls: Tls | None,
limit_concurrency: int | None,
log_config: str | None,
keepalive_timeout: int | None,
max_request_header_size: int | None,
) -> None: ...
def _uvicorn(
app: FastAPI,
*,
host: str,
port: int,
tls: Tls | None,
limit_concurrency: int | None,
log_config: str | None,
keepalive_timeout: int | None,
max_request_header_size: int | None,
uvicorn_runner: _UvicornRunner | None = None,
json_logs: bool | None = None,
environ: Mapping[str, str] | None = None,
) -> None:
import uvicorn
import litellm
from litellm._logging import (
_get_uvicorn_json_log_config, # pyright: ignore[reportPrivateUsage] # shared proxy logging config
resolve_log_level,
)
from litellm.constants import DEFAULT_MAX_REQUEST_HEADER_SIZE
from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager
keyfile: Final = None if tls is None else tls.keyfile
certfile: Final = None if tls is None else tls.certfile
json_logs_enabled: Final = litellm.json_logs if json_logs is None else json_logs
environment: Final = os.environ if environ is None else environ
litellm_log: Final = environment.get("LITELLM_LOG")
log_config_kwargs: Final[_LogConfigKwargs] = (
{"log_config": log_config}
if log_config is not None
else _LogConfigKwargs(log_config=dict(_get_uvicorn_json_log_config()))
if json_logs_enabled
else {"log_level": resolve_log_level(litellm_log)}
if litellm_log
else {}
)
runner: Final[_UvicornRunner] = uvicorn.run if uvicorn_runner is None else uvicorn_runner
runner(
app=app,
host=host,
port=port,
server_header=False,
timeout_graceful_shutdown=round(GracefulShutdownManager.get_timeout()),
**log_config_kwargs,
timeout_keep_alive=5 if keepalive_timeout is None else keepalive_timeout,
h11_max_incomplete_event_size=(
DEFAULT_MAX_REQUEST_HEADER_SIZE if max_request_header_size is None else max_request_header_size
),
ssl_keyfile=keyfile,
ssl_certfile=certfile,
limit_concurrency=limit_concurrency,
)
def build_app(settings: ControlPlaneSettings, license_check: License) -> FastAPI:
return create_app(build_container(settings), license_check)
def serve(
host: str,
port: int,
tls: Tls | None = None,
settings: Callable[[], ControlPlaneSettings] = load_settings,
license_check: Callable[[], License] = LicenseCheck,
run: Runner = _uvicorn,
limit_concurrency: int | None = None,
log_config: str | None = None,
keepalive_timeout: int | None = None,
max_request_header_size: int | None = None,
) -> None:
app: Final = build_app(settings(), license_check())
run(
app,
host=host,
port=port,
tls=tls,
limit_concurrency=limit_concurrency,
log_config=log_config,
keepalive_timeout=keepalive_timeout,
max_request_header_size=max_request_header_size,
)

View file

@ -10,9 +10,10 @@ import subprocess
import sys
import urllib.parse as urlparse
from collections.abc import Callable, Iterable, Mapping, Sequence
from dataclasses import dataclass
from pathlib import Path
from types import MappingProxyType
from typing import TYPE_CHECKING, Annotated, Any, Final
from typing import TYPE_CHECKING, Annotated, Any, Final, Protocol
import click
import httpx
@ -37,6 +38,7 @@ from litellm.types.llms.base import LiteLLMBaseModel
if TYPE_CHECKING:
from fastapi import FastAPI
from litellm.proxy.control_plane.main import Tls
from litellm.proxy.rust_forward.peer import ASGIApp
from litellm.proxy.rust_forward.server import PythonServer
else:
@ -597,8 +599,8 @@ class ProxyInitializationHelpers:
app: FastAPI,
host: str,
port: int,
ssl_certfile_path: str,
ssl_keyfile_path: str,
ssl_certfile_path: str | None,
ssl_keyfile_path: str | None,
ciphers: str | None = None,
):
"""
@ -685,8 +687,8 @@ class ProxyInitializationHelpers:
port: int,
app: "ASGIApp",
num_workers: int,
ssl_certfile_path: str,
ssl_keyfile_path: str,
ssl_certfile_path: str | None,
ssl_keyfile_path: str | None,
max_requests_before_restart: int | None = None,
max_requests_before_restart_jitter: int | None = None,
max_request_header_size: int | None = None,
@ -958,6 +960,50 @@ class ProxyInitializationHelpers:
_maybe_setup_prometheus_multiproc_dir = maybe_setup_prometheus_multiproc_dir
def control_plane_requested(control_plane: bool, skip_server_startup: bool, utility_command: bool = False) -> bool:
return control_plane and not skip_server_startup and not utility_command
class _ControlPlaneLicense(Protocol):
def is_premium(self) -> bool: ...
def _decrypt_env_var() -> Mapping[str, str]:
from litellm.secret_managers.aws_secret_manager import decrypt_env_var
return decrypt_env_var()
def _control_plane_license_check() -> _ControlPlaneLicense:
from litellm.proxy.auth.litellm_license import LicenseCheck
return LicenseCheck()
@dataclass(frozen=True, slots=True)
class CliDependencies:
decrypt_env: Callable[[], Mapping[str, str]] = _decrypt_env_var
license_check: Callable[[], _ControlPlaneLicense] = _control_plane_license_check
def _decrypt_kms_env_vars(decrypt_env: Callable[[], Mapping[str, str]]) -> None:
if os.getenv("USE_AWS_KMS") != "True":
return
for key, value in decrypt_env().items():
os.environ[key] = value
def control_plane_tls(ssl_keyfile_path: str | None, ssl_certfile_path: str | None) -> "Tls | None":
from litellm.proxy.control_plane.main import Tls
if ssl_keyfile_path is None and ssl_certfile_path is None:
return None
if ssl_keyfile_path is None or ssl_certfile_path is None:
raise click.UsageError("--control-plane needs both --ssl_keyfile_path and --ssl_certfile_path for TLS")
return Tls(keyfile=ssl_keyfile_path, certfile=ssl_certfile_path)
@click.command()
@click.argument("cli_args", nargs=-1)
@click.option("--host", default="0.0.0.0", help="Host for the server to listen on.", envvar="HOST")
@ -1265,6 +1311,17 @@ class ProxyInitializationHelpers:
),
envvar="USE_V2_MIGRATION_RESOLVER",
)
@click.option(
"--control-plane",
"control_plane",
is_flag=True,
default=False,
envvar="LITELLM_CONTROL_PLANE",
help=(
"Serve the global control plane API instead of the proxy, on the same image. Enterprise only, "
"refuses to start without LITELLM_LICENSE. Configure with CONTROL_PLANE_* env vars."
),
)
@click.option(
"--use_legacy_migration_resolver",
is_flag=True,
@ -1295,7 +1352,9 @@ class ProxyInitializationHelpers:
),
envvar="PROMETHEUS_METRICS_PORT",
)
@click.pass_context
def run_server(
ctx: click.Context,
cli_args,
host: str,
port: int,
@ -1315,7 +1374,7 @@ def run_server(
add_function_to_prompt,
config,
max_budget,
test,
test: str | bool | None,
local,
num_workers: int,
granian_threads,
@ -1324,18 +1383,18 @@ def run_server(
azure_postgresql_auth: bool,
num_requests,
use_queue,
health,
health: bool,
setup,
version,
version: bool,
run_gunicorn,
run_hypercorn,
run_granian,
ssl_keyfile_path,
ssl_certfile_path,
ssl_keyfile_path: str | None,
ssl_certfile_path: str | None,
ciphers,
log_config,
log_config: str | None,
use_prisma_db_push: bool,
skip_server_startup,
skip_server_startup: bool,
validate_config: bool,
check_migrations_against_existing_data: bool,
keepalive_timeout: int | None,
@ -1349,7 +1408,9 @@ def run_server(
use_legacy_migration_resolver: bool,
reload: bool,
prometheus_metrics_port: int | None,
control_plane: bool,
):
dependencies: Final = ctx.obj if isinstance(ctx.obj, CliDependencies) else CliDependencies()
if cli_args:
if cli_args == ("xai-oauth", "login"):
from litellm.llms.xai.oauth import XAIOAuthAuthenticator
@ -1367,6 +1428,30 @@ def run_server(
run_setup_wizard()
return
utility_command: Final = bool(
version or validate_config or health or test or check_migrations_against_existing_data
)
if control_plane_requested(control_plane, skip_server_startup, utility_command):
if num_workers > 1:
raise click.UsageError("scale the control plane with replicas.")
from litellm.proxy.control_plane.main import serve
if run_gunicorn or run_hypercorn or run_granian or reload:
raise click.UsageError(
"--control-plane serves with plain uvicorn, drop --run_gunicorn, --run_hypercorn, --run_granian and --reload"
)
_decrypt_kms_env_vars(dependencies.decrypt_env)
serve(
host=host,
port=port,
tls=control_plane_tls(ssl_keyfile_path, ssl_certfile_path),
license_check=dependencies.license_check,
limit_concurrency=limit_concurrency,
log_config=log_config,
keepalive_timeout=keepalive_timeout,
max_request_header_size=max_request_header_size,
)
return
from pydantic import ValidationError
@ -1513,17 +1598,7 @@ def run_server(
if wants_rds_iam or wants_azure_entra:
DatabaseURLSettings.from_env().apply_writer_url_to_env()
### DECRYPT ENV VAR ###
from litellm.secret_managers.aws_secret_manager import decrypt_env_var
if os.getenv("USE_AWS_KMS", None) is not None and os.getenv("USE_AWS_KMS") == "True":
## V2 IMPLEMENTATION OF AWS KMS - USER WANTS TO DECRYPT MULTIPLE KEYS IN THEIR ENV
new_env_var: Final = decrypt_env_var()
for k, v in new_env_var.items():
os.environ[k] = v
_decrypt_kms_env_vars(dependencies.decrypt_env)
litellm_settings = None
env_config_file: Final = os.environ.get("CONFIG_FILE_PATH")
if (

View file

@ -0,0 +1,237 @@
import logging
from dataclasses import dataclass, field
from typing import Final
import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
from litellm._logging import _get_uvicorn_json_log_config
from litellm.proxy.control_plane.app import LicenseRequired
from litellm.proxy.control_plane.main import Tls, _uvicorn, serve
from litellm.proxy.control_plane.settings import ControlPlaneSettings
from tests.unit.proxy.control_plane.support import MASTER_KEY
@dataclass(frozen=True)
class _License:
premium: bool
def is_premium(self) -> bool:
return self.premium
@dataclass
class _Runner: # mutable-ok: records the app it was asked to serve
served: list[tuple[FastAPI, str, int, Tls | None, int | None, str | None, int | None, int | None]] = field(
default_factory=list
)
def __call__(
self,
app: FastAPI,
*,
host: str,
port: int,
tls: Tls | None,
limit_concurrency: int | None,
log_config: str | None,
keepalive_timeout: int | None,
max_request_header_size: int | None,
) -> None:
self.served.append(
(app, host, port, tls, limit_concurrency, log_config, keepalive_timeout, max_request_header_size)
)
@dataclass(frozen=True)
class _LogConfigUnset:
pass
_LOG_CONFIG_UNSET: Final = _LogConfigUnset()
@dataclass(frozen=True)
class _LogLevelUnset:
pass
_LOG_LEVEL_UNSET: Final = _LogLevelUnset()
@dataclass(frozen=True)
class _UvicornCall:
app: FastAPI
host: str
port: int
server_header: bool
timeout_graceful_shutdown: int | None
timeout_keep_alive: int
h11_max_incomplete_event_size: int | None
ssl_keyfile: str | None
ssl_certfile: str | None
limit_concurrency: int | None
log_config: str | dict[str, object] | None | _LogConfigUnset
log_level: int | None | _LogLevelUnset
@dataclass
class _UvicornCallRecorder: # mutable-ok: records calls made to the injected Uvicorn runner
calls: list[_UvicornCall] = field(default_factory=list)
def __call__(
self,
*,
app: FastAPI,
host: str,
port: int,
server_header: bool,
timeout_graceful_shutdown: int | None,
timeout_keep_alive: int,
h11_max_incomplete_event_size: int | None,
ssl_keyfile: str | None,
ssl_certfile: str | None,
limit_concurrency: int | None,
log_config: str | dict[str, object] | None | _LogConfigUnset = _LOG_CONFIG_UNSET,
log_level: int | None | _LogLevelUnset = _LOG_LEVEL_UNSET,
) -> None:
self.calls.append(
_UvicornCall(
app,
host,
port,
server_header,
timeout_graceful_shutdown,
timeout_keep_alive,
h11_max_incomplete_event_size,
ssl_keyfile,
ssl_certfile,
limit_concurrency,
log_config,
log_level,
)
)
def _settings() -> ControlPlaneSettings:
return ControlPlaneSettings.model_validate({"LITELLM_MASTER_KEY": MASTER_KEY})
def test_serve_builds_the_control_plane_app_and_hands_it_to_the_runner() -> None:
runner: Final = _Runner()
serve("127.0.0.1", 8100, settings=_settings, license_check=lambda: _License(True), run=runner)
[(app, host, port, _tls, _limit_concurrency, _log_config, _keepalive, _header_size)] = runner.served
assert (host, port) == ("127.0.0.1", 8100)
client: Final = TestClient(app)
assert client.get("/healthz").json() == {"status": "ok"}
assert client.get("/api/v1/dataplanes", headers={"Authorization": f"Bearer {MASTER_KEY}"}).json() == {"items": []}
assert client.get("/v1/chat/completions").status_code == 404, "the proxy's routes must not be on this app"
def test_serve_refuses_without_a_license_and_never_binds() -> None:
runner: Final = _Runner()
with pytest.raises(LicenseRequired):
serve("0.0.0.0", 4000, settings=_settings, license_check=lambda: _License(False), run=runner)
assert runner.served == []
def test_serve_hands_configured_tls_to_the_runner() -> None:
runner: Final = _Runner()
tls: Final = Tls(keyfile="/certs/tls.key", certfile="/certs/tls.crt")
serve("0.0.0.0", 8443, tls=tls, settings=_settings, license_check=lambda: _License(True), run=runner)
assert [entry[3] for entry in runner.served] == [tls]
def test_serve_hands_limit_concurrency_to_the_runner() -> None:
runner: Final = _Runner()
serve(
"0.0.0.0",
8100,
settings=_settings,
license_check=lambda: _License(True),
run=runner,
limit_concurrency=250,
)
assert [entry[4] for entry in runner.served] == [250]
def test_uvicorn_keeps_its_default_logging_config_when_unset() -> None:
runner: Final = _UvicornCallRecorder()
_uvicorn(
FastAPI(),
host="127.0.0.1",
port=8100,
tls=None,
limit_concurrency=None,
log_config=None,
keepalive_timeout=None,
max_request_header_size=None,
uvicorn_runner=runner,
json_logs=False,
environ={},
)
assert len(runner.calls) == 1
assert runner.calls[0].log_config is _LOG_CONFIG_UNSET
assert runner.calls[0].log_level is _LOG_LEVEL_UNSET
def test_uvicorn_uses_json_logging_before_the_environment_log_level() -> None:
runner: Final = _UvicornCallRecorder()
_uvicorn(
FastAPI(),
host="127.0.0.1",
port=8100,
tls=None,
limit_concurrency=None,
log_config=None,
keepalive_timeout=None,
max_request_header_size=None,
uvicorn_runner=runner,
json_logs=True,
environ={"LITELLM_LOG": "DEBUG"},
)
call: Final = runner.calls[0]
assert call.log_config == _get_uvicorn_json_log_config()
assert call.log_level is _LOG_LEVEL_UNSET
def test_uvicorn_uses_the_environment_log_level_when_json_logging_is_disabled() -> None:
runner: Final = _UvicornCallRecorder()
_uvicorn(
FastAPI(),
host="127.0.0.1",
port=8100,
tls=None,
limit_concurrency=None,
log_config=None,
keepalive_timeout=None,
max_request_header_size=None,
uvicorn_runner=runner,
json_logs=False,
environ={"LITELLM_LOG": "WARNING"},
)
assert runner.calls[0].log_config is _LOG_CONFIG_UNSET
assert runner.calls[0].log_level == logging.WARNING
def test_uvicorn_forwards_an_explicit_logging_config() -> None:
runner: Final = _UvicornCallRecorder()
_uvicorn(
FastAPI(),
host="127.0.0.1",
port=8100,
tls=None,
limit_concurrency=None,
log_config="logging.yaml",
keepalive_timeout=None,
max_request_header_size=None,
uvicorn_runner=runner,
json_logs=True,
environ={"LITELLM_LOG": "DEBUG"},
)
assert len(runner.calls) == 1
assert runner.calls[0].log_config == "logging.yaml"
assert runner.calls[0].log_level is _LOG_LEVEL_UNSET

View file

@ -5,6 +5,10 @@ import inspect
import json
import logging
import os
import runpy
import sys
import types
import urllib.parse as urlparse
from collections.abc import Iterator
from contextlib import nullcontext
from pathlib import Path
@ -13,20 +17,14 @@ from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import click
import fastapi
import httpx
import pytest
import builtins
import runpy
import sys
import types
import urllib.parse as urlparse
import uvicorn
import yaml
from gunicorn.config import Config as GunicornConfig
from click.testing import CliRunner
from fastapi.testclient import TestClient
from pydantic import TypeAdapter
from uvicorn.config import LOOP_FACTORIES
from uvicorn.importer import import_from_string
@ -34,10 +32,13 @@ import litellm
from litellm._logging import turn_on_json
from litellm.constants import DEFAULT_MAX_REQUEST_HEADER_SIZE, GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS
from litellm.proxy import proxy_cli
from litellm.proxy.proxy_cli import ProxyInitializationHelpers, run_server
from litellm.proxy.proxy_cli import CliDependencies, ProxyInitializationHelpers, run_server
from litellm.proxy.rust_forward.server import Listener, Server
from litellm.proxy.server_listener import AssignedPort, InheritedFd
from litellm.proxy.shutdown.graceful_shutdown_manager import DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT, GracefulShutdownManager
from litellm.proxy.control_plane.main import serve as serve_control_plane
from litellm.proxy.control_plane.settings import ControlPlaneSettings
from tests.unit.proxy.control_plane.support import MASTER_KEY, Licensed
@pytest.fixture(autouse=True)
@ -370,10 +371,10 @@ class TestProxyInitializationHelpers:
py_file = tmp_path / "module.py"
py_file.write_text("x = 1\n")
applied = ProxyInitializationHelpers._patch_statreload_extra_paths([str(config_file)])
applied: Final = ProxyInitializationHelpers._patch_statreload_extra_paths([str(config_file)])
assert applied is True
fake_self = types.SimpleNamespace(config=types.SimpleNamespace(reload_dirs=[tmp_path]))
fake_self: Final = types.SimpleNamespace(config=types.SimpleNamespace(reload_dirs=[tmp_path]))
yielded_paths = {Path(p).resolve() for p in StatReload.iter_py_files(fake_self)}
assert config_file.resolve() in yielded_paths
@ -390,10 +391,10 @@ class TestProxyInitializationHelpers:
env_file = tmp_path / ".env"
env_file.write_text("FOO=bar\n")
applied = ProxyInitializationHelpers._patch_statreload_extra_paths([str(env_file)])
applied: Final = ProxyInitializationHelpers._patch_statreload_extra_paths([str(env_file)])
assert applied is True
fake_self = types.SimpleNamespace(config=types.SimpleNamespace(reload_dirs=[tmp_path]))
fake_self: Final = types.SimpleNamespace(config=types.SimpleNamespace(reload_dirs=[tmp_path]))
yielded_paths = {Path(p).resolve() for p in StatReload.iter_py_files(fake_self)}
assert env_file.resolve() in yielded_paths
@ -423,7 +424,7 @@ class TestProxyInitializationHelpers:
for _ in range(3):
ProxyInitializationHelpers._patch_statreload_extra_paths([str(config_file)])
fake_self = types.SimpleNamespace(config=types.SimpleNamespace(reload_dirs=[tmp_path]))
fake_self: Final = types.SimpleNamespace(config=types.SimpleNamespace(reload_dirs=[tmp_path]))
yielded = list(StatReload.iter_py_files(fake_self))
assert len(yielded) == len(set(map(str, yielded)))
yielded_paths = {Path(p).resolve() for p in yielded}
@ -458,7 +459,7 @@ class TestProxyInitializationHelpers:
assert "override" in warning_text
assert ".env" in warning_text
fake_self = types.SimpleNamespace(config=types.SimpleNamespace(reload_dirs=[tmp_path]))
fake_self: Final = types.SimpleNamespace(config=types.SimpleNamespace(reload_dirs=[tmp_path]))
yielded_paths = {Path(p).resolve() for p in StatReload.iter_py_files(fake_self)}
assert env_file.resolve() in yielded_paths
assert config_file.resolve() in yielded_paths
@ -657,12 +658,12 @@ class TestProxyInitializationHelpers:
database_name_enc = urllib.parse.quote_plus(database_name)
# Construct DATABASE_URL from the provided variables
database_url = (
database_url: Final = (
f"postgresql://{database_username_enc}:{database_password_enc}@{database_host}/{database_name_enc}"
)
# Assert the correct URL was constructed with properly escaped characters
expected_url = (
expected_url: Final = (
"postgresql://user%40with%2Bspecial:test-password-special-chars@localhost:5432/db_name%2Ftest"
)
assert database_url == expected_url
@ -700,7 +701,7 @@ class TestProxyInitializationHelpers:
)
# Remove DATABASE_URL/DIRECT_URL so the CLI doesn't attempt
# real prisma operations when these are set in CI.
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env: Final = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env["PROMETHEUS_MULTIPROC_DIR"] = str(tmp_path)
live_proxy_samples = tmp_path / "counter_123.db"
live_proxy_samples.write_bytes(b"samples of a proxy that is still running")
@ -800,7 +801,7 @@ class TestProxyInitializationHelpers:
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")}
clean_env: Final = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
@ -1001,7 +1002,7 @@ class TestProxyInitializationHelpers:
}
)
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env: Final = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
with (
patch.dict(os.environ, clean_env, clear=True),
@ -1121,7 +1122,7 @@ class TestProxyInitializationHelpers:
}
)
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env: Final = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
with (
patch.dict(os.environ, clean_env, clear=True),
@ -1232,7 +1233,7 @@ class TestProxyInitializationHelpers:
}
)
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env: Final = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
with (
patch.dict(os.environ, clean_env, clear=True),
@ -1444,7 +1445,7 @@ class TestProxyInitializationHelpers:
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")}
clean_env: Final = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
@ -1488,7 +1489,7 @@ class TestProxyInitializationHelpers:
# DB-setup block (un-timeout'd `subprocess.run(["prisma"])` +
# migrate-deploy retry loop) — same isolation every other run_server
# test in this file uses.
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env: Final = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
with (
patch.dict(os.environ, clean_env, clear=True),
@ -1554,7 +1555,7 @@ class TestProxyInitializationHelpers:
# DB-setup block (un-timeout'd `subprocess.run(["prisma"])` +
# migrate-deploy retry loop) — same isolation every other run_server
# test in this file uses.
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env: Final = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
with (
patch.dict(os.environ, clean_env, clear=True),
@ -1581,7 +1582,7 @@ class TestProxyInitializationHelpers:
"port": 8000,
}
result = runner.invoke(run_server, ["--local", "--timeout_worker_healthcheck", "15"])
result: Final = runner.invoke(run_server, ["--local", "--timeout_worker_healthcheck", "15"])
assert result.exit_code == 0
mock_get_args.assert_called_once_with(
@ -1609,7 +1610,7 @@ class TestProxyInitializationHelpers:
mock_key_mgmt = MagicMock()
mock_save_worker_config = MagicMock()
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env: Final = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
with (
patch.dict(
os.environ,
@ -1635,7 +1636,7 @@ class TestProxyInitializationHelpers:
"port": 8000,
}
result = runner.invoke(run_server, ["--local", "--max_requests_before_restart", "123"])
result: Final = runner.invoke(run_server, ["--local", "--max_requests_before_restart", "123"])
assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}"
mock_uvicorn_run.assert_called_once()
@ -1658,7 +1659,7 @@ class TestProxyInitializationHelpers:
pass
runner = CliRunner()
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env: Final = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
with (
patch.dict(os.environ, clean_env, clear=True),
patch("uvicorn.Config", _NewUvicornConfig),
@ -1711,7 +1712,7 @@ class TestProxyInitializationHelpers:
from litellm.proxy.proxy_cli import run_server
runner = CliRunner()
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env: Final = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
@ -1798,7 +1799,7 @@ class TestProxyInitializationHelpers:
from litellm.proxy.proxy_cli import run_server
runner = CliRunner()
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env: Final = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
@ -1820,7 +1821,7 @@ class TestProxyInitializationHelpers:
"port": 8000,
}
result = runner.invoke(run_server, ["--local", "--max_requests_before_restart", "123"])
result: Final = runner.invoke(run_server, ["--local", "--max_requests_before_restart", "123"])
assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}"
assert "--max_requests_before_restart 123 with --num_workers 1" in result.output
@ -1839,7 +1840,7 @@ class TestProxyInitializationHelpers:
from litellm.proxy.proxy_cli import run_server
runner = CliRunner()
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env: Final = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
@ -1881,7 +1882,7 @@ class TestProxyInitializationHelpers:
from litellm.proxy.proxy_cli import run_server
runner = CliRunner()
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env: Final = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
@ -1925,7 +1926,7 @@ class TestProxyInitializationHelpers:
from litellm.proxy.proxy_cli import run_server
runner = CliRunner()
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env: Final = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
with (
patch.dict(os.environ, clean_env, clear=True),
patch.dict(
@ -2280,7 +2281,6 @@ class TestProxyInitializationHelpers:
@patch("builtins.print")
def test_run_server_no_config_passed(self, mock_print, mock_uvicorn_run):
"""Test that run_server properly handles the case when no config is passed"""
import asyncio
from click.testing import CliRunner
@ -2373,7 +2373,7 @@ class TestQueryEngineReaperWiring:
from litellm.proxy.proxy_cli import run_server
runner = CliRunner()
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL", "NUM_WORKERS")}
clean_env: Final = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL", "NUM_WORKERS")}
if env:
clean_env.update(env)
with (
@ -2481,7 +2481,7 @@ class TestRunServerDbSetup:
save_worker_config=MagicMock(),
)
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env: Final = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test"
with (
@ -2559,7 +2559,7 @@ class TestRunServerDbSetup:
save_worker_config=MagicMock(),
)
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env: Final = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test"
clean_env["PATH"] = str(empty_bin)
@ -2719,7 +2719,7 @@ class TestRunServerDbSetup:
save_worker_config=MagicMock(),
)
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env: Final = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test"
with (
@ -2774,7 +2774,7 @@ class TestRunServerDbSetup:
save_worker_config=MagicMock(),
)
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env: Final = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env["DATABASE_URL"] = "sqlite:///data/litellm.db"
with (
@ -2830,7 +2830,7 @@ class TestRunServerDbSetup:
save_worker_config=MagicMock(),
)
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env: Final = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test"
clean_env["USE_V2_MIGRATION_RESOLVER"] = "true"
@ -3198,7 +3198,7 @@ class TestWorkerStartupHooks:
"LITELLM_WORKER_STARTUP_HOOKS": "tests.unit.proxy.test_proxy_cli:_dummy_hook",
}
# Remove DATABASE_URL to avoid real DB setup
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env: Final = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env.update(env_overrides)
with patch.dict(os.environ, clean_env, clear=True):
@ -3221,7 +3221,7 @@ class TestWorkerStartupHooks:
env_overrides = {
"LITELLM_WORKER_STARTUP_HOOKS": "tests.unit.proxy.test_proxy_cli:_dummy_async_hook",
}
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env: Final = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env.update(env_overrides)
with patch.dict(os.environ, clean_env, clear=True):
@ -3241,7 +3241,7 @@ class TestWorkerStartupHooks:
env_overrides = {
"LITELLM_WORKER_STARTUP_HOOKS": "tests.unit.proxy.test_proxy_cli:_failing_hook",
}
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env: Final = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env.update(env_overrides)
with patch.dict(os.environ, clean_env, clear=True):
@ -3270,11 +3270,11 @@ class TestWorkerStartupHooks:
from litellm.proxy.proxy_server import proxy_startup_event
hooks = "tests.unit.proxy.test_proxy_cli:_dummy_hook,tests.unit.proxy.test_proxy_cli:_dummy_async_hook"
hooks: Final = "tests.unit.proxy.test_proxy_cli:_dummy_hook,tests.unit.proxy.test_proxy_cli:_dummy_async_hook"
env_overrides = {
"LITELLM_WORKER_STARTUP_HOOKS": hooks,
}
clean_env = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env: Final = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")}
clean_env.update(env_overrides)
with patch.dict(os.environ, clean_env, clear=True):
@ -3422,7 +3422,7 @@ class TestPostgresStatementTimeoutOptions:
config_path = tmp_path / "config.yaml"
config_path.write_text(yaml.dump({"model_list": [], "general_settings": {"database_statement_timeout": 60}}))
captured = _run_server_and_capture_urls(str(config_path), direct_url="postgresql://t:t@localhost:5432/t")
captured: Final = _run_server_and_capture_urls(str(config_path), direct_url="postgresql://t:t@localhost:5432/t")
assert "options" in urlparse.parse_qs(urlparse.urlparse(captured["DATABASE_URL"]).query)
assert "options" not in urlparse.parse_qs(urlparse.urlparse(captured["DIRECT_URL"]).query)
@ -4535,3 +4535,196 @@ class TestRustForwardServer:
server.graceful_timeout_seconds,
server.trusted_proxies,
) == self._served_by_python(params)
def test_control_plane_rejects_multiple_workers() -> None:
from litellm.proxy.proxy_cli import run_server
result: Final = CliRunner().invoke(run_server, ["--control-plane", "--port", "8100", "--num_workers", "3"])
assert result.exit_code == 2
assert "scale the control plane with replicas." in result.output
@patch("uvicorn.run")
def test_migration_data_check_precedes_control_plane_startup(
mock_uvicorn_run: MagicMock, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.delenv("DATABASE_URL", raising=False)
monkeypatch.delenv("USE_AWS_KMS", raising=False)
from litellm.proxy.proxy_cli import run_server
result: Final = CliRunner().invoke(run_server, ["--control-plane", "--check_migrations_against_existing_data"])
assert result.exit_code == 2
assert "requires DATABASE_URL" in result.output
mock_uvicorn_run.assert_not_called()
def test_control_plane_cli_uses_decrypted_environment(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setenv("USE_AWS_KMS", "True")
monkeypatch.delenv("LITELLM_MASTER_KEY", raising=False)
monkeypatch.delenv("KMS_TEST_SETTING", raising=False)
def decrypt_env() -> dict[str, str]:
return {"LITELLM_MASTER_KEY": MASTER_KEY, "KMS_TEST_SETTING": "decrypted"}
with patch("uvicorn.run") as uvicorn_run:
result: Final = CliRunner().invoke(
run_server,
["--control-plane", "--host", "127.0.0.1", "--port", "8100"],
obj=CliDependencies(decrypt_env=decrypt_env, license_check=Licensed),
)
assert result.exit_code == 0, result.output
assert os.environ["KMS_TEST_SETTING"] == "decrypted"
assert os.environ["LITELLM_MASTER_KEY"] == MASTER_KEY
app: Final = uvicorn_run.call_args.kwargs["app"]
assert TestClient(app).get("/healthz").status_code == 200
def test_control_plane_cli_serves_the_real_app_with_uvicorn_options(monkeypatch: pytest.MonkeyPatch) -> None:
context: Final = run_server.make_context(
"litellm",
[
"--control-plane",
"--host",
"127.0.0.1",
"--port",
"8100",
"--log_config",
"logging.yaml",
"--keepalive_timeout",
"30",
"--max_request_header_size",
"8192",
"--limit_concurrency",
"250",
],
)
params: Final = TypeAdapter(dict[str, object]).validate_python(context.params)
host: Final = params["host"]
port: Final = params["port"]
limit_concurrency: Final = params["limit_concurrency"]
log_config: Final = params["log_config"]
keepalive_timeout: Final = params["keepalive_timeout"]
max_request_header_size: Final = params["max_request_header_size"]
assert params["control_plane"] is True
assert isinstance(host, str) and isinstance(port, int)
assert limit_concurrency == 250
assert log_config == "logging.yaml" and isinstance(log_config, str)
assert keepalive_timeout == 30 and isinstance(keepalive_timeout, int)
assert max_request_header_size == 8192 and isinstance(max_request_header_size, int)
settings: Final = ControlPlaneSettings.model_validate({"LITELLM_MASTER_KEY": MASTER_KEY})
uvicorn_run: Final = MagicMock()
with patch("uvicorn.run", uvicorn_run):
serve_control_plane(
host=host,
port=port,
settings=lambda: settings,
license_check=Licensed,
limit_concurrency=limit_concurrency,
log_config=log_config,
keepalive_timeout=keepalive_timeout,
max_request_header_size=max_request_header_size,
)
call: Final = uvicorn_run.call_args.kwargs
assert (call["host"], call["port"]) == (host, port)
assert (call["log_config"], call["timeout_keep_alive"], call["h11_max_incomplete_event_size"]) == (
"logging.yaml",
30,
8192,
)
assert call["limit_concurrency"] == 250
response: Final = TestClient(call["app"]).get("/healthz")
assert (response.status_code, response.json()) == (200, {"status": "ok"})
def test_limit_concurrency_environment_reaches_control_plane_serve(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setenv("LIMIT_CONCURRENCY", "321")
context: Final = run_server.make_context("litellm", ["--control-plane"])
params: Final = TypeAdapter(dict[str, object]).validate_python(context.params)
host: Final = params["host"]
port: Final = params["port"]
limit_concurrency: Final = params["limit_concurrency"]
assert isinstance(host, str) and isinstance(port, int)
assert limit_concurrency == 321 and isinstance(limit_concurrency, int)
settings: Final = ControlPlaneSettings.model_validate({"LITELLM_MASTER_KEY": MASTER_KEY})
uvicorn_run: Final = MagicMock()
with patch("uvicorn.run", uvicorn_run):
serve_control_plane(
host=host,
port=port,
settings=lambda: settings,
license_check=Licensed,
limit_concurrency=limit_concurrency,
)
assert uvicorn_run.call_args.kwargs["limit_concurrency"] == 321
assert TestClient(uvicorn_run.call_args.kwargs["app"]).get("/healthz").status_code == 200
@pytest.mark.parametrize("args", [("--control-plane", "--version"), ("--control-plane", "--health")])
def test_utility_commands_do_not_decrypt_kms_environment(
args: tuple[str, ...], monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setenv("USE_AWS_KMS", "True")
def fail_if_called() -> dict[str, str]:
raise AssertionError("utility command called the decryptor")
with patch(
"httpx.get",
return_value=httpx.Response(
200,
json={"status": "ok"},
request=httpx.Request("GET", "http://localhost/health"),
),
):
result: Final = CliRunner().invoke(
run_server,
args,
obj=CliDependencies(decrypt_env=fail_if_called),
)
assert result.exit_code == 0, result.output
@pytest.mark.parametrize(
("control_plane", "skip_server_startup", "utility_command", "expected"),
[
(True, False, False, True),
(True, True, False, False),
(True, False, True, False),
(False, False, False, False),
(False, True, True, False),
],
)
def test_migrations_and_utility_commands_win_over_an_inherited_control_plane_flag(
control_plane: bool, skip_server_startup: bool, utility_command: bool, expected: bool
) -> None:
from litellm.proxy.proxy_cli import control_plane_requested
assert control_plane_requested(control_plane, skip_server_startup, utility_command) is expected
@patch("uvicorn.run")
def test_version_prints_the_version_even_with_the_control_plane_flag(mock_uvicorn_run: MagicMock) -> None:
result: Final = CliRunner().invoke(run_server, ["--control-plane", "--version"])
assert result.exit_code == 0, result.output
assert "version" in result.output.lower() and mock_uvicorn_run.call_count == 0
@pytest.mark.parametrize("flag", ["--run_hypercorn", "--run_gunicorn", "--run_granian", "--reload"])
def test_other_servers_and_reload_are_a_usage_error_with_the_control_plane_flag(flag: str) -> None:
from litellm.proxy.proxy_cli import run_server
result: Final = CliRunner().invoke(run_server, ["--control-plane", flag])
assert result.exit_code == 2 and "plain uvicorn" in result.output
def test_control_plane_rejects_a_lone_tls_file() -> None:
result: Final = CliRunner().invoke(run_server, ["--control-plane", "--ssl_keyfile_path", "/c/tls.key"])
assert result.exit_code == 2 and "both" in result.output