refactor(proxy): let the gunicorn server take a private listener (#50871)

Extract serving_port() and configure_database() as they were inline, turn the
LITELLM_RUST parser into a reusable RustSwitch type, and reformat
tests/unit/proxy/test_proxy_cli.py to the 120 column limit.

run_gunicorn_server takes private_listener, an inherited socket or an assigned
loopback port. With one, gunicorn binds it instead of host:port and trusts no
peer's forwarding headers, since every peer on it is the local process in front.
Both are gunicorn's own settings, bind and forwarded_allow_ips, so the uvicorn
worker is unchanged. The reload options take factory. No caller passes either
yet, so every entry point serves as before.

gateway/launch.py

@@ -75,7 +75,7 @@ def _serve(argv: Sequence[str]) -> None:
     uvicorn_main(tuple(argv), prog_name="uvicorn")
 
 
-def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -> None:
+def configure_database() -> None:
     from litellm.proxy.db.db_url_settings import DatabaseURLSettings
     from litellm.proxy.db.pgbouncer import PgBouncerError, PgBouncerSettings, export_pooled_database_url
 
@@ -86,6 +86,10 @@ def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -
         sys.exit(f"LiteLLM gateway: in-container pgbouncer could not start: {pooled_url.reason}")
     if pooled_url is not None:
         export_pooled_database_url(pooled_url)
+
+
+def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -> None:
+    configure_database()
     serve(uvicorn_argv(argv, os.environ))
 
 

litellm/proxy/dev_reload.py

@@ -13,7 +13,7 @@ from socket import socket
 from typing import Final, Literal
 
 from pydantic import ConfigDict, ImportString, JsonValue, TypeAdapter, with_config
-from typing_extensions import ReadOnly, Required, TypedDict
+from typing_extensions import NotRequired, ReadOnly, Required, TypedDict
 from uvicorn import Config, Server
 from uvicorn import config as uvicorn_config
 from uvicorn.supervisors import ChangeReload, Multiprocess
@@ -42,6 +42,7 @@ class ServerOptions(TypedDict, total=False):
 @with_config(ConfigDict(extra="allow"))
 class ReloadOptions(ServerOptions):
     app: Required[ReadOnly[str]]
+    factory: NotRequired[ReadOnly[bool]]
     reload: ReadOnly[bool]
     reload_dirs: ReadOnly[tuple[str, ...]]
     reload_includes: ReadOnly[tuple[str, ...]]
@@ -92,7 +93,9 @@ def load_worker_app(import_path: str, worker_app_factory: bool) -> Callable[...,
 def pool_config(options: ReloadOptions, workers: int, worker_overrides: Mapping[str, object] | None = None) -> Config:
     overrides: Final = worker_overrides or {}
     worker_app: Final = TypeAdapter(str).validate_python(overrides.get("app", options["app"]))
-    worker_app_factory: Final = TypeAdapter(bool).validate_python(overrides.get("factory", False))
+    worker_app_factory: Final = TypeAdapter(bool).validate_python(
+        overrides.get("factory", options.get("factory", False))
+    )
     return Config(
         partial(load_worker_app, worker_app, worker_app_factory),
         factory=True,

litellm/proxy/proxy_cli.py

@@ -31,6 +31,7 @@ from litellm.proxy.db.pgbouncer import (
     start_in_container_pgbouncer,
 )
 from litellm.proxy.db.query_engine_reaper import start_query_engine_reaper
+from litellm.proxy.server_listener import PrivateListener, gunicorn_overrides
 from litellm.types.llms.base import LiteLLMBaseModel
 
 if TYPE_CHECKING:
@@ -686,6 +687,7 @@ class ProxyInitializationHelpers:
         max_requests_before_restart: int | None = None,
         max_requests_before_restart_jitter: int | None = None,
         max_request_header_size: int | None = None,
+        private_listener: PrivateListener | None = None,
     ):
         """
         Run litellm with `gunicorn`
@@ -766,6 +768,7 @@ class ProxyInitializationHelpers:
             "accesslog": "-",  # Log to stdout
             "timeout": 600,  # default to very high number, bedrock/anthropic.claude-v2:1 can take 30+ seconds for the 1st chunk to come in
             "access_log_format": '%(h)s %(l)s %(u)s %(t)s "%(r)s" %(s)s %(b)s',
+            **gunicorn_overrides(private_listener),
         }
 
         # Optional: recycle workers after N requests to mitigate memory growth
@@ -821,7 +824,7 @@ class ProxyInitializationHelpers:
     _run_ollama_serve = run_ollama_serve
 
     @staticmethod
-    def is_port_in_use(port):
+    def is_port_in_use(port: int):
         import socket
 
         with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
@@ -829,6 +832,10 @@ class ProxyInitializationHelpers:
 
     _is_port_in_use = is_port_in_use
 
+    @staticmethod
+    def serving_port(port: int) -> int:
+        return random.randint(1024, 49152) if port == 4000 and ProxyInitializationHelpers.is_port_in_use(port) else port
+
     @staticmethod
     def get_loop_type():
         """Helper function to determine the event loop type based on platform"""
@@ -1228,8 +1235,8 @@ class ProxyInitializationHelpers:
 )
 def run_server(
     cli_args,
-    host,
-    port,
+    host: str,
+    port: int,
     api_base,
     api_version,
     model,
@@ -1717,8 +1724,7 @@ def run_server(
                 )
                 sys.exit(1)
             export_pooled_database_url(pooled_database_url)
-        if port == 4000 and ProxyInitializationHelpers.is_port_in_use(port):
-            port = random.randint(1024, 49152)
+        port = ProxyInitializationHelpers.serving_port(port)
         if prometheus_metrics_port == port:
             raise click.UsageError("--prometheus_metrics_port must differ from --port")
 

litellm/proxy/server_listener.py

@@ -0,0 +1,38 @@
+from collections.abc import Mapping
+from dataclasses import dataclass
+from types import MappingProxyType
+from typing import Final, TypeAlias
+
+from typing_extensions import assert_never
+
+LOOPBACK_HOST: Final = "127.0.0.1"
+NO_TRUSTED_PROXIES: Final = ""
+
+
+@dataclass(frozen=True, slots=True)
+class InheritedFd:
+    fd: int
+
+
+@dataclass(frozen=True, slots=True)
+class AssignedPort:
+    port: int
+
+
+PrivateListener: TypeAlias = InheritedFd | AssignedPort
+
+
+def _gunicorn_bind(listener: PrivateListener) -> str:
+    match listener:
+        case InheritedFd(fd):
+            return f"fd://{fd}"
+        case AssignedPort(port):
+            return f"{LOOPBACK_HOST}:{port}"
+        case _:
+            assert_never(listener)
+
+
+def gunicorn_overrides(listener: PrivateListener | None) -> Mapping[str, str]:
+    if listener is None:
+        return MappingProxyType({})
+    return MappingProxyType({"bind": _gunicorn_bind(listener), "forwarded_allow_ips": NO_TRUSTED_PROXIES})

litellm/rust_bridge/configuration.py

@@ -3,13 +3,29 @@ from __future__ import annotations
 import os
 from enum import Enum, auto
 from functools import lru_cache
-from typing import Final
+from typing import Annotated, Final, TypeAlias
 
-from pydantic import TypeAdapter, ValidationError
+from pydantic import BeforeValidator, TypeAdapter, ValidationError
 from typing_extensions import assert_never
 
-_GLOBAL_ENV_NAME: Final = "LITELLM_RUST"
-_ENV_BOOL: Final = TypeAdapter(bool)
+ENV_NAME: Final = "LITELLM_RUST"
+
+
+def _blank_as_unset(value: object) -> object:
+    stripped: Final = value.strip() if isinstance(value, str) else value
+    return None if stripped == "" else stripped
+
+
+RustSwitch: TypeAlias = Annotated[bool | None, BeforeValidator(_blank_as_unset)]
+_RUST_SWITCH: Final = TypeAdapter[bool | None](RustSwitch)
+
+
+@lru_cache(maxsize=16)
+def _switch_or_unset(raw: str | None) -> bool | None:
+    try:
+        return _RUST_SWITCH.validate_python(raw)
+    except ValidationError:
+        return None
 
 
 class Rollout(Enum):
@@ -33,17 +49,6 @@ class _RustConfiguration:
 _CONFIGURATION: Final = _RustConfiguration()
 
 
-@lru_cache(maxsize=16)
-def _parse_env_bool(value: str | None) -> bool | None:
-    """`LITELLM_RUST` as a bool; cached by raw value because `decision` runs per tokenizer call."""
-    if value is None:
-        return None
-    try:
-        return _ENV_BOOL.validate_python(value.strip())
-    except ValidationError:
-        return None
-
-
 def decide(
     rollout: Rollout,
     *,
@@ -72,7 +77,7 @@ def decision(rollout: Rollout) -> Decision:
     return decide(
         rollout,
         process_override=_CONFIGURATION.override,
-        environment_override=_parse_env_bool(os.getenv(_GLOBAL_ENV_NAME)),
+        environment_override=_switch_or_unset(os.getenv(ENV_NAME)),
     )
 
 

tests/unit/proxy/test_dev_reload.py

@@ -86,6 +86,39 @@ def test_a_broken_factory_exits_the_worker_as_a_startup_failure(
     assert exit_info.value.code == uvicorn_config.STARTUP_FAILURE
 
 
+@pytest.mark.parametrize("worker_overrides", [None, {"factory": True}])
+def test_a_factory_flag_in_the_options_reaches_the_workers(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, worker_overrides: dict[str, object] | None
+) -> None:
+    (tmp_path / "reload_factory_flag_app.py").write_text(
+        "async def served(scope, receive, send):\n"
+        "    await send({'type': 'http.response.start', 'status': 204, 'headers': []})\n"
+        "    await send({'type': 'http.response.body', 'body': b'', 'more_body': False})\n"
+        "\n"
+        "def factory():\n"
+        "    return served\n"
+    )
+    monkeypatch.syspath_prepend(str(tmp_path))
+    config: Final = pool_config(
+        RELOAD_OPTIONS.validate_python({**_options("reload_factory_flag_app:factory"), "factory": True}),
+        workers=POOL_WORKERS,
+        worker_overrides=worker_overrides,
+    )
+    config.load()
+    loaded_app: Final = config.loaded_app
+    assert loaded_app is not None
+    sent: Final[list[Message]] = []
+
+    async def receive() -> Message:
+        return {"type": "http.request", "body": b"", "more_body": False}
+
+    async def send(message: Message) -> None:
+        sent.append(message)
+
+    asyncio.run(loaded_app({"type": "http", "method": "GET", "path": "/", "headers": []}, receive, send))
+    assert sent[0] == {"type": "http.response.start", "status": 204, "headers": []}
+
+
 @REQUIRES_UVICORN_STARTUP_FAILURE
 @pytest.mark.parametrize(
     ("module", "source"),
@@ -118,7 +151,9 @@ def test_the_reload_pool_defaults_the_worker_healthcheck_and_keeps_an_explicit_o
     explicit: Final = RELOAD_WORKER_HEALTHCHECK_SECONDS + 7
     assert pool_config(_options("app:app"), workers=2).timeout_worker_healthcheck == RELOAD_WORKER_HEALTHCHECK_SECONDS
     assert (
-        pool_config({**_options("app:app"), "timeout_worker_healthcheck": explicit}, workers=2).timeout_worker_healthcheck
+        pool_config(
+            {**_options("app:app"), "timeout_worker_healthcheck": explicit}, workers=2
+        ).timeout_worker_healthcheck
         == explicit
     )
 
@@ -142,7 +177,9 @@ def test_the_reload_pool_accepts_and_keeps_every_uvicorn_arg_the_proxy_cli_build
 
 
 def test_the_reload_pool_passes_a_uvicorn_arg_it_does_not_list_through_to_the_workers() -> None:
-    config: Final = pool_config(RELOAD_OPTIONS.validate_python({**_options("app:app"), "backlog": 123}), workers=POOL_WORKERS)
+    config: Final = pool_config(
+        RELOAD_OPTIONS.validate_python({**_options("app:app"), "backlog": 123}), workers=POOL_WORKERS
+    )
     assert (config.backlog, config.reload) == (123, False)
 
 

tests/unit/proxy/test_server_listener.py

@@ -0,0 +1,28 @@
+import os
+from typing import Final
+
+import pytest
+
+from litellm.proxy.server_listener import AssignedPort, InheritedFd, PrivateListener, gunicorn_overrides
+
+
+def test_gunicorn_keeps_its_own_listener_and_trust_without_a_private_listener() -> None:
+    assert dict(gunicorn_overrides(None)) == {}
+
+
+@pytest.mark.skipif(os.name == "nt", reason="gunicorn does not run on Windows")
+@pytest.mark.parametrize(
+    ("listener", "address"),
+    [
+        (InheritedFd(17), 17),
+        (AssignedPort(45556), ("127.0.0.1", 45556)),
+    ],
+)
+def test_gunicorn_serves_a_private_listener_and_trusts_no_peer_to_name_the_client(
+    listener: PrivateListener, address: int | tuple[str, int]
+) -> None:
+    gunicorn_config: Final = pytest.importorskip("gunicorn.config").Config()
+    gunicorn_config.set("forwarded_allow_ips", "*")
+    for name, value in gunicorn_overrides(listener).items():
+        gunicorn_config.set(name, value)
+    assert (gunicorn_config.address, gunicorn_config.forwarded_allow_ips) == ([address], [])

tests/unit/rust_bridge/test_configuration.py

@@ -106,8 +106,10 @@ def test_process_override_applies_when_environment_is_unset() -> None:
     assert configuration.rust_enabled() is True
 
 
-@pytest.mark.parametrize("value", ("", " ", "sometimes", "2"))
-def test_invalid_environment_value_is_ignored(monkeypatch: pytest.MonkeyPatch, value: str) -> None:
+@pytest.mark.parametrize("value", ("", " ", "sometimes", "2", "enabled"))
+def test_blank_or_invalid_environment_value_leaves_the_switch_unset(
+    monkeypatch: pytest.MonkeyPatch, value: str
+) -> None:
     monkeypatch.setenv("LITELLM_RUST", value)
 
     assert configuration.rust_enabled() is False
@@ -116,6 +118,13 @@ def test_invalid_environment_value_is_ignored(monkeypatch: pytest.MonkeyPatch, v
     assert configuration.rust_enabled() is True
 
 
+def test_environment_is_read_on_every_decision(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("LITELLM_RUST", "1")
+    assert configuration.rust_enabled() is True
+    monkeypatch.setenv("LITELLM_RUST", "0")
+    assert configuration.rust_enabled() is False
+
+
 def test_process_override_and_reset_apply_to_existing_threads() -> None:
     with ThreadPoolExecutor(max_workers=1) as executor:
         assert executor.submit(configuration.rust_enabled).result() is False