From c6afb67f64f4e220020430f61cf238ae29761587 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 13 Oct 2026 21:30:40 +0000 Subject: [PATCH] feat(control_plane): --control-plane flag serves the control plane API from the existing image Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/control_plane/main.py | 60 +++++++++++++++++ litellm/proxy/proxy_cli.py | 44 +++++++++++- tests/unit/proxy/control_plane/test_main.py | 72 ++++++++++++++++++++ tests/unit/proxy/test_proxy_cli.py | 75 ++++++++++++++++++--- 4 files changed, 240 insertions(+), 11 deletions(-) create mode 100644 litellm/proxy/control_plane/main.py create mode 100644 tests/unit/proxy/control_plane/test_main.py diff --git a/litellm/proxy/control_plane/main.py b/litellm/proxy/control_plane/main.py new file mode 100644 index 00000000000..76e830ecfca --- /dev/null +++ b/litellm/proxy/control_plane/main.py @@ -0,0 +1,60 @@ +"""Entrypoint for `litellm --control-plane`: the same image and CLI, a different app. Refuses to serve unlicensed.""" + +from collections.abc import Callable +from dataclasses import dataclass +from typing import Final, Protocol + +from fastapi import FastAPI + +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 + +APP_FACTORY: Final = "litellm.proxy.control_plane.main:app_factory" + + +@dataclass(frozen=True, slots=True) +class Tls: + keyfile: str + certfile: str + + +class Runner(Protocol): + def __call__(self, app: FastAPI, *, host: str, port: int, workers: int, tls: Tls | None) -> None: ... + + +def _uvicorn(app: FastAPI, *, host: str, port: int, workers: int, tls: Tls | None) -> None: + """One worker serves the app already built (and license checked) by ``serve``. More workers need an importable + factory because uvicorn re-imports the app in every worker process, each of which runs the license gate again.""" + import uvicorn + + keyfile: Final = None if tls is None else tls.keyfile + certfile: Final = None if tls is None else tls.certfile + if workers <= 1: + uvicorn.run(app, host=host, port=port, ssl_keyfile=keyfile, ssl_certfile=certfile) + return + uvicorn.run( + APP_FACTORY, host=host, port=port, workers=workers, factory=True, ssl_keyfile=keyfile, ssl_certfile=certfile + ) + + +def build_app(settings: ControlPlaneSettings, license_check: License) -> FastAPI: + return create_app(build_container(settings), license_check) + + +def app_factory() -> FastAPI: + return build_app(load_settings(), LicenseCheck()) + + +def serve( + host: str, + port: int, + workers: int = 1, + tls: Tls | None = None, + settings: Callable[[], ControlPlaneSettings] = load_settings, + license_check: Callable[[], License] = LicenseCheck, + run: Runner = _uvicorn, +) -> None: + app: Final = build_app(settings(), license_check()) + run(app, host=host, port=port, workers=max(1, workers), tls=tls) diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 0fbfa0b06a7..0f86003619c 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -35,6 +35,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 @@ -883,6 +885,22 @@ 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: + """A migration job or a utility command (--version, --validate_config, --health, --test) inherits the + deployment's LITELLM_CONTROL_PLANE, so --skip_server_startup and those commands must win over it.""" + return control_plane and not skip_server_startup and not utility_command + + +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") @@ -1178,6 +1196,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, @@ -1210,8 +1239,8 @@ class ProxyInitializationHelpers: ) def run_server( cli_args, - host, - port, + host: str, + port: int, api_base, api_version, model, @@ -1261,6 +1290,7 @@ def run_server( use_legacy_migration_resolver: bool, reload: bool, prometheus_metrics_port: int | None, + control_plane: bool, ): if cli_args: if cli_args == ("xai-oauth", "login"): @@ -1279,6 +1309,16 @@ def run_server( run_setup_wizard() return + utility_command: Final = bool(version or validate_config or health or test) + if control_plane_requested(control_plane, skip_server_startup, utility_command): + 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" + ) + serve(host=host, port=port, workers=num_workers, tls=control_plane_tls(ssl_keyfile_path, ssl_certfile_path)) + return args: Final = locals() try: 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..bbb6d7786b4 --- /dev/null +++ b/tests/unit/proxy/control_plane/test_main.py @@ -0,0 +1,72 @@ +import importlib +from dataclasses import dataclass, field + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from litellm.proxy.control_plane.app import LicenseRequired +from litellm.proxy.control_plane.main import APP_FACTORY, Tls, serve +from litellm.proxy.control_plane.settings import ControlPlaneSettings + + +@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, int, Tls | None]] = field(default_factory=list) + + def __call__(self, app: FastAPI, *, host: str, port: int, workers: int, tls: Tls | None) -> None: + self.served.append((app, host, port, workers, tls)) + + +def _settings() -> ControlPlaneSettings: + return ControlPlaneSettings.model_validate({"LITELLM_MASTER_KEY": "sk-cp"}) + + +def test_serve_builds_the_control_plane_app_and_hands_it_to_the_runner() -> None: + runner = _Runner() + serve("127.0.0.1", 8100, settings=_settings, license_check=lambda: _License(True), run=runner) + [(app, host, port, workers, _tls)] = runner.served + assert (host, port, workers) == ("127.0.0.1", 8100, 1) + client = TestClient(app) + assert client.get("/healthz").json() == {"status": "ok"} + assert client.get("/api/v1/dataplanes", headers={"Authorization": "Bearer sk-cp"}).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 = _Runner() + with pytest.raises(LicenseRequired): + serve("0.0.0.0", 4000, settings=_settings, license_check=lambda: _License(False), run=runner) + assert runner.served == [] + + +@pytest.mark.parametrize(("requested", "expected"), [(4, 4), (0, 1), (-2, 1)]) +def test_serve_passes_the_worker_count_through_to_the_runner(requested: int, expected: int) -> None: + runner = _Runner() + serve("0.0.0.0", 8100, workers=requested, settings=_settings, license_check=lambda: _License(True), run=runner) + [(_, _, _, workers, _)] = runner.served + assert workers == expected + + +def test_app_factory_is_importable_by_uvicorn_workers_and_builds_the_same_app(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LITELLM_MASTER_KEY", "sk-cp") + monkeypatch.delenv("LITELLM_LICENSE", raising=False) + module_path, attribute = APP_FACTORY.split(":") + factory = getattr(importlib.import_module(module_path), attribute) + with pytest.raises(LicenseRequired): + factory() # no LITELLM_LICENSE in the test environment, so every worker refuses too + + +def test_serve_hands_configured_tls_to_the_runner() -> None: + runner = _Runner() + tls = 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[4] for entry in runner.served] == [tls] diff --git a/tests/unit/proxy/test_proxy_cli.py b/tests/unit/proxy/test_proxy_cli.py index 643c01e3e7f..6d5502260de 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -5,6 +5,10 @@ import json import logging import os from collections.abc import Iterator +import runpy +import sys +import types +import urllib.parse as urlparse from contextlib import nullcontext from pathlib import Path from types import SimpleNamespace @@ -15,16 +19,9 @@ 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 click.testing import CliRunner from uvicorn.config import LOOP_FACTORIES from uvicorn.importer import import_from_string @@ -2295,7 +2292,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 @@ -4455,3 +4451,64 @@ def test_migration_data_check_requires_database_url(monkeypatch): with pytest.raises(click.UsageError, match="requires DATABASE_URL"): ProxyInitializationHelpers.run_migration_data_check(use_migrate=True) + + +class TestControlPlaneFlag: + """`--control-plane` swaps the app served from the existing image, but migration jobs must still win.""" + + @patch("litellm.proxy.control_plane.main.serve") + def test_control_plane_flag_serves_the_control_plane_with_the_worker_count(self, mock_serve): + from litellm.proxy.proxy_cli import run_server + + result = CliRunner().invoke(run_server, ["--control-plane", "--port", "8100", "--num_workers", "3"]) + assert result.exit_code == 0, result.output + mock_serve.assert_called_once() + assert mock_serve.call_args.kwargs["port"] == 8100 + assert mock_serve.call_args.kwargs["workers"] == 3 + + @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( + self, control_plane: bool, skip_server_startup: bool, utility_command: bool, expected: bool + ): + """A migration job or `litellm --version` shares the deployment's env, LITELLM_CONTROL_PLANE=true must not + turn either into a server.""" + from litellm.proxy.proxy_cli import control_plane_requested + + assert control_plane_requested(control_plane, skip_server_startup, utility_command) is expected + + @patch("litellm.proxy.control_plane.main.serve") + def test_version_prints_the_version_even_with_the_control_plane_flag(self, mock_serve): + from litellm.proxy.proxy_cli import run_server + + result = CliRunner().invoke(run_server, ["--control-plane", "--version"]) + assert result.exit_code == 0, result.output + assert "version" in result.output.lower() and mock_serve.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(self, flag: str): + from litellm.proxy.proxy_cli import run_server + + result = CliRunner().invoke(run_server, ["--control-plane", flag]) + assert result.exit_code == 2 and "plain uvicorn" in result.output + + @patch("litellm.proxy.control_plane.main.serve") + def test_tls_flags_reach_the_control_plane_and_a_lone_one_is_a_usage_error(self, mock_serve): + from litellm.proxy.control_plane.main import Tls + from litellm.proxy.proxy_cli import run_server + + both = CliRunner().invoke( + run_server, ["--control-plane", "--ssl_keyfile_path", "/c/tls.key", "--ssl_certfile_path", "/c/tls.crt"] + ) + assert both.exit_code == 0, both.output + assert mock_serve.call_args.kwargs["tls"] == Tls(keyfile="/c/tls.key", certfile="/c/tls.crt") + lone = CliRunner().invoke(run_server, ["--control-plane", "--ssl_keyfile_path", "/c/tls.key"]) + assert lone.exit_code == 2 and "both" in lone.output