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

@@ -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)

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:

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]

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