Merge 165f061b81 into 19bb31f60e

litellm/_logging.py

@@ -1122,7 +1122,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.
 

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,
+    )

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
@@ -35,6 +36,8 @@ from litellm.types.llms.base import LiteLLMBaseModel
 
 if TYPE_CHECKING:
     from fastapi import FastAPI
+
+    from litellm.proxy.control_plane.main import Tls
 else:
     FastAPI = Any
 
@@ -889,6 +892,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")
@@ -1196,6 +1243,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,
@@ -1226,10 +1284,12 @@ class ProxyInitializationHelpers:
     ),
     envvar="PROMETHEUS_METRICS_PORT",
 )
+@click.pass_context
 def run_server(
+    ctx: click.Context,
     cli_args,
-    host,
-    port,
+    host: str,
+    port: int,
     api_base,
     api_version,
     model,
@@ -1280,7 +1340,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
@@ -1298,6 +1360,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
 
     args: Final = locals()
     try:
@@ -1401,17 +1487,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 (

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