diff --git a/litellm/_logging.py b/litellm/_logging.py index ec58239e21b..d1aab417c88 100644 --- a/litellm/_logging.py +++ b/litellm/_logging.py @@ -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. diff --git a/litellm/proxy/control_plane/main.py b/litellm/proxy/control_plane/main.py new file mode 100644 index 00000000000..db8ec39c7de --- /dev/null +++ b/litellm/proxy/control_plane/main.py @@ -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, + ) diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index f252e69d739..cadd45e7f9d 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -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 ( diff --git a/tests/unit/proxy/control_plane/test_main.py b/tests/unit/proxy/control_plane/test_main.py new file mode 100644 index 00000000000..519d803557e --- /dev/null +++ b/tests/unit/proxy/control_plane/test_main.py @@ -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 diff --git a/tests/unit/proxy/test_proxy_cli.py b/tests/unit/proxy/test_proxy_cli.py index d85d685c352..92991f39cd8 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -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