From 1ca5b247d04cd070aebf88a711fb0487539cb1a3 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Wed, 4 Nov 2026 04:02:09 +0900 Subject: [PATCH 01/10] feat(proxy): supervise the opt-in Rust forwarding server Co-Authored-By: Claude Opus 5.5 --- .github/workflows/test-unit.yml | 1 + gateway/launch.py | 71 +- litellm/proxy/dev_reload.py | 7 +- litellm/proxy/proxy_cli.py | 107 ++- litellm/proxy/rust_forward/AGENTS.md | 18 + litellm/proxy/rust_forward/__init__.py | 0 litellm/proxy/rust_forward/mode.py | 177 +++++ litellm/proxy/rust_forward/peer.py | 57 ++ .../proxy/rust_forward/platform/__init__.py | 52 ++ litellm/proxy/rust_forward/platform/posix.py | 45 ++ .../proxy/rust_forward/platform/windows.py | 56 ++ litellm/proxy/rust_forward/server.py | 30 + litellm/proxy/rust_forward/supervisor.py | 332 +++++++++ .../proxy/shutdown/graceful_uvicorn_worker.py | 3 + litellm/rust_bridge/configuration.py | 4 +- litellm_cli/main.py | 8 +- .../configuration/test_rust_forward.py | 180 +++++ .../test_rust_forward_traffic.py | 457 ++++++++++++ tests/unit/gateway/test_launch.py | 110 ++- tests/unit/proxy/rust_forward/__init__.py | 0 tests/unit/proxy/rust_forward/test_mode.py | 112 +++ tests/unit/proxy/rust_forward/test_peer.py | 143 ++++ .../proxy/rust_forward/test_supervisor.py | 241 ++++++ tests/unit/proxy/test_dev_reload.py | 41 +- tests/unit/proxy/test_proxy_cli.py | 684 +++++++----------- 25 files changed, 2475 insertions(+), 461 deletions(-) create mode 100644 litellm/proxy/rust_forward/AGENTS.md create mode 100644 litellm/proxy/rust_forward/__init__.py create mode 100644 litellm/proxy/rust_forward/mode.py create mode 100644 litellm/proxy/rust_forward/peer.py create mode 100644 litellm/proxy/rust_forward/platform/__init__.py create mode 100644 litellm/proxy/rust_forward/platform/posix.py create mode 100644 litellm/proxy/rust_forward/platform/windows.py create mode 100644 litellm/proxy/rust_forward/server.py create mode 100644 litellm/proxy/rust_forward/supervisor.py create mode 100644 tests/integration/configuration/test_rust_forward.py create mode 100644 tests/integration/configuration/test_rust_forward_traffic.py create mode 100644 tests/unit/proxy/rust_forward/__init__.py create mode 100644 tests/unit/proxy/rust_forward/test_mode.py create mode 100644 tests/unit/proxy/rust_forward/test_peer.py create mode 100644 tests/unit/proxy/rust_forward/test_supervisor.py diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 8ac157c94f5..cd15af01c18 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -589,6 +589,7 @@ jobs: tests/unit/proxy/credential_endpoints tests/unit/proxy/discovery_endpoints tests/unit/proxy/shutdown + tests/unit/proxy/rust_forward tests/unit/proxy/public_endpoints tests/unit/proxy/prompts tests/unit/proxy/rag_endpoints diff --git a/gateway/launch.py b/gateway/launch.py index 59f7121e513..f7f98bd1156 100644 --- a/gateway/launch.py +++ b/gateway/launch.py @@ -22,15 +22,16 @@ import sys from collections.abc import Callable, Mapping, Sequence from typing import TYPE_CHECKING, Final +from typing_extensions import assert_never from uvicorn.main import main as uvicorn_main from litellm_cli.main import without_gc_pauses if TYPE_CHECKING: - from fastapi import FastAPI - from litellm.proxy.db.db_url_settings import DatabaseURLSettings from litellm.proxy.db.pgbouncer import PgBouncerError, PgBouncerSettings + from litellm.proxy.rust_forward.peer import ASGIApp + from litellm.proxy.rust_forward.server import PythonServer GATEWAY_APP: Final = "gateway.launch:worker_app" KEEPALIVE_FLAG: Final = "--timeout-keep-alive" @@ -65,17 +66,58 @@ def pool_database_url( return start_in_container_pgbouncer(pgbouncer, upstream_url, token_auth=settings.token_auth()) -def worker_app() -> "FastAPI": +def worker_app() -> "ASGIApp": with without_gc_pauses(): from gateway.main import app - return app + from litellm.proxy.rust_forward.mode import rust_forward_mode + from litellm.proxy.rust_forward.peer import wrap + return wrap(app, rust_forward_mode(os.environ)) 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 python_server(argv: Sequence[str]) -> "PythonServer": + """What uvicorn would serve the public port with, read from the same flags and environment it reads.""" + from pydantic import TypeAdapter + + from litellm.proxy.rust_forward.server import Listener, PythonServer, Server + from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager + + params: Final[Mapping[str, object]] = uvicorn_main.make_context( + "uvicorn", list(uvicorn_argv(argv, os.environ)) + ).params + tls: Final = ( + params["ssl_keyfile"] is not None + or params["ssl_certfile"] is not None + or any(arg.split("=")[0] == "--ssl-ciphers" for arg in argv) + ) + listener: Final = ( + Listener.TLS + if tls + else Listener.UNIX_SOCKET + if params["uds"] is not None + else Listener.INHERITED_FD + if params["fd"] is not None + else Listener.TCP + ) + graceful_timeout: Final = TypeAdapter[int | None](int | None).validate_python(params["timeout_graceful_shutdown"]) + return PythonServer( + server=Server.UVICORN, + listener=listener, + host=TypeAdapter(str).validate_python(params["host"]), + port=TypeAdapter(int).validate_python(params["port"]), + keepalive_seconds=TypeAdapter(int).validate_python(params["timeout_keep_alive"]), + graceful_timeout_seconds=GracefulShutdownManager.get_timeout() + if graceful_timeout is None + else graceful_timeout, + proxy_headers=TypeAdapter(bool).validate_python(params["proxy_headers"]), + trusted_proxies=TypeAdapter[str | None](str | None).validate_python(params["forwarded_allow_ips"]), + ) + + +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,7 +128,24 @@ 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) - serve(uvicorn_argv(argv, os.environ)) + + +def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -> None: + from litellm.proxy.rust_forward.mode import Disabled, Supervised, Supervisor, rust_forward_mode, uvicorn_flags + from litellm.proxy.rust_forward.supervisor import supervise + + mode: Final = rust_forward_mode(os.environ) + match mode: + case Disabled(): + configure_database() + serve(uvicorn_argv(argv, os.environ)) + case Supervisor(): + command: Final = (sys.executable, "-m", "gateway.launch", *argv) + sys.exit(supervise(python_server(argv), command, before_start=configure_database)) + case Supervised(): + serve((*uvicorn_argv(argv, os.environ), *uvicorn_flags(mode))) + case _: + assert_never(mode) if __name__ == "__main__": diff --git a/litellm/proxy/dev_reload.py b/litellm/proxy/dev_reload.py index 500f581b34d..7b057271cb6 100644 --- a/litellm/proxy/dev_reload.py +++ b/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, diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 0c89627e74d..d3b3e073fa5 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -35,6 +35,9 @@ from litellm.types.llms.base import LiteLLMBaseModel if TYPE_CHECKING: from fastapi import FastAPI + + from litellm.proxy.rust_forward.peer import ASGIApp + from litellm.proxy.rust_forward.server import PythonServer else: FastAPI = Any @@ -678,13 +681,15 @@ class ProxyInitializationHelpers: def run_gunicorn_server( host: str, port: int, - app: FastAPI, + app: "ASGIApp", num_workers: int, ssl_certfile_path: str, ssl_keyfile_path: str, max_requests_before_restart: int | None = None, max_requests_before_restart_jitter: int | None = None, max_request_header_size: int | None = None, + bind: str | None = None, + proxy_headers: bool = True, ): """ Run litellm with `gunicorn` @@ -757,9 +762,11 @@ class ProxyInitializationHelpers: ) graceful_shutdown_timeout: Final = GracefulShutdownManager.get_timeout() gunicorn_options: Final = { - "bind": f"{host}:{port}", + "bind": f"{host}:{port}" if bind is None else bind, "workers": num_workers, # default is 1 - "worker_class": build_graceful_shutdown_uvicorn_worker(graceful_shutdown_timeout, max_request_header_size), + "worker_class": build_graceful_shutdown_uvicorn_worker( + graceful_shutdown_timeout, max_request_header_size, proxy_headers=proxy_headers + ), "graceful_timeout": math.ceil(graceful_shutdown_timeout) + GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS, "preload": True, # Add the preload flag, "accesslog": "-", # Log to stdout @@ -820,7 +827,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: @@ -828,6 +835,57 @@ 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 python_server( + host: str, + port: int, + keepalive_timeout: int | None, + run_gunicorn: bool, + run_hypercorn: bool, + run_granian: bool, + tls: bool, + ) -> "PythonServer": + """``run_gunicorn_server`` never hands ``--keepalive_timeout`` or forwarded IPs to gunicorn, so its workers keep + gunicorn's keep-alive and gunicorn's trusted proxies, which include ``::1`` where uvicorn's do not.""" + import uvicorn + from gunicorn.config import Config as GunicornConfig + + from litellm.proxy.rust_forward.server import Listener, PythonServer, Server + from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager + + server: Final = ( + Server.HYPERCORN + if run_hypercorn + else Server.GRANIAN + if run_granian + else Server.GUNICORN + if run_gunicorn + else Server.UVICORN + ) + keepalive: Final = ( + TypeAdapter(int).validate_python(GunicornConfig().keepalive) + if server is Server.GUNICORN + else uvicorn.Config(app="", log_config=None).timeout_keep_alive + if keepalive_timeout is None + else keepalive_timeout + ) + return PythonServer( + server=server, + listener=Listener.TLS if tls else Listener.TCP, + host=host, + port=ProxyInitializationHelpers.serving_port(port), + keepalive_seconds=keepalive, + graceful_timeout_seconds=GracefulShutdownManager.get_timeout(), + proxy_headers=True, + trusted_proxies=",".join(TypeAdapter(list[str]).validate_python(GunicornConfig().forwarded_allow_ips)) + if server is Server.GUNICORN + else None, + ) + @staticmethod def get_loop_type(): """Helper function to determine the event loop type based on platform""" @@ -1227,8 +1285,8 @@ class ProxyInitializationHelpers: ) def run_server( cli_args, - host, - port, + host: str, + port: int, api_base, api_version, model, @@ -1298,6 +1356,33 @@ def run_server( run_setup_wizard() return + from litellm.proxy.rust_forward.mode import ( + Supervisor, + gunicorn_bind, + rust_forward_mode, + server_proxy_headers, + uvicorn_options, + ) + + forward_mode: Final = rust_forward_mode(os.environ) + if isinstance(forward_mode, Supervisor) and not any( + (version, validate_config, health, test, skip_server_startup, check_migrations_against_existing_data) + ): + from litellm.proxy.rust_forward.supervisor import supervise + + server: Final = ProxyInitializationHelpers.python_server( + host, + port, + keepalive_timeout, + run_gunicorn=run_gunicorn, + run_hypercorn=run_hypercorn, + run_granian=run_granian, + tls=any((ssl_keyfile_path, ssl_certfile_path, ciphers)), + ) + raise SystemExit( + supervise(server, (sys.executable, "-m", "litellm.proxy.proxy_cli", *sys.argv[1:])) + ) + args: Final = locals() try: from litellm.proxy.proxy_server import ( @@ -1717,8 +1802,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") @@ -1767,6 +1851,7 @@ def run_server( timeout_worker_healthcheck=(timeout_worker_healthcheck if running_uvicorn else None), max_request_header_size=max_request_header_size, ) + uvicorn_args.update(uvicorn_options(forward_mode, WORKER_APP)) # Optional: recycle uvicorn workers after N requests if max_requests_before_restart is not None: uvicorn_args["limit_max_requests"] = max_requests_before_restart @@ -1831,16 +1916,20 @@ def run_server( workers=num_workers, ) elif run_gunicorn is True: + from litellm.proxy.rust_forward.peer import wrap + ProxyInitializationHelpers.run_gunicorn_server( host=host, port=port, - app=app, + app=wrap(app, forward_mode), num_workers=num_workers, ssl_certfile_path=ssl_certfile_path, ssl_keyfile_path=ssl_keyfile_path, max_requests_before_restart=max_requests_before_restart, max_requests_before_restart_jitter=max_requests_before_restart_jitter, max_request_header_size=max_request_header_size, + bind=gunicorn_bind(forward_mode, f"{host}:{port}"), + proxy_headers=server_proxy_headers(forward_mode), ) elif run_hypercorn is True: ProxyInitializationHelpers.init_hypercorn_server( diff --git a/litellm/proxy/rust_forward/AGENTS.md b/litellm/proxy/rust_forward/AGENTS.md new file mode 100644 index 00000000000..294d94dd21b --- /dev/null +++ b/litellm/proxy/rust_forward/AGENTS.md @@ -0,0 +1,18 @@ +# Rust forwarding, Python half + +With `LITELLM_RUST` on, the supervisor here runs the Rust gateway on the public port and Python on a private loopback listener. This package is the Python half of forwarding and pairs with the crate `litellm-rust/crates/gateway-forward`. Once Rust serves every route, both get deleted together, along with one dispatch branch in each entry point + +Every forwarding decision lives here. Entry points (`litellm/proxy/proxy_cli.py`, `gateway/launch.py`, `litellm_cli/main.py`) only translate between their own flags and this package: + +- In the supervisor they describe what they would serve with as a `server.PythonServer` and hand it to `supervisor.supervise`, which decides whether forwarding can start, what Rust and the Python child receive, and how to report a refusal +- In the supervised child they merge `mode.uvicorn_options`, `mode.uvicorn_flags` or `mode.gunicorn_bind` and `mode.server_proxy_headers` into their server settings, and wrap the app with `peer.wrap` in their worker factory. These are identities when forwarding is off, so entry points never branch on it + +Every OS-specific operation lives behind `platform.Platform`, which `supervise` picks once with `platform.current_platform()`: the gateway binary name, how the private listener reaches Python, the spawn options, and how children are signalled and stopped. `PosixPlatform` hands Python an inherited socket (`mode.InheritedFd`). `WindowsPlatform` hands it an assigned loopback port (`mode.AssignedPort`), starts each child in its own process group and stops it with `CTRL_BREAK_EVENT`. The supervisor and `mode.py` never check the OS themselves + +Windows support exists only for running the proxy locally from the CLI. Nobody deploys the proxy on Windows, and the Windows design leans on that. A port picked by the supervisor can be taken by another process before Python binds it, `CTRL_BREAK_EVENT` needs the console a CLI run always has, and gunicorn is refused. Revisit those choices before treating Windows as a deployment target + +Each environment contract is written and read in one place. `mode.supervised_environment` and `mode.rust_forward_mode` are the two halves of the Python child's, and `supervisor.rust_spec` writes the `LITELLM_FORWARD_*` variables that `gateway-forward`'s `settings.rs` reads. Python owns every value Rust needs, so a new one gets a new variable in `rust_spec` + +The supervisor's whole lifecycle is `supervisor.step`, a pure function from state and event to the next state and the actions to run. Children start in `START_ORDER`, each once the one before it reports `Ready`, and on a stop every started child is signalled at once in reverse start order. Keep decisions in `step` and keep `_run` executing actions, so a new child is an entry in `START_ORDER` and the launch table rather than a new state + +Python must see each request as if the client had connected directly. `x-litellm-rust-peer` is the only Rust-specific input, and `peer.py` honors it only from the loopback hop. The proxy app itself never knows Rust is in front of it diff --git a/litellm/proxy/rust_forward/__init__.py b/litellm/proxy/rust_forward/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/proxy/rust_forward/mode.py b/litellm/proxy/rust_forward/mode.py new file mode 100644 index 00000000000..a95a66cf4bc --- /dev/null +++ b/litellm/proxy/rust_forward/mode.py @@ -0,0 +1,177 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Annotated, Final, TypeAlias + +from pydantic import BaseModel, BeforeValidator, ConfigDict, Field +from typing_extensions import assert_never + +from litellm.rust_bridge.configuration import ENV_NAME + +GATEWAY_FD: Final = "_LITELLM_GATEWAY_FD" +GATEWAY_PORT: Final = "_LITELLM_GATEWAY_PORT" +PROXY_HEADERS: Final = "_LITELLM_GATEWAY_PROXY_HEADERS" +TRUSTED_PROXIES: Final = "FORWARDED_ALLOW_IPS" +DEFAULT_TRUSTED_PROXIES: Final = "127.0.0.1" +RUST_HOP: Final = "127.0.0.1" + + +def _blank_as_unset(value: object) -> object: + stripped: Final = value.strip() if isinstance(value, str) else value + return None if stripped == "" else stripped + + +_Switch: TypeAlias = Annotated[bool | None, BeforeValidator(_blank_as_unset)] + + +class _ForwardEnvironment(BaseModel): + model_config = ConfigDict(frozen=True) + + rust: _Switch = Field(None, validation_alias=ENV_NAME) + listener_fd: int | None = Field(None, validation_alias=GATEWAY_FD) + listener_port: int | None = Field(None, validation_alias=GATEWAY_PORT) + proxy_headers: bool = Field(True, validation_alias=PROXY_HEADERS) + trusted_proxies: str = Field(DEFAULT_TRUSTED_PROXIES, validation_alias=TRUSTED_PROXIES) + + +@dataclass(frozen=True, slots=True) +class Disabled: + pass + + +@dataclass(frozen=True, slots=True) +class Supervisor: + pass + + +@dataclass(frozen=True, slots=True) +class InheritedFd: + fd: int + + +@dataclass(frozen=True, slots=True) +class AssignedPort: + port: int + + +ChildListener: TypeAlias = InheritedFd | AssignedPort + + +@dataclass(frozen=True, slots=True) +class Supervised: + listener: ChildListener + proxy_headers: bool + trusted_proxies: str + + +RustForwardMode: TypeAlias = Disabled | Supervisor | Supervised + + +def _listener_environment(listener: ChildListener) -> tuple[str, str]: + match listener: + case InheritedFd(fd): + return GATEWAY_FD, str(fd) + case AssignedPort(port): + return GATEWAY_PORT, str(port) + case _: + assert_never(listener) + + +def supervised_environment( + listener: ChildListener, proxy_headers: bool, trusted_proxies: str | None +) -> Mapping[str, str]: + name, value = _listener_environment(listener) + return MappingProxyType( + { + name: value, + PROXY_HEADERS: str(proxy_headers).lower(), + **({} if trusted_proxies is None else {TRUSTED_PROXIES: trusted_proxies}), + } + ) + + +def rust_forward_mode(environ: Mapping[str, str]) -> RustForwardMode: + parsed: Final = _ForwardEnvironment.model_validate(environ) + if parsed.listener_fd is not None: + return Supervised(InheritedFd(parsed.listener_fd), parsed.proxy_headers, parsed.trusted_proxies) + if parsed.listener_port is not None: + return Supervised(AssignedPort(parsed.listener_port), parsed.proxy_headers, parsed.trusted_proxies) + if parsed.rust: + return Supervisor() + return Disabled() + + +# A supervised server's TCP peer is always Rust on loopback, so server-level proxy headers would trust any +# client's forwarding headers. They stay off there and ``peer.wrap`` applies the operator's trust instead. + + +def _uvicorn_listener(listener: ChildListener) -> Mapping[str, object]: + match listener: + case InheritedFd(fd): + return MappingProxyType({"fd": fd}) + case AssignedPort(port): + return MappingProxyType({"host": RUST_HOP, "port": port}) + case _: + assert_never(listener) + + +def uvicorn_options(mode: RustForwardMode, factory: str) -> Mapping[str, object]: + match mode: + case Supervised(listener): + return MappingProxyType( + {"app": factory, "factory": True, **_uvicorn_listener(listener), "proxy_headers": False} + ) + case Disabled() | Supervisor(): + return MappingProxyType({}) + case _: + assert_never(mode) + + +def _uvicorn_listener_flags(listener: ChildListener) -> tuple[str, ...]: + match listener: + case InheritedFd(fd): + return ("--fd", str(fd)) + case AssignedPort(port): + return ("--host", RUST_HOP, "--port", str(port)) + case _: + assert_never(listener) + + +def uvicorn_flags(mode: RustForwardMode) -> tuple[str, ...]: + match mode: + case Supervised(listener): + return (*_uvicorn_listener_flags(listener), "--no-proxy-headers") + case Disabled() | Supervisor(): + return () + case _: + assert_never(mode) + + +def _gunicorn_listener(listener: ChildListener) -> str: + match listener: + case InheritedFd(fd): + return f"fd://{fd}" + case AssignedPort(port): + return f"{RUST_HOP}:{port}" + case _: + assert_never(listener) + + +def gunicorn_bind(mode: RustForwardMode, public: str) -> str: + match mode: + case Supervised(listener): + return _gunicorn_listener(listener) + case Disabled() | Supervisor(): + return public + case _: + assert_never(mode) + + +def server_proxy_headers(mode: RustForwardMode) -> bool: + match mode: + case Supervised(): + return False + case Disabled() | Supervisor(): + return True + case _: + assert_never(mode) diff --git a/litellm/proxy/rust_forward/peer.py b/litellm/proxy/rust_forward/peer.py new file mode 100644 index 00000000000..2d03d6ca35c --- /dev/null +++ b/litellm/proxy/rust_forward/peer.py @@ -0,0 +1,57 @@ +from collections.abc import Awaitable, Callable +from typing import TYPE_CHECKING, Final, TypeAlias + +from typing_extensions import assert_never + +from litellm.proxy.rust_forward.mode import RUST_HOP, Disabled, RustForwardMode, Supervised, Supervisor + +if TYPE_CHECKING: + from uvicorn._types import ASGIReceiveCallable, ASGISendCallable, Scope, WWWScope + +PEER_HEADER: Final = b"x-litellm-rust-peer" + +ASGIApp: TypeAlias = Callable[..., Awaitable[None]] + + +def _restore_peer(scope: "WWWScope") -> None: + headers: Final = tuple(scope["headers"]) + peer: Final = next((value for name, value in headers if name == PEER_HEADER), None) + if peer is None: + return + remaining: Final = tuple((name, value) for name, value in headers if name != PEER_HEADER) + scope["headers"] = remaining # rebind-ok: the peer header is private to the Rust hop + client: Final = scope["client"] + if client is None or client[0] != RUST_HOP: + return + host, _, port = peer.decode("latin-1").rpartition(":") + scope["client"] = (host.strip("[]"), int(port)) # rebind-ok: uvicorn's access log reads this same scope + + +class RustPeerMiddleware: + """Make Python see each request as if the client had connected to it directly. + + Rust owns the public port, so Python's peer is always Rust on loopback. Rust passes the real + TCP peer in ``x-litellm-rust-peer`` and leaves every forwarding header untouched. This puts the + real peer back into the scope, then runs uvicorn's own ``ProxyHeadersMiddleware`` with the trust + settings Python would have used on the public port. + """ + + def __init__(self, app: ASGIApp, mode: Supervised) -> None: + from uvicorn.middleware.proxy_headers import ProxyHeadersMiddleware + + self.app: Final = ProxyHeadersMiddleware(app, mode.trusted_proxies) if mode.proxy_headers else app + + async def __call__(self, scope: "Scope", receive: "ASGIReceiveCallable", send: "ASGISendCallable") -> None: + if scope["type"] != "lifespan": + _restore_peer(scope) + await self.app(scope, receive, send) + + +def wrap(app: ASGIApp, mode: RustForwardMode) -> ASGIApp: + match mode: + case Supervised(): + return RustPeerMiddleware(app, mode) + case Disabled() | Supervisor(): + return app + case _: + assert_never(mode) diff --git a/litellm/proxy/rust_forward/platform/__init__.py b/litellm/proxy/rust_forward/platform/__init__.py new file mode 100644 index 00000000000..b40ff5cb015 --- /dev/null +++ b/litellm/proxy/rust_forward/platform/__init__.py @@ -0,0 +1,52 @@ +import asyncio +import os +import signal +from collections.abc import Callable +from contextlib import AbstractContextManager +from dataclasses import dataclass +from typing import Protocol + +from litellm.proxy.rust_forward.mode import ChildListener +from litellm.proxy.rust_forward.server import Server + + +@dataclass(frozen=True, slots=True) +class PrivateListener: + listener: ChildListener + port: int + pass_fds: tuple[int, ...] + + +@dataclass(frozen=True, slots=True) +class SpawnOptions: + pass_fds: tuple[int, ...] + start_new_session: bool + creationflags: int + + +class Platform(Protocol): + @property + def name(self) -> str: ... + + @property + def gateway_name(self) -> str: ... + + @property + def servers(self) -> frozenset[Server]: ... + + def private_listener(self) -> AbstractContextManager[PrivateListener]: ... + + def spawn_options(self, pass_fds: tuple[int, ...]) -> SpawnOptions: ... + + def catch_stop(self, loop: asyncio.AbstractEventLoop, on_stop: Callable[[signal.Signals], None]) -> None: ... + + def signal_child(self, process: asyncio.subprocess.Process, signum: signal.Signals) -> None: ... + + def stop_group(self, process: asyncio.subprocess.Process) -> None: ... + + +def current_platform() -> Platform: + from litellm.proxy.rust_forward.platform.posix import PosixPlatform + from litellm.proxy.rust_forward.platform.windows import WindowsPlatform + + return WindowsPlatform() if os.name == "nt" else PosixPlatform() diff --git a/litellm/proxy/rust_forward/platform/posix.py b/litellm/proxy/rust_forward/platform/posix.py new file mode 100644 index 00000000000..b1b2e500657 --- /dev/null +++ b/litellm/proxy/rust_forward/platform/posix.py @@ -0,0 +1,45 @@ +import asyncio +import contextlib +import os +import signal +import socket +from collections.abc import Callable, Generator +from dataclasses import dataclass +from typing import Final + +from pydantic import TypeAdapter + +from litellm.proxy.rust_forward.mode import RUST_HOP, InheritedFd +from litellm.proxy.rust_forward.platform import PrivateListener, SpawnOptions +from litellm.proxy.rust_forward.server import Server + +_ADDRESS: Final = TypeAdapter(tuple[str, int]) + + +@dataclass(frozen=True, slots=True) +class PosixPlatform: + name: str = "POSIX" + gateway_name: str = "litellm-gateway" + servers: frozenset[Server] = frozenset({Server.UVICORN, Server.GUNICORN}) + + @contextlib.contextmanager + def private_listener(self) -> Generator[PrivateListener]: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as private: + private.bind((RUST_HOP, 0)) + _, port = _ADDRESS.validate_python(private.getsockname()) + yield PrivateListener(InheritedFd(private.fileno()), port, (private.fileno(),)) + + def spawn_options(self, pass_fds: tuple[int, ...]) -> SpawnOptions: + return SpawnOptions(pass_fds=pass_fds, start_new_session=True, creationflags=0) + + def catch_stop(self, loop: asyncio.AbstractEventLoop, on_stop: Callable[[signal.Signals], None]) -> None: + for signum in (signal.SIGTERM, signal.SIGINT): + loop.add_signal_handler(signum, on_stop, signum) + + def signal_child(self, process: asyncio.subprocess.Process, signum: signal.Signals) -> None: + with contextlib.suppress(ProcessLookupError): + process.send_signal(signum) + + def stop_group(self, process: asyncio.subprocess.Process) -> None: + with contextlib.suppress(ProcessLookupError): + os.killpg(process.pid, signal.SIGTERM) diff --git a/litellm/proxy/rust_forward/platform/windows.py b/litellm/proxy/rust_forward/platform/windows.py new file mode 100644 index 00000000000..f32063428e3 --- /dev/null +++ b/litellm/proxy/rust_forward/platform/windows.py @@ -0,0 +1,56 @@ +import asyncio +import contextlib +import signal +import socket +from collections.abc import Callable, Generator +from dataclasses import dataclass +from types import FrameType +from typing import Final + +from pydantic import TypeAdapter + +from litellm.proxy.rust_forward.mode import RUST_HOP, AssignedPort +from litellm.proxy.rust_forward.platform import PrivateListener, SpawnOptions +from litellm.proxy.rust_forward.server import Server + +_ADDRESS: Final = TypeAdapter(tuple[str, int]) +# signal.CTRL_BREAK_EVENT, signal.SIGBREAK and subprocess.CREATE_NEW_PROCESS_GROUP exist only on Windows, so they are +# spelled out to keep this module importable and type checked everywhere. A Windows-only test pins them to the stdlib. +CTRL_BREAK_EVENT: Final = 1 +SIGBREAK: Final = 21 +CREATE_NEW_PROCESS_GROUP: Final = 0x00000200 + + +# Windows support exists only for running the proxy locally from the CLI, never for deploying it. That is why the +# supervisor picks a free loopback port and hands Python its number instead of an open socket, which Windows uvicorn +# cannot serve. Another process could take that port before Python binds it, and the readiness probe would then +# accept it as Python. Python's failed bind still ends the run, which is acceptable on a developer's machine. +@dataclass(frozen=True, slots=True) +class WindowsPlatform: + name: str = "Windows" + gateway_name: str = "litellm-gateway.exe" + servers: frozenset[Server] = frozenset({Server.UVICORN}) + + @contextlib.contextmanager + def private_listener(self) -> Generator[PrivateListener]: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: + probe.bind((RUST_HOP, 0)) + _, port = _ADDRESS.validate_python(probe.getsockname()) + yield PrivateListener(AssignedPort(port), port, ()) + + def spawn_options(self, pass_fds: tuple[int, ...]) -> SpawnOptions: + return SpawnOptions(pass_fds=pass_fds, start_new_session=False, creationflags=CREATE_NEW_PROCESS_GROUP) + + def catch_stop(self, loop: asyncio.AbstractEventLoop, on_stop: Callable[[signal.Signals], None]) -> None: + def handle(signum: int, _frame: FrameType | None) -> None: + loop.call_soon_threadsafe(on_stop, signal.Signals(signum)) + + for signum in (signal.SIGINT, SIGBREAK): + signal.signal(signum, handle) + + def signal_child(self, process: asyncio.subprocess.Process, signum: signal.Signals) -> None: + self.stop_group(process) + + def stop_group(self, process: asyncio.subprocess.Process) -> None: + with contextlib.suppress(ProcessLookupError): + process.send_signal(CTRL_BREAK_EVENT) diff --git a/litellm/proxy/rust_forward/server.py b/litellm/proxy/rust_forward/server.py new file mode 100644 index 00000000000..8398cbda533 --- /dev/null +++ b/litellm/proxy/rust_forward/server.py @@ -0,0 +1,30 @@ +from dataclasses import dataclass +from enum import Enum + + +class Server(Enum): + UVICORN = "uvicorn" + GUNICORN = "gunicorn" + HYPERCORN = "Hypercorn" + GRANIAN = "Granian" + + +class Listener(Enum): + TCP = "plain TCP" + TLS = "TLS" + UNIX_SOCKET = "a Unix socket" + INHERITED_FD = "an inherited file descriptor" + + +@dataclass(frozen=True, slots=True) +class PythonServer: + """What an entry point would serve with on the public port if Rust were not in front of it.""" + + server: Server + listener: Listener + host: str + port: int + keepalive_seconds: int + graceful_timeout_seconds: float + proxy_headers: bool + trusted_proxies: str | None diff --git a/litellm/proxy/rust_forward/supervisor.py b/litellm/proxy/rust_forward/supervisor.py new file mode 100644 index 00000000000..3331b3f0c0b --- /dev/null +++ b/litellm/proxy/rust_forward/supervisor.py @@ -0,0 +1,332 @@ +import asyncio +import contextlib +import os +import shutil +import signal +import sys +import sysconfig +from collections.abc import AsyncGenerator, Callable, Coroutine, Mapping, Sequence +from dataclasses import dataclass +from enum import Enum, IntEnum +from functools import partial +from types import MappingProxyType +from typing import Final, TypeAlias + +from typing_extensions import assert_never + +from litellm.constants import GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS +from litellm.proxy.rust_forward.mode import RUST_HOP, supervised_environment +from litellm.proxy.rust_forward.platform import Platform, PrivateListener, SpawnOptions, current_platform +from litellm.proxy.rust_forward.server import Listener, PythonServer, Server +from litellm.rust_bridge.configuration import ENV_NAME + +PROBE_INTERVAL_SECONDS: Final = 0.05 + + +def gateway_binary(platform: Platform, scripts_dir: str) -> str | None: + return shutil.which(platform.gateway_name, path=scripts_dir) + + +@dataclass(frozen=True, slots=True) +class StartupError: + reason: str + + +def startup_check(platform: Platform, server: PythonServer, binary: str | None) -> str | StartupError: + if server.server in (Server.HYPERCORN, Server.GRANIAN): + return StartupError(f"{ENV_NAME} does not support {server.server.value}") + if server.server not in platform.servers: + return StartupError(f"{ENV_NAME} does not support {server.server.value} on {platform.name}") + if server.listener is not Listener.TCP: + return StartupError(f"{ENV_NAME} does not support {server.listener.value} on the public port") + if binary is None: + return StartupError( + f"{ENV_NAME} requires the litellm-gateway binary. Install litellm[proxy] in this environment" + ) + return binary + + +class Child(Enum): + RUST = "rust" + PYTHON = "python" + + +# Each child starts once the one before it is ready. Python alone takes no connections until it has booted, so Rust +# only opens the public port once Python listens. +START_ORDER: Final = (Child.PYTHON, Child.RUST) + + +class ExitCode(IntEnum): + CLEAN = 0 + CHILD_FAILED = 1 + + +@dataclass(frozen=True, slots=True) +class Ready: + child: Child + + +@dataclass(frozen=True, slots=True) +class Stop: + signum: signal.Signals + + +@dataclass(frozen=True, slots=True) +class Exited: + child: Child + code: int + + +@dataclass(frozen=True, slots=True) +class Running: + up: tuple[Child, ...] + + +@dataclass(frozen=True, slots=True) +class Stopping: + up: frozenset[Child] + + +@dataclass(frozen=True, slots=True) +class Finished: + code: ExitCode + + +@dataclass(frozen=True, slots=True) +class Start: + child: Child + + +@dataclass(frozen=True, slots=True) +class Signal: + child: Child + signum: signal.Signals + + +Event: TypeAlias = Ready | Stop | Exited +State: TypeAlias = Running | Stopping | Finished +Action: TypeAlias = Start | Signal +Transition: TypeAlias = tuple[State, tuple[Action, ...]] + + +def _start_after(up: tuple[Child, ...]) -> Transition: + following: Final = START_ORDER[len(up) : len(up) + 1] + return Running((*up, *following)), tuple(Start(child) for child in following) + + +def boot() -> Transition: + return _start_after(()) + + +# Every child gets the signal at once, as Python alone stopped. Python ends its long streams within its graceful +# timeout, which Rust's drain already covers, so stopping Python only after Rust exits would hold Python's shutdown, +# spend flush included, behind streams only Python can end. Reverse start order lets Rust stop accepting first. +def _stop(up: tuple[Child, ...], signum: signal.Signals) -> Transition: + return Stopping(frozenset(up)), tuple(Signal(child, signum) for child in reversed(up)) + + +def step(state: State, event: Event) -> Transition: + match state: + case Running(up): + match event: + case Ready(child) if up[-1:] == (child,): + return _start_after(up) + case Ready(): + return state, () + case Stop(signum): + return _stop(up, signum) + case Exited(): + return Finished(ExitCode.CHILD_FAILED), () + case _: + assert_never(event) + case Stopping(up): + match event: + case Ready() | Stop(): + return state, () + case Exited(child, code) if child in up and code == 0: + remaining: Final = up - {child} + return (Stopping(remaining) if remaining else Finished(ExitCode.CLEAN)), () + case Exited(): + return Finished(ExitCode.CHILD_FAILED), () + case _: + assert_never(event) + case Finished(): + return state, () + case _: + assert_never(state) + + +@dataclass(frozen=True, slots=True) +class ChildSpec: + argv: tuple[str, ...] + env_overrides: Mapping[str, str] + spawn: SpawnOptions + + +def python_spec( + platform: Platform, command: Sequence[str], private: PrivateListener, server: PythonServer +) -> ChildSpec: + return ChildSpec( + argv=(*command, "--port", str(private.port)), + env_overrides=supervised_environment(private.listener, server.proxy_headers, server.trusted_proxies), + spawn=platform.spawn_options(private.pass_fds), + ) + + +def rust_spec(platform: Platform, binary: str, upstream_port: int, server: PythonServer) -> ChildSpec: + return ChildSpec( + argv=(binary,), + env_overrides=MappingProxyType( + { + "LITELLM_FORWARD_UPSTREAM": f"http://{RUST_HOP}:{upstream_port}", + "LITELLM_FORWARD_HOST": server.host, + "LITELLM_FORWARD_PORT": str(server.port), + "LITELLM_FORWARD_UPSTREAM_KEEPALIVE": str(server.keepalive_seconds), + "LITELLM_FORWARD_SHUTDOWN_DRAIN": str( + server.graceful_timeout_seconds + GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS + ), + } + ), + spawn=platform.spawn_options(()), + ) + + +async def _spawn(spec: ChildSpec) -> asyncio.subprocess.Process: + return await asyncio.create_subprocess_exec( + *spec.argv, + env={**os.environ, **spec.env_overrides}, + pass_fds=spec.spawn.pass_fds, + start_new_session=spec.spawn.start_new_session, + creationflags=spec.spawn.creationflags, + ) + + +Launch: TypeAlias = Callable[[asyncio.Queue[Event]], contextlib.AbstractAsyncContextManager[asyncio.subprocess.Process]] + + +async def _report_exit(child: Child, process: asyncio.subprocess.Process, events: asyncio.Queue[Event]) -> None: + events.put_nowait(Exited(child, await process.wait())) + + +async def _accepts(port: int) -> bool: + try: + _, writer = await asyncio.wait_for(asyncio.open_connection(RUST_HOP, port), PROBE_INTERVAL_SECONDS) + except (OSError, asyncio.TimeoutError): + return False + writer.close() + return True + + +async def _report_ready_once_listening(child: Child, port: int, events: asyncio.Queue[Event]) -> None: + while not await _accepts(port): + await asyncio.sleep(PROBE_INTERVAL_SECONDS) + events.put_nowait(Ready(child)) + + +@contextlib.asynccontextmanager +async def _in_background(work: Coroutine[None, None, None]) -> AsyncGenerator[None]: + task: Final = asyncio.create_task(work) + try: + yield + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + +@contextlib.asynccontextmanager +async def _launch( + platform: Platform, + child: Child, + spec: ChildSpec, + events: asyncio.Queue[Event], +) -> AsyncGenerator[asyncio.subprocess.Process]: + process: Final = await _spawn(spec) + waiter: Final = asyncio.create_task(_report_exit(child, process, events)) + try: + yield process + finally: + platform.stop_group(process) + await waiter + + +@contextlib.asynccontextmanager +async def _launch_until_listening( + platform: Platform, child: Child, spec: ChildSpec, port: int, events: asyncio.Queue[Event] +) -> AsyncGenerator[asyncio.subprocess.Process]: + async with ( + _launch(platform, child, spec, events) as process, + _in_background(_report_ready_once_listening(child, port, events)), + ): + yield process + + +async def _run( + platform: Platform, + transition: Transition, + events: asyncio.Queue[Event], + children: Mapping[Child, asyncio.subprocess.Process], + launches: Mapping[Child, Launch], +) -> ExitCode: + state, actions = transition + if not actions: + if isinstance(state, Finished): + return state.code + return await _run(platform, step(state, await events.get()), events, children, launches) + action: Final = actions[0] + rest: Final = (state, actions[1:]) + match action: + case Signal(child, signum): + platform.signal_child(children[child], signum) + return await _run(platform, rest, events, children, launches) + case Start(child): + async with launches[child](events) as process: + return await _run(platform, rest, events, MappingProxyType({**children, child: process}), launches) + case _: + assert_never(action) + + +async def _supervise(platform: Platform, binary: str, command: Sequence[str], server: PythonServer) -> ExitCode: + events: Final[asyncio.Queue[Event]] = asyncio.Queue() + platform.catch_stop(asyncio.get_running_loop(), lambda signum: events.put_nowait(Stop(signum))) + with platform.private_listener() as private: + launches: Final[Mapping[Child, Launch]] = MappingProxyType( + { + Child.PYTHON: partial( + _launch_until_listening, + platform, + Child.PYTHON, + python_spec(platform, command, private, server), + private.port, + ), + Child.RUST: partial(_launch, platform, Child.RUST, rust_spec(platform, binary, private.port, server)), + } + ) + return await _run(platform, boot(), events, MappingProxyType({}), launches) + + +def _nothing_to_prepare() -> None: + pass + + +def supervise( + server: PythonServer, + command: Sequence[str], + before_start: Callable[[], None] = _nothing_to_prepare, + scripts_dir: str = sysconfig.get_path("scripts"), + platform: Platform | None = None, +) -> int: + """Serve ``server``'s public port from Rust and run ``command`` behind it, or say why that cannot start. + + ``before_start`` runs only once forwarding is known to be possible, so a refused start leaves nothing behind. + """ + chosen: Final = current_platform() if platform is None else platform + checked: Final = startup_check(chosen, server, gateway_binary(chosen, scripts_dir)) + if isinstance(checked, StartupError): + sys.stderr.write(f"{checked.reason}\n") + return ExitCode.CHILD_FAILED + before_start() + try: + return asyncio.run(_supervise(chosen, checked, command, server)) + except OSError as error: + sys.stderr.write(f"Rust gateway startup failed: {error}\n") + return ExitCode.CHILD_FAILED diff --git a/litellm/proxy/shutdown/graceful_uvicorn_worker.py b/litellm/proxy/shutdown/graceful_uvicorn_worker.py index 0bbbfe0a4a9..2836988b6ec 100644 --- a/litellm/proxy/shutdown/graceful_uvicorn_worker.py +++ b/litellm/proxy/shutdown/graceful_uvicorn_worker.py @@ -6,6 +6,8 @@ from litellm.constants import DEFAULT_MAX_REQUEST_HEADER_SIZE def build_graceful_shutdown_uvicorn_worker( timeout_graceful_shutdown: float, max_request_header_size: int | None = None, + *, + proxy_headers: bool = True, ) -> type[UvicornWorker]: class GracefulShutdownUvicornWorker(UvicornWorker): def __init__(self, *args: object) -> None: @@ -15,6 +17,7 @@ def build_graceful_shutdown_uvicorn_worker( "h11_max_incomplete_event_size": ( max_request_header_size if max_request_header_size is not None else DEFAULT_MAX_REQUEST_HEADER_SIZE ), + "proxy_headers": proxy_headers, } super().__init__(*args) diff --git a/litellm/rust_bridge/configuration.py b/litellm/rust_bridge/configuration.py index 152ba632996..5b51f155df9 100644 --- a/litellm/rust_bridge/configuration.py +++ b/litellm/rust_bridge/configuration.py @@ -8,7 +8,7 @@ from typing import Final from pydantic import TypeAdapter, ValidationError from typing_extensions import assert_never -_GLOBAL_ENV_NAME: Final = "LITELLM_RUST" +ENV_NAME: Final = "LITELLM_RUST" _ENV_BOOL: Final = TypeAdapter(bool) @@ -72,7 +72,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=_parse_env_bool(os.getenv(ENV_NAME)), ) diff --git a/litellm_cli/main.py b/litellm_cli/main.py index a9dfa4be179..ea170603d08 100644 --- a/litellm_cli/main.py +++ b/litellm_cli/main.py @@ -33,7 +33,7 @@ from contextlib import contextmanager from typing import TYPE_CHECKING, Final if TYPE_CHECKING: - from fastapi import FastAPI + from litellm.proxy.rust_forward.peer import ASGIApp SDK_PACKAGE: Final = "litellm" CLI_PACKAGE: Final = "litellm.proxy.client" @@ -214,10 +214,12 @@ def without_gc_pauses() -> Generator[None, None, None]: gc.enable() -def proxy_worker_app() -> "FastAPI": +def proxy_worker_app() -> "ASGIApp": with without_gc_pauses(): from litellm.proxy.proxy_cli import ProxyInitializationHelpers ProxyInitializationHelpers.apply_worker_log_config(os.environ) from litellm.proxy.proxy_server import app - return app + from litellm.proxy.rust_forward.mode import rust_forward_mode + from litellm.proxy.rust_forward.peer import wrap + return wrap(app, rust_forward_mode(os.environ)) diff --git a/tests/integration/configuration/test_rust_forward.py b/tests/integration/configuration/test_rust_forward.py new file mode 100644 index 00000000000..b75da286bf7 --- /dev/null +++ b/tests/integration/configuration/test_rust_forward.py @@ -0,0 +1,180 @@ +import json +import os +import subprocess +import sys +from pathlib import Path +from typing import Final + +import httpx +import pytest + +from tests.integration._support.client import Gateway, eventually +from tests.integration._support.ports import allocate_port +from tests.integration._support.process import stop_root_process + +TRUSTED_BALANCER: Final = "127.0.0.3" + + +def ready(client: httpx.Client) -> bool: + try: + return client.get("/health/liveliness").status_code == 200 + except httpx.TransportError: + return False + + +def client_ip(base_url: str, local_address: str, key: str, forwarded_for: str) -> httpx.Response: + with httpx.Client( + base_url=base_url, + transport=httpx.HTTPTransport(local_address=local_address), + trust_env=False, + timeout=5, + ) as client: + return client.get( + "/v1/mcp/network/client-ip", + headers={"Authorization": f"Bearer {key}", "X-Forwarded-For": forwarded_for}, + ) + + +@pytest.mark.parametrize("server_arguments", [(), ("--run_gunicorn",)]) +def test_gateway_preserves_metrics_bind_and_mcp_client_ip( + gateway: Gateway, tmp_path: Path, server_arguments: tuple[str, ...] +) -> None: + config: Final = tmp_path / "config.yaml" + config.write_text( + json.dumps( + { + "model_list": [], + "general_settings": {"master_key": "os.environ/LITELLM_MASTER_KEY", "use_x_forwarded_for": False}, + } + ) + ) + metrics_port: Final = allocate_port() + public_port: Final = allocate_port() + environment: Final = { + **{key: value for key, value in os.environ.items() if key not in ("DATABASE_URL", "DATABASE_URL_READ_REPLICA")}, + "LITELLM_MASTER_KEY": gateway.key, + "STORE_MODEL_IN_DB": "False", + "PYTHON_DOTENV_DISABLED": "1", + "LITELLM_RUST": "true", + "FORWARDED_ALLOW_IPS": TRUSTED_BALANCER, + } + with (tmp_path / "proxy.log").open("wb") as output: + with subprocess.Popen( + ( + sys.executable, + "-I", + "-m", + "litellm.proxy.proxy_cli", + "--config", + str(config), + "--host", + "0.0.0.0", + "--port", + str(public_port), + "--num_workers", + "2", + "--timeout_worker_healthcheck", + "120", + "--prometheus_metrics_port", + str(metrics_port), + *server_arguments, + ), + env=environment, + stdout=output, + stderr=output, + ) as supervisor: + try: + with httpx.Client( + base_url=f"http://127.0.0.1:{public_port}", + transport=httpx.HTTPTransport(local_address="127.0.0.2"), + trust_env=False, + follow_redirects=True, + timeout=5, + ) as external: + assert eventually(lambda: ready(external), bool, seconds=180), (tmp_path / "proxy.log").read_text() + response: Final = external.get(f"http://127.0.0.2:{metrics_port}/metrics") + assert response.status_code == 200, response.text + public: Final = f"http://127.0.0.1:{public_port}" + spoofed: Final = client_ip(public, "127.0.0.2", gateway.key, "1.2.3.4") + assert spoofed.status_code == 200, spoofed.text + assert spoofed.json() == {"ip": "127.0.0.2"} + behind_balancer: Final = client_ip(public, TRUSTED_BALANCER, gateway.key, "203.0.113.9") + assert behind_balancer.status_code == 200, behind_balancer.text + assert behind_balancer.json() == {"ip": "203.0.113.9"} + borrowed: Final = client_ip(public, "127.0.0.2", gateway.key, f"203.0.113.9, {TRUSTED_BALANCER}") + assert borrowed.status_code == 200, borrowed.text + assert borrowed.json() == {"ip": "127.0.0.2"} + finally: + assert stop_root_process(supervisor), (tmp_path / "proxy.log").read_text() + assert supervisor.returncode == 0, (tmp_path / "proxy.log").read_text() + + +@pytest.mark.parametrize( + "server_arguments", + [("--no-proxy-headers",), ("--forwarded-allow-ips", "127.0.0.1")], +) +def test_gateway_launcher_private_hop_preserves_client_address( + tmp_path: Path, server_arguments: tuple[str, ...] +) -> None: + config: Final = tmp_path / "config.yaml" + config.write_text( + json.dumps({"model_list": [], "general_settings": {"master_key": "os.environ/LITELLM_MASTER_KEY"}}) + ) + public_port: Final = allocate_port() + environment: Final = { + **{key: value for key, value in os.environ.items() if not key.startswith("DATABASE_") and key != "DIRECT_URL"}, + "LITELLM_MASTER_KEY": "sk-rust-forwarding-integration-key-not-a-secret", + "LITELLM_RUST": "true", + "LITELLM_PGBOUNCER_ENABLED": "false", + "CONFIG_FILE_PATH": str(config), + "STORE_MODEL_IN_DB": "False", + "PYTHON_DOTENV_DISABLED": "1", + } + log: Final = tmp_path / "gateway.log" + with log.open("wb") as output: + with subprocess.Popen( + ( + sys.executable, + "-I", + "-m", + "gateway.launch", + "--host", + "0.0.0.0", + "--port", + str(public_port), + "--workers", + "2", + "--timeout-worker-healthcheck", + "120", + *server_arguments, + ), + env=environment, + stdout=output, + stderr=output, + ) as supervisor: + try: + with httpx.Client( + base_url=f"http://127.0.0.1:{public_port}", + transport=httpx.HTTPTransport(local_address="127.0.0.2"), + trust_env=False, + timeout=5, + ) as external: + assert eventually(lambda: ready(external), bool, seconds=180), log.read_text() + response: Final = external.get( + "/health/liveliness?client-ip-proof", headers={"X-Forwarded-For": "198.51.100.10"} + ) + assert response.status_code == 200, response.text + events: Final = eventually( + lambda: tuple( + line + for line in log.read_text().splitlines() + if "GET /health/liveliness?client-ip-proof " in line + ), + bool, + seconds=10, + ) + assert len(events) == 1, log.read_text() + assert "127.0.0.2:" in events[0], events + finally: + assert stop_root_process(supervisor), log.read_text() + assert supervisor.returncode == 0, log.read_text() diff --git a/tests/integration/configuration/test_rust_forward_traffic.py b/tests/integration/configuration/test_rust_forward_traffic.py new file mode 100644 index 00000000000..eeeeea366d8 --- /dev/null +++ b/tests/integration/configuration/test_rust_forward_traffic.py @@ -0,0 +1,457 @@ +import asyncio +import json +import os +import socket +import subprocess +import sys +import threading +import time +from collections.abc import AsyncIterator, Callable, Iterator, Mapping +from contextlib import AbstractContextManager, ExitStack +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final + +import anyio +import httpx +import openai +import pytest +from pydantic import BaseModel, JsonValue +from starlette.applications import Starlette +from starlette.requests import Request as StarletteRequest +from starlette.responses import StreamingResponse +from starlette.routing import Route, WebSocketRoute +from starlette.websockets import WebSocket, WebSocketDisconnect +from websockets.sync.client import connect + +from tests.integration._support.asgi import asgi_server +from tests.integration._support.client import eventually +from tests.integration._support.mcp import mcp_peer +from tests.integration._support.ports import allocate_port +from tests.integration._support.process import stop_root_process +from tests.integration._support.wire import Reply, Request, wire_server + +MASTER_KEY: Final = "sk-rust-forward-traffic-key-not-a-secret" +KEEPALIVE_SECONDS: Final = 1 +FILES_MODEL: Final = "rust-forward-files" +FILE_ID: Final = "file-rust-forward" +CHAT_MODEL: Final = "rust-forward-chat" +REALTIME_MODEL: Final = "rust-forward-realtime" +MCP_HEADERS: Final = MappingProxyType( + { + "Authorization": f"Bearer {MASTER_KEY}", + "Accept": "application/json, text/event-stream", + "mcp-protocol-version": "2025-06-18", + } +) + + +@dataclass(frozen=True, slots=True) +class ScriptedProvider: + url: str + released: threading.Event + hung_up: SimpleQueue[str] + + +@dataclass(frozen=True, slots=True) +class RustProxy: + url: str + port: int + log: Path + provider: ScriptedProvider + + +class _Message(BaseModel): + content: str + + +class _ChatRequest(BaseModel): + messages: tuple[_Message, ...] + + +class _InputText(BaseModel): + text: str + + +class _Item(BaseModel): + content: tuple[_InputText, ...] = () + + +class _RealtimeFrame(BaseModel): + type: str + item: _Item | None = None + + +def _sse(content: str | None, finish_reason: str | None = None) -> bytes: + chunk: Final = { + "id": "chatcmpl-rust-forward", + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "delta": {} if content is None else {"content": content}, "finish_reason": finish_reason} + ], + } + return b"data: " + json.dumps(chunk).encode() + b"\n\n" + + +async def _gated_stream(released: threading.Event) -> AsyncIterator[bytes]: + yield _sse("before-gate") + await anyio.to_thread.run_sync(released.wait, 30) + yield _sse("after-gate") + yield _sse(None, "stop") + yield b"data: [DONE]\n\n" + + +async def _endless_stream(marker: str, hung_up: SimpleQueue[str]) -> AsyncIterator[bytes]: + try: + for index in range(600): + yield _sse(f"tick-{index} ") + await asyncio.sleep(0.1) + except (asyncio.CancelledError, GeneratorExit): + hung_up.put(marker) + raise + + +def _provider_app(released: threading.Event, hung_up: SimpleQueue[str]) -> Starlette: + async def chat(request: StarletteRequest) -> StreamingResponse: + prompt: Final = _ChatRequest.model_validate_json(await request.body()).messages[-1].content + return StreamingResponse( + _gated_stream(released) if prompt == "gated" else _endless_stream(prompt, hung_up), + media_type="text/event-stream", + ) + + async def realtime(websocket: WebSocket) -> None: + await websocket.accept() + await websocket.send_json({"type": "session.created", "session": {"id": "sess_rust_forward"}}) + try: + while True: + frame = _RealtimeFrame.model_validate_json(await websocket.receive_text()) # rebind-ok: per frame + if frame.item is not None: + text = "".join(part.text for part in frame.item.content) # rebind-ok: per frame + await websocket.send_json({"type": "response.output_text.delta", "delta": f"echo:{text}"}) + except WebSocketDisconnect: + return + + return Starlette( + routes=[ + Route("/v1/chat/completions", chat, methods=["POST"]), + WebSocketRoute("/v1/realtime", realtime), + ] + ) + + +def _routable_address() -> str: + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as probe: + probe.connect(("192.0.2.1", 9)) + address: Final[str] = probe.getsockname()[0] + assert not address.startswith("127."), address + return address + + +def _ready(client: httpx.Client) -> bool: + try: + return client.get("/health/liveliness").status_code == 200 + except httpx.TransportError: + return False + + +def _files_peer(content: bytes) -> Callable[[Request], Reply]: + file_object: Final = json.dumps( + { + "id": FILE_ID, + "object": "file", + "purpose": "batch", + "bytes": len(content), + "created_at": 1, + "filename": "upload.jsonl", + "status": "processed", + } + ).encode() + + def respond(request: Request) -> Reply: + if request.method == "POST" and request.target.startswith("/v1/files"): + return Reply(body=file_object) + if request.method == "GET" and request.target.startswith(f"/v1/files/{FILE_ID}/content"): + return Reply( + content_type="application/octet-stream", body=content, chunks=(content[:65536], content[65536:]) + ) + return Reply(status=404, body=b'{"error": {"message": "not found"}}') + + return respond + + +UPLOAD: Final = b"".join( + json.dumps( + { + "custom_id": f"row-{index}", + "method": "POST", + "url": "/v1/chat/completions", + "body": {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": os.urandom(48).hex()}]}, + } + ).encode() + + b"\n" + for index in range(30_000) +) + + +@pytest.fixture(scope="module") +def files_upstream() -> Iterator[str]: + with wire_server(_files_peer(UPLOAD), keep_alive=True) as wire: + yield wire.url + + +@pytest.fixture(scope="module") +def rust_proxy(tmp_path_factory: pytest.TempPathFactory, files_upstream: str) -> Iterator[RustProxy]: + directory: Final = tmp_path_factory.mktemp("rust-forward-traffic") + log: Final = directory / "proxy.log" + port: Final = allocate_port() + released: Final = threading.Event() + hung_up: Final[SimpleQueue[str]] = SimpleQueue() + with ExitStack() as stack: + mcp: Final = stack.enter_context(mcp_peer()) + provider: Final = ScriptedProvider( + stack.enter_context(asgi_server(_provider_app(released, hung_up))), released, hung_up + ) + config: Final = directory / "config.yaml" + config.write_text( + json.dumps( + { + "model_list": [ + { + "model_name": CHAT_MODEL, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": f"{provider.url}/v1", + "api_key": "integration-provider-key", + }, + }, + { + "model_name": REALTIME_MODEL, + "litellm_params": { + "model": "openai/gpt-realtime", + "api_base": provider.url, + "api_key": "integration-provider-key", + }, + "model_info": {"mode": "realtime"}, + }, + { + "model_name": FILES_MODEL, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": f"{files_upstream}/v1", + "api_key": "integration-provider-key", + }, + }, + ], + "mcp_servers": {"math": {**mcp.registration(), "available_on_public_internet": False}}, + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "mcp_internal_ip_ranges": ["127.0.0.0/8"], + }, + } + ) + ) + environment: Final = { + **{ + key: value + for key, value in os.environ.items() + if not key.startswith("DATABASE_") and key != "DIRECT_URL" + }, + "LITELLM_MASTER_KEY": MASTER_KEY, + "LITELLM_RUST": "true", + "STORE_MODEL_IN_DB": "False", + "PYTHON_DOTENV_DISABLED": "1", + } + output: Final = stack.enter_context(log.open("wb")) + supervisor: Final = stack.enter_context( + subprocess.Popen( + ( + sys.executable, + "-I", + "-m", + "litellm.proxy.proxy_cli", + "--config", + str(config), + "--host", + "0.0.0.0", + "--port", + str(port), + "--num_workers", + "2", + "--keepalive_timeout", + str(KEEPALIVE_SECONDS), + "--timeout_worker_healthcheck", + "120", + ), + env=environment, + stdout=output, + stderr=output, + ) + ) + try: + with httpx.Client(base_url=f"http://127.0.0.1:{port}", trust_env=False, timeout=5) as probe: + assert eventually(lambda: _ready(probe), bool, seconds=180), log.read_text() + yield RustProxy(f"http://127.0.0.1:{port}", port, log, provider) + finally: + released.set() + assert stop_root_process(supervisor), log.read_text() + assert supervisor.returncode == 0, log.read_text() + + +def _still_open(stream: httpx.Response) -> bool: + try: + next(stream.iter_raw()) + except httpx.ReadTimeout: + return True + except (StopIteration, httpx.TransportError): + return False + return True + + +def _initialize(client: httpx.Client, headers: Mapping[str, str]) -> httpx.Response: + return client.post( + "/mcp", + headers={**MCP_HEADERS, **headers}, + json={ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "protocolVersion": "2025-06-18", + "capabilities": {}, + "clientInfo": {"name": "rust-forward", "version": "1"}, + }, + }, + ) + + +def _mcp_post(client: httpx.Client, session: str, body: Mapping[str, JsonValue]) -> httpx.Response: + return client.post("/mcp", headers={**MCP_HEADERS, "mcp-session-id": session}, json=body) + + +def test_mcp_session_keeps_working_while_its_get_stream_sits_idle(rust_proxy: RustProxy) -> None: + idle_seconds: Final = 8 * KEEPALIVE_SECONDS + with httpx.Client(base_url=rust_proxy.url, trust_env=False, timeout=30) as client: + initialized: Final = _initialize(client, {}) + assert initialized.status_code == 200, initialized.text + session: Final = initialized.headers["mcp-session-id"] + assert _mcp_post(client, session, {"jsonrpc": "2.0", "method": "notifications/initialized"}).status_code == 202 + with httpx.Client(base_url=rust_proxy.url, trust_env=False, timeout=httpx.Timeout(5, read=2)) as listener: + with listener.stream( + "GET", "/mcp", headers={**MCP_HEADERS, "Accept": "text/event-stream", "mcp-session-id": session} + ) as notifications: + assert notifications.status_code == 200 + assert notifications.headers["content-type"].startswith("text/event-stream") + time.sleep(idle_seconds) + called: Final = _mcp_post( + client, + session, + { + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": {"name": "math-add", "arguments": {"a": 2, "b": 3}}, + }, + ) + assert called.status_code == 200, called.text + assert '"text":"5"' in called.text.replace(" ", ""), called.text + assert _still_open(notifications) + + +def test_uploaded_file_and_its_download_cross_byte_for_byte(rust_proxy: RustProxy) -> None: + with openai.OpenAI(api_key=MASTER_KEY, base_url=rust_proxy.url, max_retries=0, timeout=60) as client: + created: Final = client.files.create( + file=("upload.jsonl", UPLOAD), purpose="batch", extra_body={"model": FILES_MODEL} + ) + downloaded: Final = client.files.content(created.id, extra_query={"model": FILES_MODEL}).content + assert len(UPLOAD) > 4 * 1024 * 1024 + assert downloaded == UPLOAD, (len(downloaded), len(UPLOAD)) + + +def test_internal_only_mcp_server_judges_the_connecting_peer_not_its_forwarded_for_claim( + rust_proxy: RustProxy, +) -> None: + outsider: Final = _routable_address() + with httpx.Client( + base_url=f"http://{outsider}:{rust_proxy.port}", + transport=httpx.HTTPTransport(local_address=outsider), + trust_env=False, + timeout=10, + ) as spoofing: + spoofed: Final = _initialize(spoofing, {"X-Forwarded-For": "127.0.0.1"}) + with httpx.Client(base_url=rust_proxy.url, trust_env=False, timeout=10) as internal: + admitted: Final = _initialize(internal, {}) + assert admitted.status_code == 200, admitted.text + assert spoofed.status_code == 403, spoofed.text + assert "client IP" in spoofed.text, spoofed.text + + +def test_oauth_metadata_names_the_origin_the_client_used(rust_proxy: RustProxy) -> None: + with httpx.Client(base_url=rust_proxy.url, trust_env=False, timeout=10) as client: + direct: Final = client.get("/.well-known/oauth-protected-resource") + balanced: Final = client.get( + "/.well-known/oauth-protected-resource", + headers={"Host": "gateway.example.test", "X-Forwarded-Proto": "https"}, + ) + assert direct.status_code == 200, direct.text + assert direct.json()["resource"] == rust_proxy.url + assert direct.json()["authorization_servers"] == [f"{rust_proxy.url}/mcp"] + assert balanced.status_code == 200, balanced.text + assert balanced.json()["resource"] == "https://gateway.example.test" + + +def _chat_stream(client: httpx.Client, prompt: str) -> AbstractContextManager[httpx.Response]: + return client.stream( + "POST", + "/v1/chat/completions", + headers={"Authorization": f"Bearer {MASTER_KEY}"}, + json={"model": CHAT_MODEL, "stream": True, "messages": [{"role": "user", "content": prompt}]}, + ) + + +def _content(lines: Iterator[str]) -> str: + return next(line for line in lines if '"content"' in line) + + +def test_stream_chunk_reaches_the_client_before_the_provider_sends_the_next(rust_proxy: RustProxy) -> None: + with httpx.Client(base_url=rust_proxy.url, trust_env=False, timeout=httpx.Timeout(10, read=5)) as client: + with _chat_stream(client, "gated") as response: + assert response.status_code == 200, response.read() + lines: Final = response.iter_lines() + first: Final = _content(lines) + rust_proxy.provider.released.set() + rest: Final = tuple(line for line in lines if line) + assert "before-gate" in first, first + assert any("after-gate" in line for line in rest), rest + assert rest[-1] == "data: [DONE]", rest + + +def test_client_hang_up_mid_stream_closes_the_provider_stream(rust_proxy: RustProxy) -> None: + marker: Final = f"endless-{time.monotonic_ns()}" + with httpx.Client(base_url=rust_proxy.url, trust_env=False, timeout=httpx.Timeout(10, read=5)) as client: + with _chat_stream(client, marker) as response: + assert response.status_code == 200, response.read() + assert "tick-0" in _content(response.iter_lines()) + assert rust_proxy.provider.hung_up.get(timeout=15) == marker + + +def test_realtime_websocket_frames_cross_both_ways(rust_proxy: RustProxy) -> None: + marker: Final = f"realtime-{time.monotonic_ns()}" + with connect( + f"ws://127.0.0.1:{rust_proxy.port}/v1/realtime?model={REALTIME_MODEL}", + additional_headers={"Authorization": f"Bearer {MASTER_KEY}"}, + open_timeout=10, + ) as session: + session.send( + json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "message", "role": "user", "content": [{"type": "input_text", "text": marker}]}, + } + ) + ) + replies: Final = tuple(session.recv(timeout=10) for _ in range(2)) + assert any(f"echo:{marker}" in reply for reply in replies), replies diff --git a/tests/unit/gateway/test_launch.py b/tests/unit/gateway/test_launch.py index d682fcbad2d..50dab887688 100644 --- a/tests/unit/gateway/test_launch.py +++ b/tests/unit/gateway/test_launch.py @@ -12,13 +12,17 @@ from typing import Final, cast from unittest.mock import MagicMock, patch import pytest +import uvicorn from uvicorn.importer import import_from_string from uvicorn.main import main as uvicorn_main -from gateway.launch import GATEWAY_APP, main, pool_database_url, uvicorn_argv +from gateway.launch import GATEWAY_APP, main, pool_database_url, python_server, uvicorn_argv from litellm.proxy.db.db_url_settings import DatabaseURLSettings from litellm.proxy.db.pgbouncer import PGBOUNCER_POOLED_ENV_VAR, PgBouncerError, PgBouncerSettings from litellm.proxy.proxy_server import app as proxy_app +from litellm.proxy.rust_forward.mode import AssignedPort, InheritedFd, Supervised, uvicorn_flags +from litellm.proxy.rust_forward.server import Listener +from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager DB_ENV: Final = { "DATABASE_HOST": "db.internal", @@ -114,6 +118,100 @@ def _uvicorn_params(argv: tuple[str, ...]) -> dict[str, object]: class TestUvicornArgv: + @pytest.mark.parametrize( + ("arguments", "listener"), + [ + ((), Listener.TCP), + (("--uds", "/tmp/gateway.sock"), Listener.UNIX_SOCKET), + (("--fd", "0"), Listener.INHERITED_FD), + (("--fd=17",), Listener.INHERITED_FD), + (("--ssl-keyfile", "unused"), Listener.TLS), + (("--ssl-ciphers", "TLSv1.2"), Listener.TLS), + ], + ) + def test_forwarding_sees_the_public_listener_uvicorn_would_open( + self, arguments: tuple[str, ...], listener: Listener + ) -> None: + assert python_server(arguments).listener is listener + + def test_an_unsupported_listener_stops_forwarding_before_the_database_is_prepared( + self, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] + ) -> None: + monkeypatch.setenv("LITELLM_RUST", "true") + monkeypatch.delenv("_LITELLM_GATEWAY_FD", raising=False) + monkeypatch.setenv("LITELLM_PGBOUNCER_ENABLED", "true") + monkeypatch.setenv("LITELLM_PGBOUNCER_BINARY", "/nonexistent/pgbouncer") + with pytest.raises(SystemExit) as stopped: + main(("--uds", "/tmp/gateway.sock"), serve=lambda argv: None) + assert stopped.value.code == 1 + assert capsys.readouterr().err == "LITELLM_RUST does not support a Unix socket on the public port\n" + + @pytest.mark.parametrize( + ("argv", "trust"), + [ + ((), (True, None)), + (("--forwarded-allow-ips", "10.0.0.5"), (True, "10.0.0.5")), + (("--forwarded-allow-ips=*",), (True, "*")), + (("--no-proxy-headers",), (False, None)), + ], + ) + def test_forwarding_carries_operator_proxy_header_options_to_the_child( + self, argv: tuple[str, ...], trust: tuple[bool, str | None], monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.delenv("FORWARDED_ALLOW_IPS", raising=False) + server: Final = python_server(("--host", "0.0.0.0", "--port", "4000", *argv)) + assert (server.host, server.port) == ("0.0.0.0", 4000) + assert (server.proxy_headers, server.trusted_proxies) == trust + + @pytest.mark.parametrize( + ("argv", "environ"), + [(("--timeout-keep-alive", "2"), {}), ((), {"KEEPALIVE_TIMEOUT": "7"}), ((), {})], + ) + def test_rust_retires_connections_by_the_keepalive_uvicorn_runs_with( + self, argv: tuple[str, ...], environ: dict[str, str], monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.delenv("KEEPALIVE_TIMEOUT", raising=False) + for name, value in environ.items(): + monkeypatch.setenv(name, value) + served: Final = _uvicorn_params(uvicorn_argv(argv, os.environ))["timeout_keep_alive"] + assert python_server(argv).keepalive_seconds == served + + def test_without_a_keepalive_rust_follows_uvicorns_own_default(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("KEEPALIVE_TIMEOUT", raising=False) + assert python_server(()).keepalive_seconds == uvicorn.Config(app="", log_config=None).timeout_keep_alive + + @pytest.mark.parametrize("argv", [(), ("--timeout-graceful-shutdown", "12")]) + def test_rust_drains_for_pythons_graceful_timeout_plus_teardown( + self, argv: tuple[str, ...], monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "7.5") + flag: Final = _uvicorn_params(uvicorn_argv(argv, {}))["timeout_graceful_shutdown"] + graceful: Final = GracefulShutdownManager.get_timeout() if flag is None else float(str(flag)) + assert python_server(argv).graceful_timeout_seconds == graceful + + def test_supervised_listener_uses_the_inherited_loopback_socket(self): + supervised: Final = uvicorn_flags(Supervised(InheritedFd(17), True, "127.0.0.1")) + params: Final = _uvicorn_params((*uvicorn_argv(("--workers", "2"), {}), *supervised)) + assert params["fd"] == 17 + assert params["workers"] == 2 + + def test_assigned_port_listener_binds_only_loopback(self): + supervised: Final = uvicorn_flags(Supervised(AssignedPort(45556), True, "127.0.0.1")) + params: Final = _uvicorn_params((*uvicorn_argv(("--host", "0.0.0.0", "--port", "4000"), {}), *supervised)) + assert (params["host"], params["port"], params["fd"]) == ("127.0.0.1", 45556, None) + + @pytest.mark.parametrize("argv", [(), ("--proxy-headers",), ("--forwarded-allow-ips=*",)]) + def test_private_hop_leaves_proxy_headers_to_the_app(self, argv: tuple[str, ...]): + params: Final = _uvicorn_params( + (*uvicorn_argv(argv, {}), *uvicorn_flags(Supervised(InheritedFd(17), True, "*"))) + ) + assert params["proxy_headers"] is False + + def test_direct_listener_keeps_explicit_proxy_header_options(self): + params: Final = _uvicorn_params(uvicorn_argv(("--no-proxy-headers", "--forwarded-allow-ips=*"), {})) + assert params["proxy_headers"] is False + assert params["forwarded_allow_ips"] == "*" + def test_keepalive_env_reaches_uvicorn(self): params: Final = _uvicorn_params(uvicorn_argv(("--workers", "4"), {"KEEPALIVE_TIMEOUT": "75"})) assert params["app"] == GATEWAY_APP @@ -229,6 +327,16 @@ class TestMain: DatabaseURLSettings.from_env().apply_to_env() assert os.environ["DATABASE_URL"] == pooled + def test_forwarding_disabled_keeps_uvicorn_socket_options( + self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("LITELLM_RUST", "false") + monkeypatch.delenv("_LITELLM_GATEWAY_FD", raising=False) + monkeypatch.delenv("LITELLM_PGBOUNCER_ENABLED", raising=False) + served: Final[list[tuple[str, ...]]] = [] + main(("--uds", "/tmp/gateway.sock"), serve=lambda argv: served.append(tuple(argv))) + assert _uvicorn_params(served[0])["uds"] == "/tmp/gateway.sock" + def test_a_pooler_that_cannot_start_stops_the_gateway_before_uvicorn( self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch, tmp_path: Path ): diff --git a/tests/unit/proxy/rust_forward/__init__.py b/tests/unit/proxy/rust_forward/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/proxy/rust_forward/test_mode.py b/tests/unit/proxy/rust_forward/test_mode.py new file mode 100644 index 00000000000..9d57d4806c3 --- /dev/null +++ b/tests/unit/proxy/rust_forward/test_mode.py @@ -0,0 +1,112 @@ +from typing import Final + +import pytest +import uvicorn +from pydantic import ValidationError + +from litellm.proxy.proxy_cli import WORKER_APP, ProxyInitializationHelpers +from litellm.proxy.rust_forward.mode import ( + AssignedPort, + Disabled, + InheritedFd, + RustForwardMode, + Supervised, + Supervisor, + gunicorn_bind, + rust_forward_mode, + server_proxy_headers, + supervised_environment, + uvicorn_options, +) + + +@pytest.mark.parametrize( + ("environ", "expected"), + [ + ({}, Disabled()), + ({"LITELLM_RUST": "false"}, Disabled()), + ({"LITELLM_RUST": "off"}, Disabled()), + ({"LITELLM_RUST": "0"}, Disabled()), + ({"LITELLM_RUST": "true"}, Supervisor()), + ({"LITELLM_RUST": "1"}, Supervisor()), + ({"LITELLM_RUST": "yes"}, Supervisor()), + ({"LITELLM_RUST": "on"}, Supervisor()), + ({"LITELLM_RUST": " 1 "}, Supervisor()), + ({"LITELLM_RUST": ""}, Disabled()), + ({"LITELLM_RUST": " "}, Disabled()), + ( + {"LITELLM_RUST": "true", "_LITELLM_GATEWAY_FD": "17", "FORWARDED_ALLOW_IPS": "127.0.0.1,10.0.0.5"}, + Supervised(InheritedFd(17), True, "127.0.0.1,10.0.0.5"), + ), + ], +) +def test_mode_is_resolved_from_environment(environ: dict[str, str], expected: RustForwardMode) -> None: + assert rust_forward_mode(environ) == expected + + +@pytest.mark.parametrize( + ("environ", "variable"), + [ + ({"LITELLM_RUST": "enabled"}, "LITELLM_RUST"), + ({"LITELLM_RUST": "2"}, "LITELLM_RUST"), + ({"_LITELLM_GATEWAY_FD": "seventeen"}, "_LITELLM_GATEWAY_FD"), + ({"_LITELLM_GATEWAY_FD": "17", "_LITELLM_GATEWAY_PROXY_HEADERS": "maybe"}, "_LITELLM_GATEWAY_PROXY_HEADERS"), + ], +) +def test_invalid_forwarding_value_stops_startup_naming_the_variable(environ: dict[str, str], variable: str) -> None: + with pytest.raises(ValidationError, match=variable): + rust_forward_mode(environ) + + +@pytest.mark.parametrize( + ("proxy_headers", "trusted_proxies", "inherited", "expected"), + [ + (True, None, {}, Supervised(InheritedFd(17), True, "127.0.0.1")), + (True, None, {"FORWARDED_ALLOW_IPS": "10.0.0.5"}, Supervised(InheritedFd(17), True, "10.0.0.5")), + (True, "*", {"FORWARDED_ALLOW_IPS": "10.0.0.5"}, Supervised(InheritedFd(17), True, "*")), + (False, None, {}, Supervised(InheritedFd(17), False, "127.0.0.1")), + (True, None, {}, Supervised(AssignedPort(45555), True, "127.0.0.1")), + (False, "*", {}, Supervised(AssignedPort(45555), False, "*")), + ], +) +def test_child_reads_back_the_listener_and_trust_the_supervisor_wrote( + proxy_headers: bool, trusted_proxies: str | None, inherited: dict[str, str], expected: Supervised +) -> None: + written: Final = supervised_environment(expected.listener, proxy_headers, trusted_proxies) + assert rust_forward_mode({**inherited, **written}) == expected + + +def _served_uvicorn_config(mode: RustForwardMode) -> uvicorn.Config: + options: Final = { + **ProxyInitializationHelpers.get_default_unvicorn_init_args("0.0.0.0", 45555), + **uvicorn_options(mode, WORKER_APP), + } + return uvicorn.Config(**options) + + +def test_supervised_uvicorn_serves_the_inherited_socket_without_trusting_the_rust_hop() -> None: + config: Final = _served_uvicorn_config(Supervised(InheritedFd(17), True, "*")) + assert (config.fd, config.proxy_headers) == (17, False) + assert (config.app, config.factory) == (WORKER_APP, True) + + +@pytest.mark.parametrize("mode", [Disabled(), Supervisor()]) +def test_unsupervised_servers_keep_their_own_socket_and_proxy_headers(mode: RustForwardMode) -> None: + config: Final = _served_uvicorn_config(mode) + assert (config.fd, config.proxy_headers, config.port) == (None, True, 45555) + assert gunicorn_bind(mode, "0.0.0.0:4000") == "0.0.0.0:4000" + assert server_proxy_headers(mode) is True + + +def test_supervised_gunicorn_binds_the_inherited_socket_without_trusting_the_rust_hop() -> None: + mode: Final = Supervised(InheritedFd(17), True, "*") + assert gunicorn_bind(mode, "0.0.0.0:4000") == "fd://17" + assert server_proxy_headers(mode) is False + + +def test_assigned_port_serves_only_the_loopback_hop_without_trusting_it() -> None: + mode: Final = Supervised(AssignedPort(45556), True, "*") + config: Final = _served_uvicorn_config(mode) + assert (config.host, config.port, config.fd, config.proxy_headers) == ("127.0.0.1", 45556, None, False) + assert gunicorn_bind(mode, "0.0.0.0:4000") == "127.0.0.1:45556" + assert server_proxy_headers(mode) is False diff --git a/tests/unit/proxy/rust_forward/test_peer.py b/tests/unit/proxy/rust_forward/test_peer.py new file mode 100644 index 00000000000..ee625c4a4be --- /dev/null +++ b/tests/unit/proxy/rust_forward/test_peer.py @@ -0,0 +1,143 @@ +import asyncio +import gc +from typing import TYPE_CHECKING, Final + +import httpx +import pytest +from fastapi import Request +from fastapi.routing import APIRoute +from uvicorn.importer import import_from_string +from uvicorn.middleware.proxy_headers import ProxyHeadersMiddleware + +if TYPE_CHECKING: + from uvicorn._types import ( + ASGIReceiveCallable, + ASGISendCallable, + ASGISendEvent, + HTTPDisconnectEvent, + HTTPScope, + Scope, + ) + +from litellm.proxy.proxy_server import app as proxy_app +from litellm.proxy.rust_forward.mode import InheritedFd, Supervised +from litellm.proxy.rust_forward.peer import wrap + + +Headers = tuple[tuple[bytes, bytes], ...] + + +async def _receive() -> "HTTPDisconnectEvent": + return {"type": "http.disconnect"} + + +async def _send(message: "ASGISendEvent") -> None: + return None + + +async def _seen_by_app( + behind_rust: bool, client: tuple[str, int], headers: Headers, mode: Supervised +) -> tuple[tuple[str, int] | None, str, Headers]: + seen: Final[asyncio.Future[tuple[tuple[str, int] | None, str, Headers]]] = ( + asyncio.get_running_loop().create_future() + ) + + async def probe(scope: "Scope", receive: "ASGIReceiveCallable", send: "ASGISendCallable") -> None: + assert scope["type"] == "http" + seen.set_result((scope["client"], scope["scheme"], tuple(scope["headers"]))) + + direct: Final = ProxyHeadersMiddleware(probe, mode.trusted_proxies) if mode.proxy_headers else probe + app: Final = wrap(probe, mode) if behind_rust else direct + scope: Final[HTTPScope] = { + "type": "http", + "asgi": {"version": "3.0", "spec_version": "2.3"}, + "http_version": "1.1", + "method": "GET", + "scheme": "http", + "path": "/", + "raw_path": b"/", + "query_string": b"", + "root_path": "", + "headers": headers, + "client": client, + "server": ("127.0.0.1", 4000), + } + await app(scope, _receive, _send) + return seen.result() + + +PEER: Final = ("203.0.113.7", 51234) +TRUSTED_PEER: Final = ("10.0.12.34", 40000) +FORWARDED: Final[Headers] = ( + (b"x-forwarded-for", b"198.51.100.1, 203.0.113.7"), + (b"x-forwarded-proto", b"https"), +) + + +@pytest.mark.parametrize("peer", [PEER, TRUSTED_PEER, ("2001:db8::1", 443)]) +@pytest.mark.parametrize("headers", [(), FORWARDED, ((b"x-forwarded-for", b"10.1.2.3"), (b"host", b"public.example"))]) +@pytest.mark.parametrize( + "mode", + [ + Supervised(InheritedFd(17), True, "127.0.0.1"), + Supervised(InheritedFd(17), True, "10.0.0.0/8"), + Supervised(InheritedFd(17), True, "*"), + Supervised(InheritedFd(17), False, "10.0.0.0/8"), + ], +) +async def test_python_behind_rust_sees_what_a_direct_connection_would( + peer: tuple[str, int], headers: Headers, mode: Supervised +) -> None: + rendered: Final = f"[{peer[0]}]:{peer[1]}" if ":" in peer[0] else f"{peer[0]}:{peer[1]}" + behind_rust: Final = await _seen_by_app( + True, ("127.0.0.1", 60000), (*headers, (b"x-litellm-rust-peer", rendered.encode())), mode + ) + direct: Final = await _seen_by_app(False, peer, headers, mode) + assert behind_rust == direct + + +async def test_untrusted_peer_cannot_set_the_scheme_behind_rust() -> None: + client, scheme, _ = await _seen_by_app( + True, + ("127.0.0.1", 60000), + ((b"x-forwarded-proto", b"https"), (b"x-litellm-rust-peer", b"203.0.113.7:51234")), + Supervised(InheritedFd(17), True, "127.0.0.1"), + ) + assert (client, scheme) == (("203.0.113.7", 51234), "http") + + +async def test_peer_header_from_anyone_but_the_rust_hop_is_dropped_and_ignored() -> None: + client, _, headers = await _seen_by_app( + True, + ("192.0.2.9", 1000), + ((b"x-litellm-rust-peer", b"10.0.12.34:1"),), + Supervised(InheritedFd(17), False, "127.0.0.1"), + ) + assert (client, headers) == (("192.0.2.9", 1000), ()) + + +async def _peer_probe(request: Request) -> dict[str, object]: + assert request.client is not None + return {"client": [request.client.host, request.client.port], "peer": request.headers.get("x-litellm-rust-peer")} + + +@pytest.mark.parametrize("entry_point", ["litellm_cli.main:proxy_worker_app", "gateway.launch:worker_app"]) +async def test_every_entry_point_hands_routes_the_peer_rust_saw( + entry_point: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("_LITELLM_GATEWAY_FD", "17") + monkeypatch.setattr(proxy_app.router, "lifespan_context", proxy_app.router.lifespan_context) + monkeypatch.setattr( + proxy_app.router, "routes", [APIRoute("/__rust_forward_peer_probe", _peer_probe), *proxy_app.router.routes] + ) + try: + served: Final = import_from_string(entry_point)() + finally: + gc.unfreeze() + transport: Final = httpx.ASGITransport(app=served, client=("127.0.0.1", 60000)) + async with httpx.AsyncClient(transport=transport, base_url="http://proxy") as client: + response: Final = await client.get( + "/__rust_forward_peer_probe", headers={"x-litellm-rust-peer": "203.0.113.7:51234"} + ) + assert response.status_code == 200, response.text + assert response.json() == {"client": ["203.0.113.7", 51234], "peer": None} diff --git a/tests/unit/proxy/rust_forward/test_supervisor.py b/tests/unit/proxy/rust_forward/test_supervisor.py new file mode 100644 index 00000000000..d0ba6872bbd --- /dev/null +++ b/tests/unit/proxy/rust_forward/test_supervisor.py @@ -0,0 +1,241 @@ +import functools +import signal +import socket +import subprocess +import sys +from pathlib import Path +from typing import Final + +import pytest + +from litellm.constants import GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS +from litellm.proxy.proxy_cli import run_server +from litellm.proxy.rust_forward.mode import AssignedPort, InheritedFd, Supervised, rust_forward_mode +from litellm.proxy.rust_forward.platform import Platform, PrivateListener, current_platform +from litellm.proxy.rust_forward.platform.posix import PosixPlatform +from litellm.proxy.rust_forward.platform.windows import ( + CREATE_NEW_PROCESS_GROUP, + CTRL_BREAK_EVENT, + SIGBREAK, + WindowsPlatform, +) +from litellm.proxy.rust_forward.server import Listener, PythonServer, Server +from litellm.proxy.rust_forward.supervisor import ( + Child, + Event, + ExitCode, + Exited, + Finished, + Ready, + Running, + Signal, + Start, + StartupError, + Stop, + Stopping, + Transition, + boot, + gateway_binary, + python_spec, + rust_spec, + startup_check, + step, + supervise, +) + +TERM: Final = signal.SIGTERM +INT: Final = signal.SIGINT +POSIX: Final = PosixPlatform() +WINDOWS: Final = WindowsPlatform() +POSIX_ONLY: Final = pytest.mark.skipif(sys.platform == "win32", reason="runs a POSIX shebang gateway") + + +def _server( + server: Server = Server.UVICORN, listener: Listener = Listener.TCP, trusted_proxies: str | None = None +) -> PythonServer: + return PythonServer(server, listener, "0.0.0.0", 4000, 2, 2.5, False, trusted_proxies) + + +def test_python_child_serves_the_private_socket_with_the_operators_trust() -> None: + private: Final = PrivateListener(InheritedFd(17), 45555, (17,)) + spec: Final = python_spec(POSIX, ("python",), private, _server(trusted_proxies="10.0.0.5")) + assert rust_forward_mode(spec.env_overrides) == Supervised(InheritedFd(17), False, "10.0.0.5") + assert (spec.spawn.pass_fds, spec.spawn.start_new_session) == ((17,), True) + + +def test_windows_python_child_binds_the_assigned_port_in_its_own_process_group() -> None: + with WINDOWS.private_listener() as private: + pass + spec: Final = python_spec(WINDOWS, ("python",), private, _server()) + assert rust_forward_mode(spec.env_overrides) == Supervised(AssignedPort(private.port), False, "127.0.0.1") + assert spec.argv[-2:] == ("--port", str(private.port)) + assert (spec.spawn.pass_fds, spec.spawn.creationflags) == ((), CREATE_NEW_PROCESS_GROUP) + with socket.create_server(("127.0.0.1", private.port)): + pass + + +@pytest.mark.skipif(sys.platform != "win32", reason="these names exist only on Windows") +def test_windows_constants_match_the_stdlib() -> None: + assert (CTRL_BREAK_EVENT, SIGBREAK) == (signal.CTRL_BREAK_EVENT, signal.SIGBREAK) + assert CREATE_NEW_PROCESS_GROUP == subprocess.CREATE_NEW_PROCESS_GROUP + + +def test_python_child_keeps_the_public_host_and_its_options() -> None: + arguments: Final = ("--host", "0.0.0.0", "--port", "4000", "--num_workers", "2", "--reload") + spec: Final = python_spec(POSIX, arguments, PrivateListener(InheritedFd(17), 45555, (17,)), _server()) + params: Final = run_server.make_context("litellm", list(spec.argv)).params + assert (params["host"], params["port"]) == ("0.0.0.0", 45555) + assert (params["num_workers"], params["reload"]) == (2, True) + + +def test_rust_child_serves_public_port_and_forwards_to_private_listener() -> None: + spec: Final = rust_spec(POSIX, "gateway", 45555, _server()) + assert spec.argv == ("gateway",) + assert spec.spawn.pass_fds == () + assert dict(spec.env_overrides) == { + "LITELLM_FORWARD_UPSTREAM": "http://127.0.0.1:45555", + "LITELLM_FORWARD_HOST": "0.0.0.0", + "LITELLM_FORWARD_PORT": "4000", + "LITELLM_FORWARD_UPSTREAM_KEEPALIVE": "2", + "LITELLM_FORWARD_SHUTDOWN_DRAIN": str(2.5 + GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS), + } + + +@pytest.mark.parametrize( + ("server", "reason"), + [ + (_server(Server.HYPERCORN), "LITELLM_RUST does not support Hypercorn"), + (_server(Server.GRANIAN), "LITELLM_RUST does not support Granian"), + (_server(listener=Listener.TLS), "LITELLM_RUST does not support TLS on the public port"), + (_server(listener=Listener.UNIX_SOCKET), "LITELLM_RUST does not support a Unix socket on the public port"), + ( + _server(Server.GUNICORN, Listener.INHERITED_FD), + "LITELLM_RUST does not support an inherited file descriptor on the public port", + ), + ], +) +def test_unsupported_servers_are_refused_before_anything_starts( + server: PythonServer, reason: str, capsys: pytest.CaptureFixture[str] +) -> None: + assert startup_check(POSIX, server, "gateway") == StartupError(reason) + prepared: Final[list[bool]] = [] + refused: Final = supervise(server, ("python",), before_start=lambda: prepared.append(True), platform=POSIX) + assert refused == ExitCode.CHILD_FAILED + assert (capsys.readouterr().err, prepared) == (f"{reason}\n", []) + + +@pytest.mark.parametrize("server", [Server.UVICORN, Server.GUNICORN]) +def test_plain_tcp_uvicorn_and_gunicorn_can_be_forwarded(server: Server) -> None: + assert startup_check(POSIX, _server(server), "gateway") == "gateway" + + +def test_windows_forwards_uvicorn_and_refuses_gunicorn() -> None: + assert startup_check(WINDOWS, _server(), "gateway") == "gateway" + assert startup_check(WINDOWS, _server(Server.GUNICORN), "gateway") == StartupError( + "LITELLM_RUST does not support gunicorn on Windows" + ) + + +def test_missing_binary_gives_install_instruction() -> None: + assert startup_check(POSIX, _server(), None) == StartupError( + "LITELLM_RUST requires the litellm-gateway binary. Install litellm[proxy] in this environment" + ) + + +@POSIX_ONLY +def test_binary_is_found_only_in_the_environment_scripts_dir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + scripts: Final = tmp_path / "scripts" + elsewhere: Final = tmp_path / "elsewhere" + scripts.mkdir() + elsewhere.mkdir() + binary: Final = elsewhere / "litellm-gateway" + binary.write_text("#!/bin/sh\nexit 0\n") + binary.chmod(0o755) + monkeypatch.setenv("PATH", str(elsewhere)) + assert gateway_binary(POSIX, str(scripts)) is None + assert gateway_binary(POSIX, str(elsewhere)) == str(binary) + + +PYTHON_CHILD: Final = """ +import os, pathlib, socket, time +time.sleep(0.5) +pathlib.Path(os.environ["MARKERS"], "listening").touch() +fd = os.environ.get("_LITELLM_GATEWAY_FD") +port = os.environ.get("_LITELLM_GATEWAY_PORT") +listener = socket.socket(fileno=int(fd)) if fd else socket.create_server(("127.0.0.1", int(port))) +listener.listen() +time.sleep(10) +""" +RUST_CHILD: Final = """ +import os, pathlib +markers = pathlib.Path(os.environ["MARKERS"]) +(markers / "rust_saw_python_listening").write_text(str((markers / "listening").exists())) +""" + + +def _fake_gateway(scripts: Path) -> Platform: + if sys.platform != "win32": + gateway: Final = scripts / "litellm-gateway" + gateway.write_text(f"#!{sys.executable} -I\n{RUST_CHILD}") + gateway.chmod(0o755) + return current_platform() + script: Final = scripts / "rust_child.py" + script.write_text(RUST_CHILD) + (scripts / "litellm-gateway.cmd").write_text(f'@"{sys.executable}" -I "{script}"\r\n') + return WindowsPlatform(gateway_name="litellm-gateway.cmd") + + +def test_rust_opens_the_public_port_only_once_python_listens(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("MARKERS", str(tmp_path)) + command: Final = (sys.executable, "-I", "-c", PYTHON_CHILD) + supervise(_server(), command, scripts_dir=str(tmp_path), platform=_fake_gateway(tmp_path)) + assert (tmp_path / "rust_saw_python_listening").read_text() == "True" + + +def _fold(events: tuple[Event, ...]) -> Transition: + def advance(acc: Transition, event: Event) -> Transition: + state, done = acc + following, actions = step(state, event) + return following, (*done, *actions) + + return functools.reduce(advance, events, boot()) + + +RUST: Final = Child.RUST +PYTHON: Final = Child.PYTHON +PYTHON_READY: Final = Ready(PYTHON) +STARTED: Final = (Start(PYTHON), Start(RUST)) +BOTH_TERM: Final = (*STARTED, Signal(RUST, TERM), Signal(PYTHON, TERM)) +CLEAN: Final = Finished(ExitCode.CLEAN) +FAILED: Final = Finished(ExitCode.CHILD_FAILED) + + +@pytest.mark.parametrize( + ("events", "expected"), + [ + ((), (Running((PYTHON,)), (Start(PYTHON),))), + ((Ready(RUST),), (Running((PYTHON,)), (Start(PYTHON),))), + ((PYTHON_READY,), (Running((PYTHON, RUST)), STARTED)), + ((PYTHON_READY, PYTHON_READY, Ready(RUST)), (Running((PYTHON, RUST)), STARTED)), + ((PYTHON_READY, Stop(TERM)), (Stopping(frozenset({RUST, PYTHON})), BOTH_TERM)), + ((PYTHON_READY, Stop(TERM), Exited(RUST, 0), Exited(PYTHON, 0)), (CLEAN, BOTH_TERM)), + ((PYTHON_READY, Stop(TERM), Exited(PYTHON, 0), Exited(RUST, 0)), (CLEAN, BOTH_TERM)), + ( + (PYTHON_READY, Stop(INT), Stop(INT), Exited(RUST, 0), Exited(PYTHON, 0)), + (CLEAN, (*STARTED, Signal(RUST, INT), Signal(PYTHON, INT))), + ), + ((PYTHON_READY, Stop(TERM), Exited(RUST, 0), Exited(RUST, 0)), (FAILED, BOTH_TERM)), + ((PYTHON_READY, Stop(TERM), Exited(RUST, 1)), (FAILED, BOTH_TERM)), + ((PYTHON_READY, Stop(TERM), Exited(RUST, 0), Exited(PYTHON, 1)), (FAILED, BOTH_TERM)), + ((PYTHON_READY, Exited(PYTHON, 0)), (FAILED, STARTED)), + ((PYTHON_READY, Exited(RUST, 1)), (FAILED, STARTED)), + ((PYTHON_READY, Exited(RUST, 1), Stop(TERM), PYTHON_READY, Exited(PYTHON, 0)), (FAILED, STARTED)), + ((Exited(PYTHON, 1),), (FAILED, (Start(PYTHON),))), + ((Stop(TERM), PYTHON_READY, Exited(PYTHON, 0)), (CLEAN, (Start(PYTHON), Signal(PYTHON, TERM)))), + ((Stop(TERM), Exited(PYTHON, 1)), (FAILED, (Start(PYTHON), Signal(PYTHON, TERM)))), + ], +) +def test_children_start_in_order_once_ready_and_stop_together_in_reverse( + events: tuple[Event, ...], expected: Transition +) -> None: + assert _fold(events) == expected diff --git a/tests/unit/proxy/test_dev_reload.py b/tests/unit/proxy/test_dev_reload.py index 60a7c53820b..6650c1cd3e3 100644 --- a/tests/unit/proxy/test_dev_reload.py +++ b/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) diff --git a/tests/unit/proxy/test_proxy_cli.py b/tests/unit/proxy/test_proxy_cli.py index 6347c165865..fe0e46caaf3 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -201,42 +201,30 @@ class TestProxyInitializationHelpers: # Execute with pytest.raises(ValueError, match="Invalid test value"): - ProxyInitializationHelpers._run_test_chat_completion( - "localhost", 8000, "gpt-3.5-turbo", True - ) + ProxyInitializationHelpers._run_test_chat_completion("localhost", 8000, "gpt-3.5-turbo", True) # Test with valid string test value - ProxyInitializationHelpers._run_test_chat_completion( - "localhost", 8000, "gpt-3.5-turbo", "http://test-url" - ) + ProxyInitializationHelpers._run_test_chat_completion("localhost", 8000, "gpt-3.5-turbo", "http://test-url") # Assert - mock_openai.assert_called_once_with( - api_key="My API Key", base_url="http://test-url" - ) + mock_openai.assert_called_once_with(api_key="My API Key", base_url="http://test-url") mock_client.chat.completions.create.assert_called() def test_get_default_unvicorn_init_args(self): # Test without log_config - args = ProxyInitializationHelpers._get_default_unvicorn_init_args( - "localhost", 8000 - ) + args = ProxyInitializationHelpers._get_default_unvicorn_init_args("localhost", 8000) assert args["app"] == "litellm.proxy.proxy_server:app" assert args["host"] == "localhost" assert args["port"] == 8000 assert args["server_header"] is False # Test with log_config - args = ProxyInitializationHelpers._get_default_unvicorn_init_args( - "localhost", 8000, "log_config.json" - ) + args = ProxyInitializationHelpers._get_default_unvicorn_init_args("localhost", 8000, "log_config.json") assert args["log_config"] == "log_config.json" # Test with json_logs=True with patch("litellm.json_logs", True): - args = ProxyInitializationHelpers._get_default_unvicorn_init_args( - "localhost", 8000 - ) + args = ProxyInitializationHelpers._get_default_unvicorn_init_args("localhost", 8000) # When json_logs is True, log_config should be set to the JSON log config dict assert args["log_config"] is not None assert isinstance(args["log_config"], dict) @@ -244,15 +232,11 @@ class TestProxyInitializationHelpers: assert "formatters" in args["log_config"] # Test with keepalive_timeout - args = ProxyInitializationHelpers._get_default_unvicorn_init_args( - "localhost", 8000, None, 60 - ) + args = ProxyInitializationHelpers._get_default_unvicorn_init_args("localhost", 8000, None, 60) assert args["timeout_keep_alive"] == 60 # Test with both log_config and keepalive_timeout - args = ProxyInitializationHelpers._get_default_unvicorn_init_args( - "localhost", 8000, "log_config.json", 120 - ) + args = ProxyInitializationHelpers._get_default_unvicorn_init_args("localhost", 8000, "log_config.json", 120) assert args["log_config"] == "log_config.json" assert args["timeout_keep_alive"] == 120 @@ -383,14 +367,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 = 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 = 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 @@ -407,14 +387,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 = 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 = 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 @@ -426,10 +402,7 @@ class TestProxyInitializationHelpers: StatReload._litellm_patched_config_paths.clear() assert ProxyInitializationHelpers._patch_statreload_extra_paths([]) is False - assert ( - ProxyInitializationHelpers._patch_statreload_extra_paths([None, ""]) - is False - ) + assert ProxyInitializationHelpers._patch_statreload_extra_paths([None, ""]) is False def test_patch_statreload_extra_paths_is_idempotent(self, tmp_path): from pathlib import Path @@ -447,18 +420,14 @@ 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 = 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} assert config_file.resolve() in yielded_paths assert py_file.resolve() in yielded_paths - def test_configure_dev_reload_watches_env_and_sets_override_flag( - self, tmp_path, monkeypatch - ): + def test_configure_dev_reload_watches_env_and_sets_override_flag(self, tmp_path, monkeypatch): from pathlib import Path from uvicorn.supervisors.statreload import StatReload @@ -475,9 +444,7 @@ class TestProxyInitializationHelpers: uvicorn_args: dict = {} with patch("litellm._logging.verbose_proxy_logger.warning") as mock_warning: - ProxyInitializationHelpers._configure_dev_reload( - uvicorn_args, str(config_file) - ) + ProxyInitializationHelpers._configure_dev_reload(uvicorn_args, str(config_file)) assert os.environ["LITELLM_DEV_ENV_HOT_RELOAD"] == "True" assert uvicorn_args["reload"] is True @@ -488,9 +455,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 = 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 @@ -514,17 +479,13 @@ class TestProxyInitializationHelpers: mock_app = MagicMock() # Execute - ProxyInitializationHelpers._init_hypercorn_server( - mock_app, "localhost", 8000, None, None, None - ) + ProxyInitializationHelpers._init_hypercorn_server(mock_app, "localhost", 8000, None, None, None) # Assert mock_asyncio_run.assert_called_once() # Test with SSL - ProxyInitializationHelpers._init_hypercorn_server( - mock_app, "localhost", 8000, "cert.pem", "key.pem", "ECDHE" - ) + ProxyInitializationHelpers._init_hypercorn_server(mock_app, "localhost", 8000, "cert.pem", "key.pem", "ECDHE") @patch("granian.Granian") @patch("builtins.print") @@ -693,10 +654,14 @@ class TestProxyInitializationHelpers: database_name_enc = urllib.parse.quote_plus(database_name) # Construct DATABASE_URL from the provided variables - database_url = f"postgresql://{database_username_enc}:{database_password_enc}@{database_host}/{database_name_enc}" + database_url = ( + 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 = "postgresql://user%40with%2Bspecial:test-password-special-chars@localhost:5432/db_name%2Ftest" + expected_url = ( + "postgresql://user%40with%2Bspecial:test-password-special-chars@localhost:5432/db_name%2Ftest" + ) assert database_url == expected_url # Test appending query parameters @@ -714,9 +679,7 @@ class TestProxyInitializationHelpers: @patch("uvicorn.run") @patch("atexit.register") # critical @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False - ) + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) def test_skip_server_startup( self, mock_should_update, mock_setup_db, mock_atexit_register, mock_uvicorn_run, tmp_path: Path ): @@ -734,11 +697,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 = {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") @@ -758,9 +717,7 @@ class TestProxyInitializationHelpers: "litellm.proxy.proxy_server": mock_proxy_module, }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -771,18 +728,12 @@ class TestProxyInitializationHelpers: # --- skip startup --- result = runner.invoke(run_server, ["--local", "--skip_server_startup"]) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" assert "Skipping server startup" in result.output mock_uvicorn_run.assert_not_called() - result = runner.invoke( - run_server, ["--local", "--skip_server_startup", "--telemetry", "False"] - ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + result = runner.invoke(run_server, ["--local", "--skip_server_startup", "--telemetry", "False"]) + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" assert "Skipping server startup" in result.output assert "telemetry" not in runner.invoke(run_server, ["--help"]).output assert live_proxy_samples.exists() @@ -792,9 +743,7 @@ class TestProxyInitializationHelpers: result = runner.invoke(run_server, ["--local"]) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_uvicorn_run.assert_called_once() assert not live_proxy_samples.exists() @@ -831,9 +780,7 @@ class TestProxyInitializationHelpers: @patch("uvicorn.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False - ) + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) def test_limit_concurrency_passed_to_uvicorn( self, mock_should_update, mock_setup_db, mock_atexit_register, mock_uvicorn_run ): @@ -850,11 +797,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 = {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( @@ -864,9 +807,7 @@ class TestProxyInitializationHelpers: "litellm.proxy.proxy_server": mock_proxy_module, }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.side_effect = lambda *a, **k: { "app": "litellm.proxy.proxy_server:app", @@ -874,20 +815,14 @@ class TestProxyInitializationHelpers: "port": 8000, } - result = runner.invoke( - run_server, ["--local", "--limit_concurrency", "250"] - ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + result = runner.invoke(run_server, ["--local", "--limit_concurrency", "250"]) + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_uvicorn_run.assert_called_once() assert mock_uvicorn_run.call_args.kwargs.get("limit_concurrency") == 250 mock_uvicorn_run.reset_mock() result = runner.invoke(run_server, ["--local"]) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_uvicorn_run.assert_called_once() assert "limit_concurrency" not in mock_uvicorn_run.call_args.kwargs @@ -905,7 +840,9 @@ class TestProxyInitializationHelpers: @patch("httpx.HTTPTransport.handle_request") @patch("atexit.register") @patch("subprocess.Popen") - @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above + @patch( + "litellm.proxy.db.prisma_client.PrismaManager.setup_database" + ) # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above @patch( # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False ) @@ -943,9 +880,7 @@ class TestProxyInitializationHelpers: save_worker_config=MagicMock(), ) clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL", "PROMETHEUS_METRICS_PORT") + k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL", "PROMETHEUS_METRICS_PORT") } clean_env["PROMETHEUS_MULTIPROC_DIR"] = str(tmp_path) with ( @@ -971,9 +906,7 @@ class TestProxyInitializationHelpers: run_server, ["--local", "--host", "127.0.0.1", "--port", "4000", "--prometheus_metrics_port", "4001"], ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_popen.assert_called_once() spawned = list(mock_popen.call_args.args[0]) assert spawned[1:3] == ["-m", "litellm.proxy.prometheus_metrics_server"] @@ -996,9 +929,7 @@ class TestProxyInitializationHelpers: mock_popen.reset_mock() mock_uvicorn_run.reset_mock() result = runner.invoke(run_server, ["--local"]) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_popen.assert_not_called() mock_uvicorn_run.assert_called_once() @@ -1012,9 +943,7 @@ class TestProxyInitializationHelpers: mock_popen.assert_not_called() mock_uvicorn_run.assert_not_called() - result = runner.invoke( - run_server, ["--local", "--prometheus_metrics_port", "0"] - ) + result = runner.invoke(run_server, ["--local", "--prometheus_metrics_port", "0"]) assert result.exit_code == 2 assert "Invalid value for '--prometheus_metrics_port'" in result.output mock_popen.assert_not_called() @@ -1036,9 +965,7 @@ class TestProxyInitializationHelpers: @patch("subprocess.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False - ) + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) def test_db_timeout_settings_are_forwarded_to_pool_timeout( self, mock_should_update, @@ -1071,11 +998,7 @@ class TestProxyInitializationHelpers: } ) - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {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), @@ -1086,9 +1009,7 @@ class TestProxyInitializationHelpers: "litellm.proxy.proxy_server": mock_proxy_module, }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, patch( "litellm.proxy.proxy_cli.append_query_params", side_effect=lambda url, params: ( @@ -1107,9 +1028,7 @@ class TestProxyInitializationHelpers: ["--local", "--config", "test-config.yaml", "--skip_server_startup"], ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_append_query_params.assert_called() appended_params = mock_append_query_params.call_args.args[1] assert appended_params["connection_limit"] == 5 @@ -1164,9 +1083,7 @@ class TestProxyInitializationHelpers: @patch("subprocess.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False - ) + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) def test_db_connection_extra_params_forwarded_to_url( self, mock_should_update, @@ -1201,11 +1118,7 @@ class TestProxyInitializationHelpers: } ) - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {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), @@ -1216,9 +1129,7 @@ class TestProxyInitializationHelpers: "litellm.proxy.proxy_server": mock_proxy_module, }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, patch( "litellm.proxy.proxy_cli.append_query_params", side_effect=lambda url, params: str(url), @@ -1235,9 +1146,7 @@ class TestProxyInitializationHelpers: ["--local", "--config", "test-config.yaml", "--skip_server_startup"], ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_append_query_params.assert_called() appended_params = mock_append_query_params.call_args.args[1] assert appended_params["connect_timeout"] == 15 @@ -1288,9 +1197,7 @@ class TestProxyInitializationHelpers: @patch("subprocess.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False - ) + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) def test_disable_prepared_statements_forwarded_to_url( self, mock_should_update, @@ -1322,11 +1229,7 @@ class TestProxyInitializationHelpers: } ) - clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL") - } + clean_env = {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), @@ -1337,9 +1240,7 @@ class TestProxyInitializationHelpers: "litellm.proxy.proxy_server": mock_proxy_module, }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, patch( "litellm.proxy.proxy_cli.append_query_params", side_effect=lambda url, params: str(url), @@ -1356,9 +1257,7 @@ class TestProxyInitializationHelpers: ["--local", "--config", "test-config.yaml", "--skip_server_startup"], ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_append_query_params.assert_called() appended_params = mock_append_query_params.call_args.args[1] if expect_pgbouncer: @@ -1380,9 +1279,7 @@ class TestProxyInitializationHelpers: @patch("subprocess.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False - ) + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) def test_disable_prepared_statements_env_var_forwarded_to_url( self, mock_should_update, @@ -1430,9 +1327,7 @@ class TestProxyInitializationHelpers: "litellm.proxy.proxy_server": mock_proxy_module, }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, patch( "litellm.proxy.proxy_cli.append_query_params", side_effect=lambda url, params: str(url), @@ -1449,9 +1344,7 @@ class TestProxyInitializationHelpers: ["--local", "--config", "test-config.yaml", "--skip_server_startup"], ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" appended_params = mock_append_query_params.call_args.args[1] if expect_pgbouncer: assert appended_params["pgbouncer"] == "true", appended_params @@ -1461,9 +1354,7 @@ class TestProxyInitializationHelpers: @patch("subprocess.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False - ) + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) def test_malformed_disable_prepared_statements_env_var_is_rejected_even_when_config_enables_it( self, mock_should_update, @@ -1509,9 +1400,7 @@ class TestProxyInitializationHelpers: "litellm.proxy.proxy_server": mock_proxy_module, }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, patch( "litellm.proxy.proxy_cli.append_query_params", side_effect=lambda url, params: str(url), @@ -1535,9 +1424,7 @@ class TestProxyInitializationHelpers: @patch("uvicorn.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False - ) + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) def test_proxy_default_api_version_uses_azure_default( self, mock_should_update, mock_setup_db, mock_atexit_register, mock_uvicorn_run ): @@ -1554,11 +1441,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 = {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( @@ -1568,9 +1451,7 @@ class TestProxyInitializationHelpers: "litellm.proxy.proxy_server": mock_proxy_module, }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -1578,9 +1459,7 @@ class TestProxyInitializationHelpers: "port": 8000, } result = runner.invoke(run_server, ["--local", "--skip_server_startup"]) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_proxy_module.save_worker_config.assert_called_once() call_kwargs = mock_proxy_module.save_worker_config.call_args[1] assert call_kwargs["api_version"] == litellm.AZURE_DEFAULT_API_VERSION @@ -1588,12 +1467,8 @@ class TestProxyInitializationHelpers: @patch("uvicorn.run") @patch("builtins.print") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False - ) - def test_keepalive_timeout_flag( - self, mock_should_update, mock_setup_db, mock_print, mock_uvicorn_run - ): + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) + def test_keepalive_timeout_flag(self, mock_should_update, mock_setup_db, mock_print, mock_uvicorn_run): """Test that the keepalive_timeout flag is properly passed to uvicorn""" from click.testing import CliRunner @@ -1610,11 +1485,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 = {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), @@ -1629,9 +1500,7 @@ class TestProxyInitializationHelpers: ) }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, patch( "litellm.proxy.proxy_cli.ProxyInitializationHelpers.is_port_in_use", return_value=False, @@ -1664,12 +1533,8 @@ class TestProxyInitializationHelpers: @patch("uvicorn.run") @patch("builtins.print") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - @patch( - "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False - ) - def test_timeout_worker_healthcheck_flag( - self, mock_should_update, mock_setup_db, mock_print, mock_uvicorn_run - ): + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) + def test_timeout_worker_healthcheck_flag(self, mock_should_update, mock_setup_db, mock_print, mock_uvicorn_run): """Test that the --timeout_worker_healthcheck flag is threaded through to the uvicorn init helper.""" from click.testing import CliRunner @@ -1686,11 +1551,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 = {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), @@ -1705,9 +1566,7 @@ class TestProxyInitializationHelpers: ) }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, patch( "litellm.proxy.proxy_cli.ProxyInitializationHelpers.is_port_in_use", return_value=False, @@ -1719,9 +1578,7 @@ class TestProxyInitializationHelpers: "port": 8000, } - result = runner.invoke( - run_server, ["--local", "--timeout_worker_healthcheck", "15"] - ) + result = runner.invoke(run_server, ["--local", "--timeout_worker_healthcheck", "15"]) assert result.exit_code == 0 mock_get_args.assert_called_once_with( @@ -1736,9 +1593,7 @@ class TestProxyInitializationHelpers: @patch("uvicorn.run") @patch("builtins.print") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - def test_max_requests_before_restart_flag( - self, mock_setup_db, mock_print, mock_uvicorn_run - ): + def test_max_requests_before_restart_flag(self, mock_setup_db, mock_print, mock_uvicorn_run): """Test that the max_requests_before_restart flag is passed to uvicorn as limit_max_requests""" from click.testing import CliRunner @@ -1751,11 +1606,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 = {k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL")} with ( patch.dict( os.environ, @@ -1773,9 +1624,7 @@ class TestProxyInitializationHelpers: ) }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -1783,13 +1632,9 @@ class TestProxyInitializationHelpers: "port": 8000, } - result = runner.invoke( - run_server, ["--local", "--max_requests_before_restart", "123"] - ) + result = 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 result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_uvicorn_run.assert_called_once() # Check that uvicorn.run was called with limit_max_requests parameter @@ -1799,9 +1644,7 @@ class TestProxyInitializationHelpers: @patch("uvicorn.run") @patch("builtins.print") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - def test_max_requests_before_restart_jitter_flag( - self, mock_setup_db, mock_print, mock_uvicorn_run - ): + def test_max_requests_before_restart_jitter_flag(self, mock_setup_db, mock_print, mock_uvicorn_run): """--max_requests_before_restart_jitter maps to uvicorn limit_max_requests_jitter""" from click.testing import CliRunner @@ -1812,11 +1655,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 = {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), @@ -1831,9 +1670,7 @@ class TestProxyInitializationHelpers: ) }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -1852,9 +1689,7 @@ class TestProxyInitializationHelpers: ], ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_uvicorn_run.assert_called_once() call_args = mock_uvicorn_run.call_args assert call_args[1]["limit_max_requests"] == 1000 @@ -1873,11 +1708,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 = {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( @@ -1891,9 +1722,7 @@ class TestProxyInitializationHelpers: ) }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -1913,9 +1742,7 @@ class TestProxyInitializationHelpers: ], ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_uvicorn_run.assert_not_called() mock_run_gunicorn.assert_called_once() g_kwargs = mock_run_gunicorn.call_args[1] @@ -1961,20 +1788,14 @@ class TestProxyInitializationHelpers: @patch("uvicorn.run") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") - def test_single_worker_max_requests_warning_printed_on_uvicorn( - self, mock_setup_db, mock_uvicorn_run - ): + def test_single_worker_max_requests_warning_printed_on_uvicorn(self, mock_setup_db, mock_uvicorn_run): """uvicorn single-worker startup prints the does-not-recycle warning""" from click.testing import CliRunner 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 = {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( @@ -1988,9 +1809,7 @@ class TestProxyInitializationHelpers: ) }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -1998,13 +1817,9 @@ class TestProxyInitializationHelpers: "port": 8000, } - result = runner.invoke( - run_server, ["--local", "--max_requests_before_restart", "123"] - ) + result = 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 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 assert "uvicorn exits the process" in result.output @@ -2021,11 +1836,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 = {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( @@ -2039,9 +1850,7 @@ class TestProxyInitializationHelpers: ) }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -2054,12 +1863,8 @@ class TestProxyInitializationHelpers: ["--local", "--num_workers", "2", "--max_requests_before_restart", "123"], ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" - assert not any( - "does not recycle a worker" in str(c) for c in mock_print.call_args_list - ) + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" + assert not any("does not recycle a worker" in str(c) for c in mock_print.call_args_list) @patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.run_gunicorn_server") @patch("uvicorn.run") @@ -2073,11 +1878,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 = {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( @@ -2091,9 +1892,7 @@ class TestProxyInitializationHelpers: ) }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -2106,9 +1905,7 @@ class TestProxyInitializationHelpers: ["--local", "--run_gunicorn", "--max_requests_before_restart", "123"], ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + 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 assert "gunicorn accepts no connections" in result.output @@ -2125,11 +1922,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 = {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( @@ -2143,9 +1936,7 @@ class TestProxyInitializationHelpers: ) }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -2158,13 +1949,9 @@ class TestProxyInitializationHelpers: ["--local", "--run_hypercorn", "--max_requests_before_restart", "5"], ) - assert ( - result.exit_code == 0 - ), f"exit_code={result.exit_code}, output={result.output}" + assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_init_hypercorn.assert_called_once() - assert not any( - "does not recycle a worker" in str(c) for c in mock_print.call_args_list - ) + assert not any("does not recycle a worker" in str(c) for c in mock_print.call_args_list) @pytest.mark.skipif(os.name == "nt", reason="gunicorn server path skips Windows") def test_gunicorn_options_include_max_requests_jitter(self): @@ -2370,7 +2157,9 @@ class TestProxyInitializationHelpers: with patch.dict(os.environ, test_env_special): result = construct_database_url_from_env_vars() - expected_url = "postgresql://user%40with%2Bspecial:test-password-special-chars@localhost:5432/db_name%2Ftest" + expected_url = ( + "postgresql://user%40with%2Bspecial:test-password-special-chars@localhost:5432/db_name%2Ftest" + ) assert result == expected_url # Test without password (should still work) @@ -2498,11 +2287,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") - } + clean_env = {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( @@ -2517,12 +2302,8 @@ class TestQueryEngineReaperWiring: }, ), patch("uvicorn.run") as mock_uvicorn_run, - patch( - "litellm.proxy.proxy_cli.start_query_engine_reaper" - ) as mock_start_reaper, - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.start_query_engine_reaper") as mock_start_reaper, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -2533,17 +2314,13 @@ class TestQueryEngineReaperWiring: return result, mock_uvicorn_run, mock_start_reaper def test_multi_worker_uvicorn_starts_reaper(self): - result, mock_uvicorn_run, mock_start_reaper = self._invoke_run_server( - ["--local", "--num_workers", "2"] - ) + result, mock_uvicorn_run, mock_start_reaper = self._invoke_run_server(["--local", "--num_workers", "2"]) assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_uvicorn_run.assert_called_once() mock_start_reaper.assert_called_once() def test_single_worker_uvicorn_does_not_start_reaper(self): - result, mock_uvicorn_run, mock_start_reaper = self._invoke_run_server( - ["--local", "--num_workers", "1"] - ) + result, mock_uvicorn_run, mock_start_reaper = self._invoke_run_server(["--local", "--num_workers", "1"]) assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" mock_uvicorn_run.assert_called_once() mock_start_reaper.assert_not_called() @@ -2554,9 +2331,7 @@ class TestQueryEngineReaperWiring: with ( patch("gunicorn.app.base.BaseApplication.run"), - patch( - "litellm.proxy.proxy_cli.start_query_engine_reaper" - ) as mock_start_reaper, + patch("litellm.proxy.proxy_cli.start_query_engine_reaper") as mock_start_reaper, ): ProxyInitializationHelpers._run_gunicorn_server( host="127.0.0.1", @@ -2604,11 +2379,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 = {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 ( @@ -2620,9 +2391,7 @@ class TestRunServerDbSetup: "litellm.proxy.proxy_server": mock_proxy_module, }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -2653,10 +2422,18 @@ class TestRunServerDbSetup: mock_setup_database.assert_called_with(use_migrate=False, use_v2_resolver=True) @patch("atexit.register") - @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above - @patch("litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above - @patch("litellm.proxy.db.check_migration.check_prisma_schema_diff") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above - @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above + @patch( + "litellm.proxy.db.prisma_client.PrismaManager.setup_database" + ) # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above + @patch( + "litellm.proxy.db.prisma_client.PrismaManager.build_request_log_indexes" + ) # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above + @patch( + "litellm.proxy.db.check_migration.check_prisma_schema_diff" + ) # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above + @patch( + "litellm.proxy.db.prisma_client.should_update_prisma_schema" + ) # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above def test_migrations_run_when_the_prisma_cli_is_not_on_path( self, mock_should_update_schema, @@ -2680,11 +2457,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 = {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) @@ -2844,11 +2617,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 = {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 ( @@ -2860,9 +2629,7 @@ class TestRunServerDbSetup: "litellm.proxy.proxy_server": mock_proxy_module, }, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -2905,11 +2672,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 = {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 ( @@ -2923,9 +2686,7 @@ class TestRunServerDbSetup: ), ): with pytest.raises(SystemExit) as exc_info: - run_server.main( - ["--local", "--skip_server_startup"], standalone_mode=False - ) + run_server.main(["--local", "--skip_server_startup"], standalone_mode=False) assert exc_info.value.code == 1 mock_setup_database.assert_not_called() err: Final = capsys.readouterr().err @@ -2967,11 +2728,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 = {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" @@ -2985,9 +2742,7 @@ class TestRunServerDbSetup: }, ), ): - run_server.main( - ["--local", "--skip_server_startup"], standalone_mode=False - ) + run_server.main(["--local", "--skip_server_startup"], standalone_mode=False) mock_setup_database.assert_called_once_with(use_migrate=True, use_v2_resolver=True) assert "--use_v2_migration_resolver is deprecated" not in capsys.readouterr().out @@ -3030,9 +2785,7 @@ class TestRunServerDbSetup: {"proxy_server": mock_proxy_module, "litellm.proxy.proxy_server": mock_proxy_module}, ), ): - run_server.main( - ["--local", "--skip_server_startup", "--use_v2_migration_resolver"], standalone_mode=False - ) + run_server.main(["--local", "--skip_server_startup", "--use_v2_migration_resolver"], standalone_mode=False) mock_setup_database.assert_called_once_with(use_migrate=True, use_v2_resolver=True) assert "--use_v2_migration_resolver is deprecated and has no effect" in capsys.readouterr().out @@ -3169,12 +2922,7 @@ class TestRunServerDbSetup: def test_resolve_v2_migration_resolver(self, use_legacy_flag, env_value, expected): from litellm.proxy.proxy_cli import resolve_v2_migration_resolver - assert ( - resolve_v2_migration_resolver( - use_legacy_flag=use_legacy_flag, env_value=env_value - ) - is expected - ) + assert resolve_v2_migration_resolver(use_legacy_flag=use_legacy_flag, env_value=env_value) is expected def test_deprecated_v2_flag_not_reported_outside_a_cli_invocation(self): from litellm.proxy.proxy_cli import deprecated_v2_flag_passed_on_cli @@ -3215,9 +2963,7 @@ class TestRunServerDbSetup: ) clean_env = { - k: v - for k, v in os.environ.items() - if k not in ("DATABASE_URL", "DIRECT_URL", "USE_V2_MIGRATION_RESOLVER") + k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL", "USE_V2_MIGRATION_RESOLVER") } clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test" @@ -3298,9 +3044,7 @@ class TestRunServerDbSetup: "sys.modules", {"proxy_server": mock_proxy_module, "litellm.proxy.proxy_server": mock_proxy_module}, ), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, outcome as exc_info, ): mock_get_args.return_value = {"app": "litellm.proxy.proxy_server:app", "host": "localhost", "port": 8000} @@ -3352,11 +3096,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 = {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): @@ -3379,11 +3119,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 = {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): @@ -3403,11 +3139,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 = {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): @@ -3436,18 +3168,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 = "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 = {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): @@ -3593,13 +3318,9 @@ class TestPostgresStatementTimeoutOptions: import yaml config_path = tmp_path / "config.yaml" - config_path.write_text( - yaml.dump({"model_list": [], "general_settings": {"database_statement_timeout": 60}}) - ) + 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 = _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) @@ -3771,9 +3492,7 @@ class TestReadReplicaConnectionParams: import yaml config_path: Final = tmp_path / "config.yaml" - config_path.write_text( - yaml.dump({"model_list": [], "general_settings": {"database_connection_pool_limit": 3}}) - ) + config_path.write_text(yaml.dump({"model_list": [], "general_settings": {"database_connection_pool_limit": 3}})) _run_server_and_capture_urls( str(config_path), @@ -3793,9 +3512,7 @@ class TestReadReplicaConnectionParams: import yaml config_path: Final = tmp_path / "config.yaml" - config_path.write_text( - yaml.dump({"model_list": [], "general_settings": {"database_connection_pool_limit": 3}}) - ) + config_path.write_text(yaml.dump({"model_list": [], "general_settings": {"database_connection_pool_limit": 3}})) _run_server_and_capture_urls( str(config_path), @@ -3990,9 +3707,7 @@ class TestTokenAuthCliFlags: patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False), patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database"), patch("uvicorn.run"), - patch( - "litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args" - ) as mock_get_args, + patch("litellm.proxy.proxy_cli.ProxyInitializationHelpers.get_default_unvicorn_init_args") as mock_get_args, ): mock_get_args.return_value = { "app": "litellm.proxy.proxy_server:app", @@ -4006,9 +3721,7 @@ class TestTokenAuthCliFlags: return result, database_url, toggle def test_azure_flag_assembles_a_token_bearing_database_url(self): - result, database_url, toggle = self._invoke_with_azure_host( - ["--local", "--azure_postgresql_auth"] - ) + result, database_url, toggle = self._invoke_with_azure_host(["--local", "--azure_postgresql_auth"]) assert result.exit_code == 0, f"exit_code={result.exit_code}, output={result.output}" assert database_url is not None @@ -4236,7 +3949,11 @@ def probe_logger() -> Iterator[logging.Logger]: def _log_config_source(kind: str, tmp_path: Path) -> str | dict[str, object]: - config: Final[dict[str, object]] = {"version": 1, "incremental": True, "loggers": {_PROBE_LOGGER: {"level": "ERROR"}}} + config: Final[dict[str, object]] = { + "version": 1, + "incremental": True, + "loggers": {_PROBE_LOGGER: {"level": "ERROR"}}, + } if kind == "dict": return config written: Final = tmp_path / f"log_config.{kind}" @@ -4565,3 +4282,98 @@ 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 TestRustForwardServer: + """The public listener ``run_server`` reports to the Rust supervisor matches what it would serve itself.""" + + @staticmethod + def _server(argv: list[str]): + params: Final = run_server.make_context("litellm", ["--port", "45999", *argv]).params + return params, ProxyInitializationHelpers.python_server( + str(params["host"]), + int(str(params["port"])), + params["keepalive_timeout"], + run_gunicorn=params["run_gunicorn"] is True, + run_hypercorn=params["run_hypercorn"] is True, + run_granian=params["run_granian"] is True, + tls=any((params["ssl_keyfile_path"], params["ssl_certfile_path"], params["ciphers"])), + ) + + @pytest.mark.parametrize( + ("argv", "reason"), + [ + (["--ssl_keyfile_path", "unused"], "does not support TLS"), + (["--ssl_certfile_path", "unused"], "does not support TLS"), + (["--ciphers", "unused"], "does not support TLS"), + (["--run_hypercorn"], "does not support Hypercorn"), + (["--run_granian"], "does not support Granian"), + ], + ) + def test_unsupported_servers_fail_before_starting_children( + self, argv: list[str], reason: str, monkeypatch: pytest.MonkeyPatch + ) -> None: + from click.testing import CliRunner + + monkeypatch.setenv("LITELLM_RUST", "true") + monkeypatch.delenv("_LITELLM_GATEWAY_FD", raising=False) + result: Final = CliRunner().invoke(run_server, argv) + assert result.exit_code == 1 + assert f"LITELLM_RUST {reason}" in result.output + + @pytest.mark.parametrize( + ("argv", "environ"), + [(["--keepalive_timeout", "2"], {}), ([], {"KEEPALIVE_TIMEOUT": "7"}), ([], {})], + ) + def test_rust_retires_connections_by_the_keepalive_uvicorn_runs_with( + self, argv: list[str], environ: dict[str, str], monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.delenv("KEEPALIVE_TIMEOUT", raising=False) + for name, value in environ.items(): + monkeypatch.setenv(name, value) + params, server = self._server(argv) + served: Final = uvicorn.Config( + **ProxyInitializationHelpers.get_default_unvicorn_init_args( + "127.0.0.1", 45555, keepalive_timeout=params["keepalive_timeout"] + ) + ) + assert server.keepalive_seconds == served.timeout_keep_alive + + def test_gunicorn_workers_keep_gunicorns_keepalive_and_rust_follows(self, monkeypatch: pytest.MonkeyPatch) -> None: + from gunicorn.config import Config as GunicornConfig + + monkeypatch.delenv("KEEPALIVE_TIMEOUT", raising=False) + _, server = self._server(["--run_gunicorn", "--keepalive_timeout", "9"]) + assert server.keepalive_seconds == GunicornConfig().keepalive + + @pytest.mark.parametrize("gunicorn", [True, False]) + def test_python_behind_rust_trusts_the_proxies_its_server_would_have( + self, gunicorn: bool, monkeypatch: pytest.MonkeyPatch + ) -> None: + from gunicorn.config import Config as GunicornConfig + + from litellm.proxy.rust_forward.mode import InheritedFd, Supervised, rust_forward_mode + from litellm.proxy.rust_forward.platform import PrivateListener + from litellm.proxy.rust_forward.platform.posix import PosixPlatform + from litellm.proxy.rust_forward.supervisor import python_spec + + monkeypatch.delenv("FORWARDED_ALLOW_IPS", raising=False) + _, server = self._server(["--run_gunicorn"] if gunicorn else []) + mode: Final = rust_forward_mode( + python_spec( + PosixPlatform(), ("litellm",), PrivateListener(InheritedFd(17), 45555, (17,)), server + ).env_overrides + ) + trusted: Final = ( + GunicornConfig().forwarded_allow_ips + if gunicorn + else [uvicorn.Config(app="", log_config=None).forwarded_allow_ips] + ) + assert isinstance(mode, Supervised) + assert mode.trusted_proxies.split(",") == trusted + + def test_rust_drains_for_the_graceful_timeout_uvicorn_runs_with(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "12.5") + _, server = self._server([]) + served: Final = ProxyInitializationHelpers.get_default_unvicorn_init_args("127.0.0.1", 45555) + assert server.graceful_timeout_seconds == served["timeout_graceful_shutdown"] From c18d3c3044ca8d165210fecbe67295c1f2997cba Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 6 Nov 2026 03:47:22 +0900 Subject: [PATCH 02/10] fix(proxy): serve from Python alone when Rust forwarding cannot start LITELLM_RUST already selects native paths elsewhere, so a missing gateway binary or a server Rust cannot front now warns and keeps the plain Python proxy instead of exiting Co-Authored-By: Claude Opus 5.5 --- gateway/launch.py | 14 ++++-- litellm/proxy/proxy_cli.py | 8 ++-- litellm/proxy/rust_forward/AGENTS.md | 2 +- litellm/proxy/rust_forward/supervisor.py | 13 +++-- tests/unit/gateway/test_launch.py | 35 ++++++++------ .../proxy/rust_forward/test_supervisor.py | 18 +++++-- tests/unit/proxy/test_proxy_cli.py | 48 +++++++++++++++---- 7 files changed, 91 insertions(+), 47 deletions(-) diff --git a/gateway/launch.py b/gateway/launch.py index f7f98bd1156..449276ca089 100644 --- a/gateway/launch.py +++ b/gateway/launch.py @@ -132,16 +132,20 @@ def configure_database() -> None: def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -> None: from litellm.proxy.rust_forward.mode import Disabled, Supervised, Supervisor, rust_forward_mode, uvicorn_flags - from litellm.proxy.rust_forward.supervisor import supervise + from litellm.proxy.rust_forward.supervisor import StartupError, supervise mode: Final = rust_forward_mode(os.environ) match mode: - case Disabled(): + case Supervisor() if not isinstance( + exit_code := supervise( + python_server(argv), (sys.executable, "-m", "gateway.launch", *argv), before_start=configure_database + ), + StartupError, + ): + sys.exit(exit_code) + case Disabled() | Supervisor(): configure_database() serve(uvicorn_argv(argv, os.environ)) - case Supervisor(): - command: Final = (sys.executable, "-m", "gateway.launch", *argv) - sys.exit(supervise(python_server(argv), command, before_start=configure_database)) case Supervised(): serve((*uvicorn_argv(argv, os.environ), *uvicorn_flags(mode))) case _: diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index d3b3e073fa5..e99b06d6c8e 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -1368,7 +1368,7 @@ def run_server( if isinstance(forward_mode, Supervisor) and not any( (version, validate_config, health, test, skip_server_startup, check_migrations_against_existing_data) ): - from litellm.proxy.rust_forward.supervisor import supervise + from litellm.proxy.rust_forward.supervisor import StartupError, supervise server: Final = ProxyInitializationHelpers.python_server( host, @@ -1379,9 +1379,9 @@ def run_server( run_granian=run_granian, tls=any((ssl_keyfile_path, ssl_certfile_path, ciphers)), ) - raise SystemExit( - supervise(server, (sys.executable, "-m", "litellm.proxy.proxy_cli", *sys.argv[1:])) - ) + supervised: Final = supervise(server, (sys.executable, "-m", "litellm.proxy.proxy_cli", *sys.argv[1:])) + if not isinstance(supervised, StartupError): + raise SystemExit(supervised) args: Final = locals() try: diff --git a/litellm/proxy/rust_forward/AGENTS.md b/litellm/proxy/rust_forward/AGENTS.md index 294d94dd21b..64d9fb5673a 100644 --- a/litellm/proxy/rust_forward/AGENTS.md +++ b/litellm/proxy/rust_forward/AGENTS.md @@ -4,7 +4,7 @@ With `LITELLM_RUST` on, the supervisor here runs the Rust gateway on the public Every forwarding decision lives here. Entry points (`litellm/proxy/proxy_cli.py`, `gateway/launch.py`, `litellm_cli/main.py`) only translate between their own flags and this package: -- In the supervisor they describe what they would serve with as a `server.PythonServer` and hand it to `supervisor.supervise`, which decides whether forwarding can start, what Rust and the Python child receive, and how to report a refusal +- In the supervisor they describe what they would serve with as a `server.PythonServer` and hand it to `supervisor.supervise`, which decides whether forwarding can start and what Rust and the Python child receive. A refusal (no gateway binary, or a server or listener Rust cannot front) is reported and returned, and the entry point then serves from Python alone exactly as with `LITELLM_RUST` off - In the supervised child they merge `mode.uvicorn_options`, `mode.uvicorn_flags` or `mode.gunicorn_bind` and `mode.server_proxy_headers` into their server settings, and wrap the app with `peer.wrap` in their worker factory. These are identities when forwarding is off, so entry points never branch on it Every OS-specific operation lives behind `platform.Platform`, which `supervise` picks once with `platform.current_platform()`: the gateway binary name, how the private listener reaches Python, the spawn options, and how children are signalled and stopped. `PosixPlatform` hands Python an inherited socket (`mode.InheritedFd`). `WindowsPlatform` hands it an assigned loopback port (`mode.AssignedPort`), starts each child in its own process group and stops it with `CTRL_BREAK_EVENT`. The supervisor and `mode.py` never check the OS themselves diff --git a/litellm/proxy/rust_forward/supervisor.py b/litellm/proxy/rust_forward/supervisor.py index 3331b3f0c0b..5a53d3c562e 100644 --- a/litellm/proxy/rust_forward/supervisor.py +++ b/litellm/proxy/rust_forward/supervisor.py @@ -40,9 +40,7 @@ def startup_check(platform: Platform, server: PythonServer, binary: str | None) if server.listener is not Listener.TCP: return StartupError(f"{ENV_NAME} does not support {server.listener.value} on the public port") if binary is None: - return StartupError( - f"{ENV_NAME} requires the litellm-gateway binary. Install litellm[proxy] in this environment" - ) + return StartupError(f"{ENV_NAME} requires the litellm-gateway binary from the litellm-proxy package") return binary @@ -314,16 +312,17 @@ def supervise( before_start: Callable[[], None] = _nothing_to_prepare, scripts_dir: str = sysconfig.get_path("scripts"), platform: Platform | None = None, -) -> int: +) -> int | StartupError: """Serve ``server``'s public port from Rust and run ``command`` behind it, or say why that cannot start. - ``before_start`` runs only once forwarding is known to be possible, so a refused start leaves nothing behind. + A refusal is returned for the caller to serve from Python alone, and ``before_start`` runs only once forwarding + is known to be possible, so a refused start leaves nothing behind. """ chosen: Final = current_platform() if platform is None else platform checked: Final = startup_check(chosen, server, gateway_binary(chosen, scripts_dir)) if isinstance(checked, StartupError): - sys.stderr.write(f"{checked.reason}\n") - return ExitCode.CHILD_FAILED + sys.stderr.write(f"{checked.reason}, so Python serves the public port without Rust forwarding\n") + return checked before_start() try: return asyncio.run(_supervise(chosen, checked, command, server)) diff --git a/tests/unit/gateway/test_launch.py b/tests/unit/gateway/test_launch.py index 50dab887688..8a916560c1b 100644 --- a/tests/unit/gateway/test_launch.py +++ b/tests/unit/gateway/test_launch.py @@ -134,18 +134,6 @@ class TestUvicornArgv: ) -> None: assert python_server(arguments).listener is listener - def test_an_unsupported_listener_stops_forwarding_before_the_database_is_prepared( - self, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] - ) -> None: - monkeypatch.setenv("LITELLM_RUST", "true") - monkeypatch.delenv("_LITELLM_GATEWAY_FD", raising=False) - monkeypatch.setenv("LITELLM_PGBOUNCER_ENABLED", "true") - monkeypatch.setenv("LITELLM_PGBOUNCER_BINARY", "/nonexistent/pgbouncer") - with pytest.raises(SystemExit) as stopped: - main(("--uds", "/tmp/gateway.sock"), serve=lambda argv: None) - assert stopped.value.code == 1 - assert capsys.readouterr().err == "LITELLM_RUST does not support a Unix socket on the public port\n" - @pytest.mark.parametrize( ("argv", "trust"), [ @@ -327,15 +315,32 @@ class TestMain: DatabaseURLSettings.from_env().apply_to_env() assert os.environ["DATABASE_URL"] == pooled - def test_forwarding_disabled_keeps_uvicorn_socket_options( - self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch + @pytest.mark.parametrize( + ("rust", "warning"), + [ + ("false", ""), + ( + "true", + "LITELLM_RUST does not support a Unix socket on the public port, " + "so Python serves the public port without Rust forwarding\n", + ), + ], + ) + def test_python_alone_keeps_uvicorn_socket_options( + self, + rust: str, + warning: str, + password_env: dict[str, str], + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], ) -> None: - monkeypatch.setenv("LITELLM_RUST", "false") + monkeypatch.setenv("LITELLM_RUST", rust) monkeypatch.delenv("_LITELLM_GATEWAY_FD", raising=False) monkeypatch.delenv("LITELLM_PGBOUNCER_ENABLED", raising=False) served: Final[list[tuple[str, ...]]] = [] main(("--uds", "/tmp/gateway.sock"), serve=lambda argv: served.append(tuple(argv))) assert _uvicorn_params(served[0])["uds"] == "/tmp/gateway.sock" + assert capsys.readouterr().err == warning def test_a_pooler_that_cannot_start_stops_the_gateway_before_uvicorn( self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch, tmp_path: Path diff --git a/tests/unit/proxy/rust_forward/test_supervisor.py b/tests/unit/proxy/rust_forward/test_supervisor.py index d0ba6872bbd..ea69ce86951 100644 --- a/tests/unit/proxy/rust_forward/test_supervisor.py +++ b/tests/unit/proxy/rust_forward/test_supervisor.py @@ -120,8 +120,11 @@ def test_unsupported_servers_are_refused_before_anything_starts( assert startup_check(POSIX, server, "gateway") == StartupError(reason) prepared: Final[list[bool]] = [] refused: Final = supervise(server, ("python",), before_start=lambda: prepared.append(True), platform=POSIX) - assert refused == ExitCode.CHILD_FAILED - assert (capsys.readouterr().err, prepared) == (f"{reason}\n", []) + assert refused == StartupError(reason) + assert (capsys.readouterr().err, prepared) == ( + f"{reason}, so Python serves the public port without Rust forwarding\n", + [], + ) @pytest.mark.parametrize("server", [Server.UVICORN, Server.GUNICORN]) @@ -136,9 +139,14 @@ def test_windows_forwards_uvicorn_and_refuses_gunicorn() -> None: ) -def test_missing_binary_gives_install_instruction() -> None: - assert startup_check(POSIX, _server(), None) == StartupError( - "LITELLM_RUST requires the litellm-gateway binary. Install litellm[proxy] in this environment" +def test_missing_binary_is_refused_before_anything_starts(tmp_path: Path) -> None: + prepared: Final[list[bool]] = [] + refused: Final = supervise( + _server(), ("python",), before_start=lambda: prepared.append(True), scripts_dir=str(tmp_path), platform=POSIX + ) + assert (refused, prepared) == ( + StartupError("LITELLM_RUST requires the litellm-gateway binary from the litellm-proxy package"), + [], ) diff --git a/tests/unit/proxy/test_proxy_cli.py b/tests/unit/proxy/test_proxy_cli.py index fe0e46caaf3..c7d83a00188 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -4303,23 +4303,51 @@ class TestRustForwardServer: @pytest.mark.parametrize( ("argv", "reason"), [ - (["--ssl_keyfile_path", "unused"], "does not support TLS"), - (["--ssl_certfile_path", "unused"], "does not support TLS"), - (["--ciphers", "unused"], "does not support TLS"), + (["--ssl_keyfile_path", "unused"], "does not support TLS on the public port"), + (["--ssl_certfile_path", "unused"], "does not support TLS on the public port"), + (["--ciphers", "unused"], "does not support TLS on the public port"), (["--run_hypercorn"], "does not support Hypercorn"), (["--run_granian"], "does not support Granian"), ], ) - def test_unsupported_servers_fail_before_starting_children( - self, argv: list[str], reason: str, monkeypatch: pytest.MonkeyPatch + def test_unsupported_servers_are_refused_before_starting_children(self, argv: list[str], reason: str) -> None: + from litellm.proxy.rust_forward.platform.posix import PosixPlatform + from litellm.proxy.rust_forward.supervisor import StartupError, startup_check + + _, server = self._server(argv) + assert startup_check(PosixPlatform(), server, "gateway") == StartupError(f"LITELLM_RUST {reason}") + + @patch("uvicorn.run") + @patch("atexit.register") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) + def test_a_refused_forwarding_start_serves_from_python_alone( + self, mock_should_update, mock_setup_db, mock_atexit_register, mock_uvicorn_run ) -> None: from click.testing import CliRunner - monkeypatch.setenv("LITELLM_RUST", "true") - monkeypatch.delenv("_LITELLM_GATEWAY_FD", raising=False) - result: Final = CliRunner().invoke(run_server, argv) - assert result.exit_code == 1 - assert f"LITELLM_RUST {reason}" in result.output + proxy_module: Final = MagicMock( + app=MagicMock(), + ProxyConfig=MagicMock(), + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), + ) + environ: Final = { + **{k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL", "_LITELLM_GATEWAY_FD")}, + "LITELLM_RUST": "true", + } + with ( + patch.dict(os.environ, environ, clear=True), + patch.dict("sys.modules", {"proxy_server": proxy_module, "litellm.proxy.proxy_server": proxy_module}), + ): + result: Final = CliRunner().invoke( + run_server, ["--port", "45999", "--ssl_keyfile_path", "key.pem", "--ssl_certfile_path", "cert.pem"] + ) + assert result.exit_code == 0, result.output + assert "LITELLM_RUST does not support TLS on the public port, so Python serves the public port" in result.output + mock_uvicorn_run.assert_called_once() + assert mock_uvicorn_run.call_args.kwargs["ssl_keyfile"] == "key.pem" + assert mock_uvicorn_run.call_args.kwargs["port"] == 45999 @pytest.mark.parametrize( ("argv", "environ"), From 189fe5609002edafb4f82a81cb08c618ab2c7952 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 6 Nov 2026 05:18:29 +0900 Subject: [PATCH 03/10] fix(proxy): make the Rust forwarding supervisor transparent to Python The supervisor now relays every catchable signal to Python unchanged, sends the stop signals to Rust first, and exits with Python's status, dying by the same signal when Python did. Each child's stdin is a pipe only the supervisor holds, so both children shut down gracefully when the supervisor dies, even by SIGKILL or a closed terminal. LITELLM_RUST has one strict parser with no cache. Co-Authored-By: Claude Opus 5.5 --- gateway/launch.py | 10 +- litellm-rust/crates/gateway-forward/AGENTS.md | 1 + .../crates/gateway-forward/Cargo.toml | 2 +- .../crates/gateway-forward/src/lib.rs | 2 + .../crates/gateway-forward/src/supervisor.rs | 29 ++ .../crates/gateway-forward/tests/lifecycle.rs | 53 +++- .../gateway-forward/tests/support/mod.rs | 29 +- litellm-rust/crates/gateway/src/main.rs | 6 +- litellm/proxy/proxy_cli.py | 10 +- litellm/proxy/rust_forward/AGENTS.md | 10 +- litellm/proxy/rust_forward/mode.py | 57 +++- .../proxy/rust_forward/platform/__init__.py | 8 +- litellm/proxy/rust_forward/platform/posix.py | 24 +- .../proxy/rust_forward/platform/windows.py | 19 +- litellm/proxy/rust_forward/supervisor.py | 278 +++++++++++------- litellm/rust_bridge/configuration.py | 33 ++- .../test_rust_forward_lifecycle.py | 269 +++++++++++++++++ tests/unit/proxy/rust_forward/test_mode.py | 19 ++ .../proxy/rust_forward/test_supervisor.py | 171 ++++++++--- tests/unit/proxy/test_proxy_cli.py | 9 + tests/unit/rust_bridge/test_configuration.py | 21 +- 21 files changed, 852 insertions(+), 208 deletions(-) create mode 100644 litellm-rust/crates/gateway-forward/src/supervisor.rs create mode 100644 tests/integration/configuration/test_rust_forward_lifecycle.py diff --git a/gateway/launch.py b/gateway/launch.py index 449276ca089..979c40984d9 100644 --- a/gateway/launch.py +++ b/gateway/launch.py @@ -131,7 +131,14 @@ def configure_database() -> None: def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -> None: - from litellm.proxy.rust_forward.mode import Disabled, Supervised, Supervisor, rust_forward_mode, uvicorn_flags + from litellm.proxy.rust_forward.mode import ( + Disabled, + Supervised, + Supervisor, + rust_forward_mode, + stop_with_supervisor, + uvicorn_flags, + ) from litellm.proxy.rust_forward.supervisor import StartupError, supervise mode: Final = rust_forward_mode(os.environ) @@ -147,6 +154,7 @@ def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) - configure_database() serve(uvicorn_argv(argv, os.environ)) case Supervised(): + stop_with_supervisor(mode) serve((*uvicorn_argv(argv, os.environ), *uvicorn_flags(mode))) case _: assert_never(mode) diff --git a/litellm-rust/crates/gateway-forward/AGENTS.md b/litellm-rust/crates/gateway-forward/AGENTS.md index 2135980266f..c39d4b11981 100644 --- a/litellm-rust/crates/gateway-forward/AGENTS.md +++ b/litellm-rust/crates/gateway-forward/AGENTS.md @@ -13,6 +13,7 @@ - Pooled upstream connections are retired well before Python's keep-alive timeout, so a request never lands on a connection Python is closing - Idle client connections close after Python's keep-alive timeout, as uvicorn closed them, so a client that vanishes cannot hold a socket forever. Only the wait for the next request is timed, never an exchange in progress - WebSocket is the one protocol switch Python serves. Upgrades are relayed as raw bytes after the 101 and drained on shutdown +- Drain on EOF of stdin exactly as on SIGTERM. The supervisor holds the only write end, so EOF means it is gone and nobody is left to stop the gateway. `input_closed` reads on a plain thread, since a read blocked on tokio's blocking pool would hold up the runtime's exit ## Layout diff --git a/litellm-rust/crates/gateway-forward/Cargo.toml b/litellm-rust/crates/gateway-forward/Cargo.toml index 2fef19dca34..821ea91d269 100644 --- a/litellm-rust/crates/gateway-forward/Cargo.toml +++ b/litellm-rust/crates/gateway-forward/Cargo.toml @@ -13,7 +13,7 @@ hyper = { version = "1.12.0", features = ["http1", "server"] } hyper-util = { workspace = true, features = ["client-legacy", "http1", "tokio"] } serde.workspace = true thiserror.workspace = true -tokio = { workspace = true, features = ["io-util"] } +tokio = { workspace = true, features = ["io-util", "sync"] } tokio-util = { version = "0.7", features = ["rt"] } tower-service = "0.3.3" tracing.workspace = true diff --git a/litellm-rust/crates/gateway-forward/src/lib.rs b/litellm-rust/crates/gateway-forward/src/lib.rs index 4ac17b50544..88e314b43b1 100644 --- a/litellm-rust/crates/gateway-forward/src/lib.rs +++ b/litellm-rust/crates/gateway-forward/src/lib.rs @@ -5,6 +5,7 @@ mod forward; mod rewrite; mod server; mod settings; +mod supervisor; mod websocket; pub use error::Error; @@ -12,6 +13,7 @@ pub(crate) use error::Rejection; pub use forward::Forward; pub use rewrite::PEER_HEADER; pub use settings::{Settings, Upstream}; +pub use supervisor::input_closed; /// Binds the public listener and forwards to Python until `shutdown` resolves, then drains for /// `settings.shutdown_drain`. diff --git a/litellm-rust/crates/gateway-forward/src/supervisor.rs b/litellm-rust/crates/gateway-forward/src/supervisor.rs new file mode 100644 index 00000000000..f5848ce49e8 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/src/supervisor.rs @@ -0,0 +1,29 @@ +use std::io::{self, Read}; + +use tokio::sync::oneshot; + +/// Resolves once `input` reaches end of file or fails to read. The Python supervisor holds the only +/// write end of the gateway's stdin, so EOF means it is gone and the gateway drains as on SIGTERM. +/// +/// The read runs on a plain thread because a blocked read on tokio's blocking pool would hold up +/// runtime shutdown after a signal had already stopped the gateway. +pub fn input_closed(mut input: impl Read + Send + 'static) -> impl Future + Send { + let (closed, on_closed) = oneshot::channel::<()>(); + let reader = std::thread::Builder::new() + .name("supervisor-input".into()) + .spawn(move || { + let _ = io::copy(&mut input, &mut io::sink()); + let _ = closed.send(()); + }); + async move { + match reader { + Ok(_) => { + let _ = on_closed.await; + } + Err(error) => { + tracing::warn!(%error, "could not watch stdin, so only signals stop the gateway"); + std::future::pending::<()>().await; + } + } + } +} diff --git a/litellm-rust/crates/gateway-forward/tests/lifecycle.rs b/litellm-rust/crates/gateway-forward/tests/lifecycle.rs index 85d10a37be4..97907738ad6 100644 --- a/litellm-rust/crates/gateway-forward/tests/lifecycle.rs +++ b/litellm-rust/crates/gateway-forward/tests/lifecycle.rs @@ -7,9 +7,11 @@ use std::{ use axum::{Router, body::Body, extract::Request, http::StatusCode, response::Response}; use http_body_util::BodyExt; +use litellm_gateway_forward::input_closed; use rstest::rstest; use support::{ - DRAIN, Gateway, Raw, client, raw_upstream, recording, serve, unreachable_upstream, within, + DRAIN, Gateway, Raw, UVICORN_KEEPALIVE, client, raw_upstream, recording, serve, serve_until, + unreachable_upstream, within, }; use tokio::{ net::{TcpListener, TcpStream}, @@ -287,6 +289,55 @@ async fn shutdown_lets_an_in_flight_stream_finish() { within(server).await.unwrap(); } +#[rstest] +#[tokio::test] +async fn the_supervisor_closing_stdin_drains_like_a_stop_signal() { + let release = Arc::new(Notify::new()); + let released = release.clone(); + let upstream = raw_upstream(move |mut python: Raw| { + let released = released.clone(); + async move { + python.request_head().await.unwrap(); + python + .send("HTTP/1.1 200 OK\r\ntransfer-encoding: chunked\r\n\r\n5\r\nfirst\r\n") + .await; + released.notified().await; + python.send("4\r\nlast\r\n0\r\n\r\n").await; + } + }) + .await; + let (stdin, supervisor) = std::io::pipe().unwrap(); + let (address, server) = serve_until(upstream, UVICORN_KEEPALIVE, input_closed(stdin)).await; + let response = within( + client().request( + Request::get(format!("http://{address}/")) + .body(Body::empty()) + .unwrap(), + ), + ) + .await + .unwrap(); + tokio::time::sleep(Duration::from_millis(100)).await; + assert!(TcpStream::connect(address).await.is_ok()); + + drop(supervisor); + + within(async { + while TcpStream::connect(address).await.is_ok() { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await; + assert!(!server.is_finished()); + release.notify_one(); + let body = within(response.into_body().collect()) + .await + .unwrap() + .to_bytes(); + assert_eq!(body, "firstlast"); + within(server).await.unwrap(); +} + #[rstest] #[tokio::test] async fn shutdown_stops_waiting_for_a_stream_still_open_after_the_drain() { diff --git a/litellm-rust/crates/gateway-forward/tests/support/mod.rs b/litellm-rust/crates/gateway-forward/tests/support/mod.rs index d39532a456e..7d92ca4a14a 100644 --- a/litellm-rust/crates/gateway-forward/tests/support/mod.rs +++ b/litellm-rust/crates/gateway-forward/tests/support/mod.rs @@ -32,6 +32,21 @@ pub async fn within(future: impl Future) -> T { tokio::time::timeout(WAIT, future).await.expect("timed out") } +/// A forwarder in front of `upstream` that drains once `shutdown` resolves, as `run` does. +pub async fn serve_until( + upstream: SocketAddr, + keepalive: Duration, + shutdown: impl Future + Send + 'static, +) -> (SocketAddr, JoinHandle<()>) { + let forward = Forward::new(&format!("http://{upstream}").parse().unwrap(), keepalive); + let public = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = public.local_addr().unwrap(); + ( + address, + tokio::spawn(forward.serve(public, shutdown, DRAIN)), + ) +} + /// A running forwarder in front of `upstream`, listening on a real socket like production. pub struct Gateway { pub address: SocketAddr, @@ -45,17 +60,11 @@ impl Gateway { } pub async fn with_keepalive(upstream: SocketAddr, keepalive: Duration) -> Self { - let forward = Forward::new(&format!("http://{upstream}").parse().unwrap(), keepalive); - let public = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = public.local_addr().unwrap(); let (stop, stopped) = oneshot::channel::<()>(); - let server = tokio::spawn(forward.serve( - public, - async { - stopped.await.ok(); - }, - DRAIN, - )); + let (address, server) = serve_until(upstream, keepalive, async { + stopped.await.ok(); + }) + .await; Self { address, stop, diff --git a/litellm-rust/crates/gateway/src/main.rs b/litellm-rust/crates/gateway/src/main.rs index ff528723095..ebd54fb6029 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -47,7 +47,11 @@ async fn main() -> Result<(), Box> { level: log.log_level(), }) .install_global()?; - litellm_gateway_forward::run(forward, shutdown_signal()).await?; + let stdin_closed = litellm_gateway_forward::input_closed(std::io::stdin()); + litellm_gateway_forward::run(forward, async { + tokio::select! { () = shutdown_signal() => (), () = stdin_closed => () } + }) + .await?; return Ok(()); } let settings: Settings = settings::from_iter(std::env::vars_os())?; diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index e99b06d6c8e..0338b7261b9 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -1356,15 +1356,23 @@ def run_server( run_setup_wizard() return + from pydantic import ValidationError + from litellm.proxy.rust_forward.mode import ( Supervisor, + environment_error, gunicorn_bind, rust_forward_mode, server_proxy_headers, + stop_with_supervisor, uvicorn_options, ) - forward_mode: Final = rust_forward_mode(os.environ) + try: + forward_mode: Final = rust_forward_mode(os.environ) + except ValidationError as error: + raise click.ClickException(environment_error(error)) from None + stop_with_supervisor(forward_mode) if isinstance(forward_mode, Supervisor) and not any( (version, validate_config, health, test, skip_server_startup, check_migrations_against_existing_data) ): diff --git a/litellm/proxy/rust_forward/AGENTS.md b/litellm/proxy/rust_forward/AGENTS.md index 64d9fb5673a..2c76bf73674 100644 --- a/litellm/proxy/rust_forward/AGENTS.md +++ b/litellm/proxy/rust_forward/AGENTS.md @@ -7,12 +7,18 @@ Every forwarding decision lives here. Entry points (`litellm/proxy/proxy_cli.py` - In the supervisor they describe what they would serve with as a `server.PythonServer` and hand it to `supervisor.supervise`, which decides whether forwarding can start and what Rust and the Python child receive. A refusal (no gateway binary, or a server or listener Rust cannot front) is reported and returned, and the entry point then serves from Python alone exactly as with `LITELLM_RUST` off - In the supervised child they merge `mode.uvicorn_options`, `mode.uvicorn_flags` or `mode.gunicorn_bind` and `mode.server_proxy_headers` into their server settings, and wrap the app with `peer.wrap` in their worker factory. These are identities when forwarding is off, so entry points never branch on it -Every OS-specific operation lives behind `platform.Platform`, which `supervise` picks once with `platform.current_platform()`: the gateway binary name, how the private listener reaches Python, the spawn options, and how children are signalled and stopped. `PosixPlatform` hands Python an inherited socket (`mode.InheritedFd`). `WindowsPlatform` hands it an assigned loopback port (`mode.AssignedPort`), starts each child in its own process group and stops it with `CTRL_BREAK_EVENT`. The supervisor and `mode.py` never check the OS themselves +Every OS-specific operation lives behind `platform.Platform`, which `supervise` picks once with `platform.current_platform()`: the gateway binary name, how the private listener reaches Python, the spawn options, which signals are relayed and which of them stop Rust, and how children are signalled and stopped. `PosixPlatform` hands Python an inherited socket (`mode.InheritedFd`). `WindowsPlatform` hands it an assigned loopback port (`mode.AssignedPort`), starts each child in its own process group and stops it with `CTRL_BREAK_EVENT`. The supervisor and `mode.py` never check the OS themselves Windows support exists only for running the proxy locally from the CLI. Nobody deploys the proxy on Windows, and the Windows design leans on that. A port picked by the supervisor can be taken by another process before Python binds it, `CTRL_BREAK_EVENT` needs the console a CLI run always has, and gunicorn is refused. Revisit those choices before treating Windows as a deployment target Each environment contract is written and read in one place. `mode.supervised_environment` and `mode.rust_forward_mode` are the two halves of the Python child's, and `supervisor.rust_spec` writes the `LITELLM_FORWARD_*` variables that `gateway-forward`'s `settings.rs` reads. Python owns every value Rust needs, so a new one gets a new variable in `rust_spec` -The supervisor's whole lifecycle is `supervisor.step`, a pure function from state and event to the next state and the actions to run. Children start in `START_ORDER`, each once the one before it reports `Ready`, and on a stop every started child is signalled at once in reverse start order. Keep decisions in `step` and keep `_run` executing actions, so a new child is an entry in `START_ORDER` and the launch table rather than a new state +The supervisor is transparent to Python, so `python` with `LITELLM_RUST` behaves on signals and exits exactly as Python alone. Its whole lifecycle is `supervisor.step`, a pure function from a `Lifecycle` of facts and an event to the next facts and the actions to run, and `exit_status` reads the result from those facts. `_run` only executes actions, and the `async with` around each child owns its cleanup. Three rules decide everything there, and a change that breaks one of them breaks the equivalence `tests/integration/configuration/test_rust_forward_lifecycle.py` checks against Python alone + +- Signals pass through. Every signal in `Platform.relayed_signals` goes to Python unchanged, and the ones in `Platform.stop_signals` go to Rust first. The supervisor never reads meaning into a signal, so a repeated SIGINT reaches Python again and forces its exit as it would alone. Children run in their own session (`start_new_session`) only so that a terminal Ctrl-C reaches Python once, through the supervisor +- Python's exit status is the result. When Python exits for any reason, Rust gets SIGTERM, and once Rust is gone the supervisor exits with Python's status, dying by the same signal when Python died by one. The one exception is Rust exiting while Python serves with no stop requested, which stops Python and exits 1 +- Children cannot outlive the supervisor. Each child's stdin is a pipe only the supervisor holds open. The Python process the supervisor started (`mode.stop_with_supervisor`, never its workers) and the Rust gateway (`litellm_gateway_forward::input_closed`) shut down gracefully on EOF, so a supervisor killed by SIGKILL or a closed terminal leaves nothing behind + +Rust starts once Python listens, never after a stop was requested or Python exited. Windows children run in their own process group, which can only be sent `CTRL_BREAK_EVENT`, so both relayed signals arrive as that and a second Ctrl-C does not force Python's exit there Python must see each request as if the client had connected directly. `x-litellm-rust-peer` is the only Rust-specific input, and `peer.py` honors it only from the loopback hop. The proxy app itself never knows Rust is in front of it diff --git a/litellm/proxy/rust_forward/mode.py b/litellm/proxy/rust_forward/mode.py index a95a66cf4bc..4b4cd77e05b 100644 --- a/litellm/proxy/rust_forward/mode.py +++ b/litellm/proxy/rust_forward/mode.py @@ -1,12 +1,17 @@ -from collections.abc import Mapping +import contextlib +import os +import signal +import threading +from collections.abc import Callable, Mapping from dataclasses import dataclass +from functools import partial from types import MappingProxyType -from typing import Annotated, Final, TypeAlias +from typing import Final, TypeAlias -from pydantic import BaseModel, BeforeValidator, ConfigDict, Field +from pydantic import BaseModel, ConfigDict, Field, ValidationError from typing_extensions import assert_never -from litellm.rust_bridge.configuration import ENV_NAME +from litellm.rust_bridge.configuration import ENV_NAME, RustSwitch GATEWAY_FD: Final = "_LITELLM_GATEWAY_FD" GATEWAY_PORT: Final = "_LITELLM_GATEWAY_PORT" @@ -16,18 +21,10 @@ DEFAULT_TRUSTED_PROXIES: Final = "127.0.0.1" RUST_HOP: Final = "127.0.0.1" -def _blank_as_unset(value: object) -> object: - stripped: Final = value.strip() if isinstance(value, str) else value - return None if stripped == "" else stripped - - -_Switch: TypeAlias = Annotated[bool | None, BeforeValidator(_blank_as_unset)] - - class _ForwardEnvironment(BaseModel): model_config = ConfigDict(frozen=True) - rust: _Switch = Field(None, validation_alias=ENV_NAME) + rust: RustSwitch = Field(None, validation_alias=ENV_NAME) listener_fd: int | None = Field(None, validation_alias=GATEWAY_FD) listener_port: int | None = Field(None, validation_alias=GATEWAY_PORT) proxy_headers: bool = Field(True, validation_alias=PROXY_HEADERS) @@ -101,6 +98,40 @@ def rust_forward_mode(environ: Mapping[str, str]) -> RustForwardMode: return Disabled() +def _wait_for_eof(fd: int) -> None: + with contextlib.suppress(OSError): + while os.read(fd, 4096): + pass + + +def _stop_when_closed(fd: int, stop: Callable[[], None]) -> None: + _wait_for_eof(fd) + stop() + + +def stop_with_supervisor( + mode: RustForwardMode, + fd: int = 0, + stop: Callable[[], None] = partial(signal.raise_signal, signal.SIGTERM), +) -> None: + """Shut the supervised server down gracefully once the supervisor's end of its stdin closes, so it never outlives + a supervisor that was killed. Call it only in the process the supervisor started, not in its workers.""" + match mode: + case Supervised(): + threading.Thread(target=_stop_when_closed, args=(fd, stop), name="supervisor-stdin", daemon=True).start() + case Disabled() | Supervisor(): + pass + case _: + assert_never(mode) + + +def environment_error(error: ValidationError) -> str: + return "; ".join( + f"{'.'.join(str(part) for part in detail['loc'])} is invalid: {detail['msg']}" + for detail in error.errors(include_url=False) + ) + + # A supervised server's TCP peer is always Rust on loopback, so server-level proxy headers would trust any # client's forwarding headers. They stay off there and ``peer.wrap`` applies the operator's trust instead. diff --git a/litellm/proxy/rust_forward/platform/__init__.py b/litellm/proxy/rust_forward/platform/__init__.py index b40ff5cb015..6b273864356 100644 --- a/litellm/proxy/rust_forward/platform/__init__.py +++ b/litellm/proxy/rust_forward/platform/__init__.py @@ -38,7 +38,13 @@ class Platform(Protocol): def spawn_options(self, pass_fds: tuple[int, ...]) -> SpawnOptions: ... - def catch_stop(self, loop: asyncio.AbstractEventLoop, on_stop: Callable[[signal.Signals], None]) -> None: ... + @property + def relayed_signals(self) -> tuple[signal.Signals, ...]: ... + + @property + def stop_signals(self) -> frozenset[signal.Signals]: ... + + def catch(self, loop: asyncio.AbstractEventLoop, on_signal: Callable[[signal.Signals], None]) -> None: ... def signal_child(self, process: asyncio.subprocess.Process, signum: signal.Signals) -> None: ... diff --git a/litellm/proxy/rust_forward/platform/posix.py b/litellm/proxy/rust_forward/platform/posix.py index b1b2e500657..f60791f2c06 100644 --- a/litellm/proxy/rust_forward/platform/posix.py +++ b/litellm/proxy/rust_forward/platform/posix.py @@ -32,9 +32,27 @@ class PosixPlatform: def spawn_options(self, pass_fds: tuple[int, ...]) -> SpawnOptions: return SpawnOptions(pass_fds=pass_fds, start_new_session=True, creationflags=0) - def catch_stop(self, loop: asyncio.AbstractEventLoop, on_stop: Callable[[signal.Signals], None]) -> None: - for signum in (signal.SIGTERM, signal.SIGINT): - loop.add_signal_handler(signum, on_stop, signum) + @property + def relayed_signals(self) -> tuple[signal.Signals, ...]: + return ( + signal.SIGTERM, + signal.SIGINT, + signal.SIGHUP, + signal.SIGQUIT, + signal.SIGUSR1, + signal.SIGUSR2, + signal.SIGTTIN, + signal.SIGTTOU, + signal.SIGWINCH, + ) + + @property + def stop_signals(self) -> frozenset[signal.Signals]: + return frozenset({signal.SIGTERM, signal.SIGINT, signal.SIGQUIT}) + + def catch(self, loop: asyncio.AbstractEventLoop, on_signal: Callable[[signal.Signals], None]) -> None: + for signum in self.relayed_signals: + loop.add_signal_handler(signum, on_signal, signum) def signal_child(self, process: asyncio.subprocess.Process, signum: signal.Signals) -> None: with contextlib.suppress(ProcessLookupError): diff --git a/litellm/proxy/rust_forward/platform/windows.py b/litellm/proxy/rust_forward/platform/windows.py index f32063428e3..df202d11931 100644 --- a/litellm/proxy/rust_forward/platform/windows.py +++ b/litellm/proxy/rust_forward/platform/windows.py @@ -41,13 +41,24 @@ class WindowsPlatform: def spawn_options(self, pass_fds: tuple[int, ...]) -> SpawnOptions: return SpawnOptions(pass_fds=pass_fds, start_new_session=False, creationflags=CREATE_NEW_PROCESS_GROUP) - def catch_stop(self, loop: asyncio.AbstractEventLoop, on_stop: Callable[[signal.Signals], None]) -> None: - def handle(signum: int, _frame: FrameType | None) -> None: - loop.call_soon_threadsafe(on_stop, signal.Signals(signum)) + @property + def relayed_signals(self) -> tuple[signal.Signals, ...]: + return (signal.SIGINT, signal.Signals(SIGBREAK)) - for signum in (signal.SIGINT, SIGBREAK): + @property + def stop_signals(self) -> frozenset[signal.Signals]: + return frozenset(self.relayed_signals) + + def catch(self, loop: asyncio.AbstractEventLoop, on_signal: Callable[[signal.Signals], None]) -> None: + def handle(signum: int, _frame: FrameType | None) -> None: + loop.call_soon_threadsafe(on_signal, signal.Signals(signum)) + + for signum in self.relayed_signals: signal.signal(signum, handle) + # A child in its own process group can only be sent CTRL_BREAK_EVENT, since Windows never delivers CTRL_C_EVENT + # to a new group. Both children shut down on it as on Ctrl-C, but uvicorn force-exits only on a repeated SIGINT, + # so a second Ctrl-C does not cut a Windows shutdown short the way it does on POSIX. def signal_child(self, process: asyncio.subprocess.Process, signum: signal.Signals) -> None: self.stop_group(process) diff --git a/litellm/proxy/rust_forward/supervisor.py b/litellm/proxy/rust_forward/supervisor.py index 5a53d3c562e..f02d6d4e1cd 100644 --- a/litellm/proxy/rust_forward/supervisor.py +++ b/litellm/proxy/rust_forward/supervisor.py @@ -6,8 +6,8 @@ import signal import sys import sysconfig from collections.abc import AsyncGenerator, Callable, Coroutine, Mapping, Sequence -from dataclasses import dataclass -from enum import Enum, IntEnum +from dataclasses import dataclass, replace +from enum import Enum from functools import partial from types import MappingProxyType from typing import Final, TypeAlias @@ -49,109 +49,141 @@ class Child(Enum): PYTHON = "python" -# Each child starts once the one before it is ready. Python alone takes no connections until it has booted, so Rust -# only opens the public port once Python listens. -START_ORDER: Final = (Child.PYTHON, Child.RUST) - - -class ExitCode(IntEnum): - CLEAN = 0 - CHILD_FAILED = 1 +FAILURE: Final = 1 @dataclass(frozen=True, slots=True) -class Ready: - child: Child +class NotStarted: + pass @dataclass(frozen=True, slots=True) -class Stop: - signum: signal.Signals +class Spawned: + pass + + +@dataclass(frozen=True, slots=True) +class Listening: + pass @dataclass(frozen=True, slots=True) class Exited: + status: int + + +@dataclass(frozen=True, slots=True) +class Lost: + """Rust exited while Python was serving and nobody had asked either to stop.""" + + status: int + + +PythonState: TypeAlias = Spawned | Listening | Exited +RustState: TypeAlias = NotStarted | Spawned | Exited | Lost + + +@dataclass(frozen=True, slots=True) +class Lifecycle: + python: PythonState + rust: RustState + stop_requested: bool + + +@dataclass(frozen=True, slots=True) +class Caught: + signum: signal.Signals + + +@dataclass(frozen=True, slots=True) +class PythonListening: + pass + + +@dataclass(frozen=True, slots=True) +class ChildExited: child: Child - code: int + status: int @dataclass(frozen=True, slots=True) -class Running: - up: tuple[Child, ...] - - -@dataclass(frozen=True, slots=True) -class Stopping: - up: frozenset[Child] - - -@dataclass(frozen=True, slots=True) -class Finished: - code: ExitCode - - -@dataclass(frozen=True, slots=True) -class Start: +class Spawn: child: Child @dataclass(frozen=True, slots=True) -class Signal: +class Send: child: Child signum: signal.Signals -Event: TypeAlias = Ready | Stop | Exited -State: TypeAlias = Running | Stopping | Finished -Action: TypeAlias = Start | Signal -Transition: TypeAlias = tuple[State, tuple[Action, ...]] - - -def _start_after(up: tuple[Child, ...]) -> Transition: - following: Final = START_ORDER[len(up) : len(up) + 1] - return Running((*up, *following)), tuple(Start(child) for child in following) +Event: TypeAlias = Caught | PythonListening | ChildExited +Action: TypeAlias = Spawn | Send +Transition: TypeAlias = tuple[Lifecycle, tuple[Action, ...]] def boot() -> Transition: - return _start_after(()) + return Lifecycle(Spawned(), NotStarted(), stop_requested=False), (Spawn(Child.PYTHON),) -# Every child gets the signal at once, as Python alone stopped. Python ends its long streams within its graceful -# timeout, which Rust's drain already covers, so stopping Python only after Rust exits would hold Python's shutdown, -# spend flush included, behind streams only Python can end. Reverse start order lets Rust stop accepting first. -def _stop(up: tuple[Child, ...], signum: signal.Signals) -> Transition: - return Stopping(frozenset(up)), tuple(Signal(child, signum) for child in reversed(up)) +def _relay(state: Lifecycle, signum: signal.Signals, stops: bool) -> Transition: + to_rust: Final = (Send(Child.RUST, signum),) if stops and isinstance(state.rust, Spawned) else () + to_python: Final = (Send(Child.PYTHON, signum),) if not isinstance(state.python, Exited) else () + return replace(state, stop_requested=state.stop_requested or stops), (*to_rust, *to_python) -def step(state: State, event: Event) -> Transition: - match state: - case Running(up): - match event: - case Ready(child) if up[-1:] == (child,): - return _start_after(up) - case Ready(): - return state, () - case Stop(signum): - return _stop(up, signum) - case Exited(): - return Finished(ExitCode.CHILD_FAILED), () +# Python alone takes no connections until it has booted, so Rust only opens the public port once Python listens. +def _python_listening(state: Lifecycle) -> Transition: + if not isinstance(state.python, Spawned): + return state, () + if state.stop_requested: + return replace(state, python=Listening()), () + return replace(state, python=Listening(), rust=Spawned()), (Spawn(Child.RUST),) + + +def _python_exited(state: Lifecycle, status: int) -> Transition: + if isinstance(state.python, Exited): + return state, () + exited: Final = replace(state, python=Exited(status)) + return exited, ((Send(Child.RUST, signal.SIGTERM),) if isinstance(state.rust, Spawned) else ()) + + +def _rust_exited(state: Lifecycle, status: int) -> Transition: + if not isinstance(state.rust, Spawned): + return state, () + if state.stop_requested or isinstance(state.python, Exited): + return replace(state, rust=Exited(status)), () + return replace(state, rust=Lost(status)), (Send(Child.PYTHON, signal.SIGTERM),) + + +def step(state: Lifecycle, event: Event, stop_signals: frozenset[signal.Signals]) -> Transition: + """The supervisor never gives a signal or an exit a meaning of its own. Every signal reaches Python as sent, + Python's exit status is the supervisor's, and Rust only follows Python, except that losing Rust stops Python.""" + match event: + case Caught(signum): + return _relay(state, signum, signum in stop_signals) + case PythonListening(): + return _python_listening(state) + case ChildExited(child, status): + match child: + case Child.PYTHON: + return _python_exited(state, status) + case Child.RUST: + return _rust_exited(state, status) case _: - assert_never(event) - case Stopping(up): - match event: - case Ready() | Stop(): - return state, () - case Exited(child, code) if child in up and code == 0: - remaining: Final = up - {child} - return (Stopping(remaining) if remaining else Finished(ExitCode.CLEAN)), () - case Exited(): - return Finished(ExitCode.CHILD_FAILED), () - case _: - assert_never(event) - case Finished(): - return state, () + assert_never(child) case _: - assert_never(state) + assert_never(event) + + +def exit_status(state: Lifecycle) -> int | None: + match state.python, state.rust: + case Exited(status), NotStarted() | Exited(): + return status + case Exited(), Lost(): + return FAILURE + case _: + return None @dataclass(frozen=True, slots=True) @@ -192,6 +224,7 @@ def rust_spec(platform: Platform, binary: str, upstream_port: int, server: Pytho async def _spawn(spec: ChildSpec) -> asyncio.subprocess.Process: return await asyncio.create_subprocess_exec( *spec.argv, + stdin=asyncio.subprocess.PIPE, env={**os.environ, **spec.env_overrides}, pass_fds=spec.spawn.pass_fds, start_new_session=spec.spawn.start_new_session, @@ -203,7 +236,7 @@ Launch: TypeAlias = Callable[[asyncio.Queue[Event]], contextlib.AbstractAsyncCon async def _report_exit(child: Child, process: asyncio.subprocess.Process, events: asyncio.Queue[Event]) -> None: - events.put_nowait(Exited(child, await process.wait())) + events.put_nowait(ChildExited(child, await process.wait())) async def _accepts(port: int) -> bool: @@ -215,10 +248,10 @@ async def _accepts(port: int) -> bool: return True -async def _report_ready_once_listening(child: Child, port: int, events: asyncio.Queue[Event]) -> None: +async def _report_once_listening(port: int, events: asyncio.Queue[Event]) -> None: while not await _accepts(port): await asyncio.sleep(PROBE_INTERVAL_SECONDS) - events.put_nowait(Ready(child)) + events.put_nowait(PythonListening()) @contextlib.asynccontextmanager @@ -231,6 +264,8 @@ async def _in_background(work: Coroutine[None, None, None]) -> AsyncGenerator[No await asyncio.gather(task, return_exceptions=True) +# A child's stdin is a pipe only the supervisor holds open, so the child sees EOF as soon as the supervisor is gone, +# however it went, and shuts down instead of outliving it. @contextlib.asynccontextmanager async def _launch( platform: Platform, @@ -245,56 +280,78 @@ async def _launch( finally: platform.stop_group(process) await waiter + if process.stdin is not None: + process.stdin.close() @contextlib.asynccontextmanager -async def _launch_until_listening( - platform: Platform, child: Child, spec: ChildSpec, port: int, events: asyncio.Queue[Event] +async def _launch_python( + platform: Platform, spec: ChildSpec, port: int, events: asyncio.Queue[Event] ) -> AsyncGenerator[asyncio.subprocess.Process]: async with ( - _launch(platform, child, spec, events) as process, - _in_background(_report_ready_once_listening(child, port, events)), + _launch(platform, Child.PYTHON, spec, events) as process, + _in_background(_report_once_listening(port, events)), ): yield process +@dataclass(frozen=True, slots=True) +class _Launching: + state: Lifecycle + child: Child + rest: tuple[Action, ...] + + +async def _advance( + platform: Platform, + transition: Transition, + events: asyncio.Queue[Event], + children: Mapping[Child, asyncio.subprocess.Process], +) -> int | _Launching: + state, actions = transition + while True: + for index, action in enumerate(actions): + match action: + case Send(child, signum): + platform.signal_child(children[child], signum) + case Spawn(child): + return _Launching(state, child, actions[index + 1 :]) + case _: + assert_never(action) + status = exit_status(state) + if status is not None: + return status + state, actions = step(state, await events.get(), platform.stop_signals) + + async def _run( platform: Platform, transition: Transition, events: asyncio.Queue[Event], children: Mapping[Child, asyncio.subprocess.Process], launches: Mapping[Child, Launch], -) -> ExitCode: - state, actions = transition - if not actions: - if isinstance(state, Finished): - return state.code - return await _run(platform, step(state, await events.get()), events, children, launches) - action: Final = actions[0] - rest: Final = (state, actions[1:]) - match action: - case Signal(child, signum): - platform.signal_child(children[child], signum) - return await _run(platform, rest, events, children, launches) - case Start(child): +) -> int: + advanced: Final = await _advance(platform, transition, events, children) + match advanced: + case int(): + return advanced + case _Launching(state, child, rest): async with launches[child](events) as process: - return await _run(platform, rest, events, MappingProxyType({**children, child: process}), launches) + return await _run( + platform, (state, rest), events, MappingProxyType({**children, child: process}), launches + ) case _: - assert_never(action) + assert_never(advanced) -async def _supervise(platform: Platform, binary: str, command: Sequence[str], server: PythonServer) -> ExitCode: +async def _supervise(platform: Platform, binary: str, command: Sequence[str], server: PythonServer) -> int: events: Final[asyncio.Queue[Event]] = asyncio.Queue() - platform.catch_stop(asyncio.get_running_loop(), lambda signum: events.put_nowait(Stop(signum))) + platform.catch(asyncio.get_running_loop(), lambda signum: events.put_nowait(Caught(signum))) with platform.private_listener() as private: launches: Final[Mapping[Child, Launch]] = MappingProxyType( { Child.PYTHON: partial( - _launch_until_listening, - platform, - Child.PYTHON, - python_spec(platform, command, private, server), - private.port, + _launch_python, platform, python_spec(platform, command, private, server), private.port ), Child.RUST: partial(_launch, platform, Child.RUST, rust_spec(platform, binary, private.port, server)), } @@ -302,6 +359,16 @@ async def _supervise(platform: Platform, binary: str, command: Sequence[str], se return await _run(platform, boot(), events, MappingProxyType({}), launches) +def _exit_like(status: int) -> int: + """A child killed by a signal kills the supervisor with the same signal, as uvicorn re-raises the one it caught.""" + if status >= 0: + return status + signum: Final = signal.Signals(-status) + signal.signal(signum, signal.SIG_DFL) + signal.raise_signal(signum) + return 128 + signum + + def _nothing_to_prepare() -> None: pass @@ -325,7 +392,8 @@ def supervise( return checked before_start() try: - return asyncio.run(_supervise(chosen, checked, command, server)) + status: Final = asyncio.run(_supervise(chosen, checked, command, server)) except OSError as error: sys.stderr.write(f"Rust gateway startup failed: {error}\n") - return ExitCode.CHILD_FAILED + return FAILURE + return _exit_like(status) diff --git a/litellm/rust_bridge/configuration.py b/litellm/rust_bridge/configuration.py index 5b51f155df9..858e2a6a766 100644 --- a/litellm/rust_bridge/configuration.py +++ b/litellm/rust_bridge/configuration.py @@ -2,14 +2,26 @@ 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 from typing_extensions import assert_never ENV_NAME: Final = "LITELLM_RUST" -_ENV_BOOL: Final = TypeAdapter(bool) + + +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) + + +def rust_switch(raw: str | None) -> bool | None: + """``LITELLM_RUST`` as a switch, unset when absent or blank. Raises ``ValidationError`` on any other value.""" + return _RUST_SWITCH.validate_python(raw) class Rollout(Enum): @@ -33,17 +45,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 +73,7 @@ def decision(rollout: Rollout) -> Decision: return decide( rollout, process_override=_CONFIGURATION.override, - environment_override=_parse_env_bool(os.getenv(ENV_NAME)), + environment_override=rust_switch(os.getenv(ENV_NAME)), ) diff --git a/tests/integration/configuration/test_rust_forward_lifecycle.py b/tests/integration/configuration/test_rust_forward_lifecycle.py new file mode 100644 index 00000000000..cc186af98d7 --- /dev/null +++ b/tests/integration/configuration/test_rust_forward_lifecycle.py @@ -0,0 +1,269 @@ +import os +import signal +import socket +import subprocess +import sys +import time +from collections.abc import Callable +from dataclasses import dataclass +from enum import Enum +from pathlib import Path +from typing import Final + +import httpx +import psutil +import pytest + +from tests.integration._support.client import eventually +from tests.integration._support.ports import allocate_port + +pytestmark = pytest.mark.skipif(sys.platform == "win32", reason="drives POSIX signals") + +APP: Final = """ +import asyncio + + +async def app(scope, receive, send): + if scope["type"] != "http": + return + await send({"type": "http.response.start", "status": 200, "headers": [(b"content-type", b"text/plain")]}) + if scope["path"] != "/stream": + await send({"type": "http.response.body", "body": b"ok"}) + return + while True: + await send({"type": "http.response.body", "body": b"tick\\n", "more_body": True}) + await asyncio.sleep(0.1) +""" + +SERVE: Final = """ +import os +import pathlib +import sys + +import uvicorn + + + +def listener(): + fd = os.environ.get("_LITELLM_GATEWAY_FD") + if fd is None: + return {"host": "127.0.0.1", "port": int(sys.argv[2])} + from litellm.proxy.rust_forward.mode import rust_forward_mode, stop_with_supervisor + + stop_with_supervisor(rust_forward_mode(os.environ)) + pathlib.Path(os.environ["MARKERS"], "python.pid").write_text(str(os.getpid())) + return {"fd": int(fd)} + + +if __name__ == "__main__": + uvicorn.run("app:app", app_dir=os.path.dirname(__file__), workers=int(sys.argv[1]), log_level="warning", **listener()) +""" + +GATEWAY: Final = """ +import asyncio +import os +import pathlib +import signal +import threading + + +def wait_for_eof(loop, stop): + while os.read(0, 4096): + pass + loop.call_soon_threadsafe(stop.set) + + +async def pipe(reader, writer): + try: + while data := await reader.read(65536): + writer.write(data) + await writer.drain() + except OSError: + pass + finally: + writer.close() + + +async def main(): + pathlib.Path(os.environ["MARKERS"], "rust.pid").write_text(str(os.getpid())) + host, port = os.environ["LITELLM_FORWARD_UPSTREAM"].removeprefix("http://").rsplit(":", 1) + loop = asyncio.get_running_loop() + stop = asyncio.Event() + for signum in (signal.SIGTERM, signal.SIGINT): + loop.add_signal_handler(signum, stop.set) + threading.Thread(target=wait_for_eof, args=(loop, stop), daemon=True).start() + relays = set() + + async def relay(client_reader, client_writer): + relays.add(asyncio.current_task()) + try: + upstream_reader, upstream_writer = await asyncio.open_connection(host, int(port)) + await asyncio.gather(pipe(client_reader, upstream_writer), pipe(upstream_reader, client_writer)) + finally: + relays.discard(asyncio.current_task()) + + server = await asyncio.start_server( + relay, os.environ["LITELLM_FORWARD_HOST"], int(os.environ["LITELLM_FORWARD_PORT"]) + ) + await stop.wait() + server.close() + if relays: + await asyncio.wait(relays, timeout=float(os.environ["LITELLM_FORWARD_SHUTDOWN_DRAIN"])) + + +asyncio.run(main()) +""" + +SUPERVISE: Final = """ +import sys +from pathlib import Path + +from litellm.proxy.rust_forward.server import Listener, PythonServer, Server +from litellm.proxy.rust_forward.supervisor import supervise + +workers, port, scripts = sys.argv[1], int(sys.argv[2]), Path(sys.argv[3]) +server = PythonServer(Server.UVICORN, Listener.TCP, "127.0.0.1", port, 5, 30.0, True, None) +result = supervise(server, (sys.executable, str(scripts / "serve.py"), workers, str(port)), scripts_dir=str(scripts)) +raise SystemExit(result if isinstance(result, int) else 3) +""" + +EXIT_SECONDS: Final = 15 +TIME_TOLERANCE_SECONDS: Final = 5 + + +class Scenario(Enum): + TERM = "SIGTERM" + INT = "SIGINT" + INT_TWICE_WHILE_STREAMING = "SIGINT twice during an open stream" + HUP = "SIGHUP" + + +@dataclass(frozen=True, slots=True) +class Outcome: + status: int | None + survived_hup: bool + + +@dataclass(frozen=True, slots=True) +class Run: + outcome: Outcome + seconds: float + + +def _scripts(directory: Path) -> Path: + (directory / "app.py").write_text(APP) + (directory / "serve.py").write_text(SERVE) + (directory / "supervise.py").write_text(SUPERVISE) + gateway: Final = directory / "litellm-gateway" + gateway.write_text(f"#!{sys.executable}\n{GATEWAY}") + gateway.chmod(0o755) + return directory + + +def _serving(port: int) -> bool: + try: + return httpx.get(f"http://127.0.0.1:{port}/", timeout=1, trust_env=False).text == "ok" + except httpx.TransportError: + return False + + +def _exit_status(process: subprocess.Popen[bytes], seconds: float) -> int | None: + try: + return process.wait(timeout=seconds) + except subprocess.TimeoutExpired: + return None + + +def _streaming(port: int) -> socket.socket: + stream: Final = socket.create_connection(("127.0.0.1", port), timeout=10) + stream.sendall(b"GET /stream HTTP/1.1\r\nHost: test\r\n\r\n") + assert eventually(lambda: stream.recv(4096), lambda received: b"tick" in received, seconds=10) + return stream + + +def _signal_after_ready(process: subprocess.Popen[bytes], port: int, scenario: Scenario) -> Run: + assert eventually(lambda: _serving(port), bool, seconds=120) + started: Final = time.monotonic() + match scenario: + case Scenario.TERM | Scenario.INT: + process.send_signal(signal.SIGTERM if scenario is Scenario.TERM else signal.SIGINT) + return Run(Outcome(_exit_status(process, EXIT_SECONDS), False), time.monotonic() - started) + case Scenario.INT_TWICE_WHILE_STREAMING: + with _streaming(port): + process.send_signal(signal.SIGINT) + time.sleep(0.5) + process.send_signal(signal.SIGINT) + status: Final = _exit_status(process, EXIT_SECONDS) + return Run(Outcome(status, False), time.monotonic() - started) + case Scenario.HUP: + process.send_signal(signal.SIGHUP) + after_hup: Final = _exit_status(process, 3) + if after_hup is not None: + return Run(Outcome(after_hup, False), time.monotonic() - started) + assert eventually(lambda: _serving(port), bool, seconds=30) + process.send_signal(signal.SIGTERM) + return Run(Outcome(_exit_status(process, EXIT_SECONDS), True), time.monotonic() - started) + + +def _run(argv: tuple[str, ...], environment: dict[str, str], act: Callable[[subprocess.Popen[bytes]], Run]) -> Run: + with subprocess.Popen(argv, env=environment, start_new_session=True) as process: + try: + return act(process) + finally: + if process.poll() is None: + for leftover in (*psutil.Process(process.pid).children(recursive=True), psutil.Process(process.pid)): + leftover.kill() + + +def _environment(markers: Path) -> dict[str, str]: + return {**os.environ, "MARKERS": str(markers), "PYTHONUNBUFFERED": "1"} + + +@pytest.mark.parametrize("workers", [1, 2]) +@pytest.mark.parametrize("scenario", list(Scenario), ids=[scenario.name for scenario in Scenario]) +def test_the_supervisor_ends_as_python_alone_would(tmp_path: Path, workers: int, scenario: Scenario) -> None: + scripts: Final = _scripts(tmp_path) + alone_port: Final = allocate_port() + alone: Final = _run( + (sys.executable, str(scripts / "serve.py"), str(workers), str(alone_port)), + _environment(tmp_path), + lambda process: _signal_after_ready(process, alone_port, scenario), + ) + supervised_port: Final = allocate_port() + supervised: Final = _run( + (sys.executable, str(scripts / "supervise.py"), str(workers), str(supervised_port), str(scripts)), + _environment(tmp_path), + lambda process: _signal_after_ready(process, supervised_port, scenario), + ) + assert supervised.outcome == alone.outcome + assert abs(supervised.seconds - alone.seconds) < TIME_TOLERANCE_SECONDS, (alone, supervised) + + +def _alive(pid: int) -> bool: + try: + return psutil.Process(pid).status() != psutil.STATUS_ZOMBIE + except psutil.NoSuchProcess: + return False + + +@pytest.mark.parametrize("workers", [1, 2]) +def test_children_exit_when_the_supervisor_is_killed(tmp_path: Path, workers: int) -> None: + scripts: Final = _scripts(tmp_path) + port: Final = allocate_port() + + def kill_supervisor(process: subprocess.Popen[bytes]) -> Run: + assert eventually(lambda: _serving(port), bool, seconds=120) + python: Final = int((tmp_path / "python.pid").read_text()) + rust: Final = int((tmp_path / "rust.pid").read_text()) + descendants: Final = (python, rust, *(child.pid for child in psutil.Process(python).children(recursive=True))) + started: Final = time.monotonic() + process.kill() + process.wait() + assert eventually(lambda: tuple(pid for pid in descendants if _alive(pid)), lambda alive: not alive) == () + return Run(Outcome(None, False), time.monotonic() - started) + + _run( + (sys.executable, str(scripts / "supervise.py"), str(workers), str(port), str(scripts)), + _environment(tmp_path), + kill_supervisor, + ) diff --git a/tests/unit/proxy/rust_forward/test_mode.py b/tests/unit/proxy/rust_forward/test_mode.py index 9d57d4806c3..34329bacbb2 100644 --- a/tests/unit/proxy/rust_forward/test_mode.py +++ b/tests/unit/proxy/rust_forward/test_mode.py @@ -1,3 +1,5 @@ +import os +import threading from typing import Final import pytest @@ -15,6 +17,7 @@ from litellm.proxy.rust_forward.mode import ( gunicorn_bind, rust_forward_mode, server_proxy_headers, + stop_with_supervisor, supervised_environment, uvicorn_options, ) @@ -110,3 +113,19 @@ def test_assigned_port_serves_only_the_loopback_hop_without_trusting_it() -> Non assert (config.host, config.port, config.fd, config.proxy_headers) == ("127.0.0.1", 45556, None, False) assert gunicorn_bind(mode, "0.0.0.0:4000") == "127.0.0.1:45556" assert server_proxy_headers(mode) is False + + +@pytest.mark.parametrize( + ("mode", "stops"), + [(Supervised(InheritedFd(17), False, "127.0.0.1"), True), (Disabled(), False), (Supervisor(), False)], +) +def test_only_a_supervised_server_stops_once_its_supervisor_is_gone(mode: RustForwardMode, stops: bool) -> None: + read_end, write_end = os.pipe() + stopped: Final = threading.Event() + try: + stop_with_supervisor(mode, read_end, stopped.set) + assert not stopped.wait(0.2) + os.close(write_end) + assert stopped.wait(5 if stops else 0.5) is stops + finally: + os.close(read_end) diff --git a/tests/unit/proxy/rust_forward/test_supervisor.py b/tests/unit/proxy/rust_forward/test_supervisor.py index ea69ce86951..970d654ebde 100644 --- a/tests/unit/proxy/rust_forward/test_supervisor.py +++ b/tests/unit/proxy/rust_forward/test_supervisor.py @@ -3,10 +3,13 @@ import signal import socket import subprocess import sys +from dataclasses import dataclass, replace from pathlib import Path from typing import Final import pytest +from hypothesis import given +from hypothesis import strategies as st from litellm.constants import GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS from litellm.proxy.proxy_cli import run_server @@ -21,20 +24,20 @@ from litellm.proxy.rust_forward.platform.windows import ( ) from litellm.proxy.rust_forward.server import Listener, PythonServer, Server from litellm.proxy.rust_forward.supervisor import ( + FAILURE, + Action, + Caught, Child, + ChildExited, Event, - ExitCode, - Exited, - Finished, - Ready, - Running, - Signal, - Start, + Lifecycle, + PythonListening, + Send, + Spawn, + Spawned, StartupError, - Stop, - Stopping, - Transition, boot, + exit_status, gateway_binary, python_spec, rust_spec, @@ -200,50 +203,124 @@ def test_rust_opens_the_public_port_only_once_python_listens(tmp_path: Path, mon assert (tmp_path / "rust_saw_python_listening").read_text() == "True" -def _fold(events: tuple[Event, ...]) -> Transition: - def advance(acc: Transition, event: Event) -> Transition: - state, done = acc - following, actions = step(state, event) - return following, (*done, *actions) +@dataclass(frozen=True, slots=True) +class World: + """What the operating system knows about the children, kept independently of ``Lifecycle``.""" - return functools.reduce(advance, events, boot()) + alive: frozenset[Child] = frozenset({Child.PYTHON}) + spawned: frozenset[Child] = frozenset({Child.PYTHON}) + reported_listening: bool = False + stop_requested: bool = False + python_status: int | None = None + rust_lost: bool = False -RUST: Final = Child.RUST -PYTHON: Final = Child.PYTHON -PYTHON_READY: Final = Ready(PYTHON) -STARTED: Final = (Start(PYTHON), Start(RUST)) -BOTH_TERM: Final = (*STARTED, Signal(RUST, TERM), Signal(PYTHON, TERM)) -CLEAN: Final = Finished(ExitCode.CLEAN) -FAILED: Final = Finished(ExitCode.CHILD_FAILED) +def _expected_exit(world: World) -> int | None: + if world.alive or world.python_status is None: + return None + return FAILURE if world.rust_lost else world.python_status + + +def _expected_relay(world: World, signum: signal.Signals) -> tuple[Action, ...]: + to_rust: Final = (Send(Child.RUST, signum),) if signum in STOPS and Child.RUST in world.alive else () + to_python: Final = (Send(Child.PYTHON, signum),) if Child.PYTHON in world.alive else () + return (*to_rust, *to_python) + + +def _expected_on_exit(world: World, child: Child) -> tuple[Action, ...]: + match child: + case Child.PYTHON: + return (Send(Child.RUST, TERM),) if Child.RUST in world.alive else () + case Child.RUST: + serving: Final = Child.PYTHON in world.alive and not world.stop_requested + return (Send(Child.PYTHON, TERM),) if serving else () + + +def _after(world: World, event: Event) -> World: + match event: + case Caught(signum): + return replace(world, stop_requested=world.stop_requested or signum in STOPS) + case PythonListening(): + return replace(world, reported_listening=world.reported_listening or Child.PYTHON in world.alive) + case ChildExited(child, status) if child in world.alive: + lost: Final = child is Child.RUST and Child.PYTHON in world.alive and not world.stop_requested + return replace( + world, + alive=world.alive - {child}, + python_status=status if child is Child.PYTHON else world.python_status, + rust_lost=world.rust_lost or lost, + ) + case ChildExited(): + return world + + +def _spawned(world: World, actions: tuple[Action, ...]) -> World: + started: Final = frozenset(action.child for action in actions if isinstance(action, Spawn)) + return replace(world, alive=world.alive | started, spawned=world.spawned | started) + + +SIGNALS: Final = POSIX.relayed_signals +STOPS: Final = POSIX.stop_signals +EVENTS: Final = st.lists( + st.one_of( + st.sampled_from(SIGNALS).map(Caught), + st.just(PythonListening()), + st.builds(ChildExited, st.sampled_from(Child), st.integers(-15, 3)), + ), + max_size=40, +) + + +@given(EVENTS) +def test_lifecycle_relays_every_signal_and_ends_with_pythons_status(events: list[Event]) -> None: + state, actions = boot() + world = World() + assert actions == (Spawn(Child.PYTHON),) + for event in events: + before = _after(world, event) + state, actions = step(state, event, STOPS) + match event: + case Caught(signum): + assert actions == _expected_relay(before, signum) + case PythonListening(): + starts_rust = Child.PYTHON in world.alive and not world.reported_listening and not world.stop_requested + assert actions == ((Spawn(Child.RUST),) if starts_rust else ()) + case ChildExited(child) if child in world.alive: + assert actions == _expected_on_exit(world, child) + case ChildExited(): + assert actions == () + assert all(action.child in before.alive for action in actions if isinstance(action, Send)) + assert all(action.child not in before.spawned for action in actions if isinstance(action, Spawn)) + world = _spawned(before, actions) + assert exit_status(state) == _expected_exit(world) + + +@given(EVENTS) +def test_rust_starts_at_most_once_and_only_while_python_listens_unstopped(events: list[Event]) -> None: + def advance(acc: tuple[Lifecycle, tuple[tuple[Event, Lifecycle, Action], ...]], event: Event): + state, seen = acc + following, actions = step(state, event, STOPS) + return following, (*seen, *((event, state, action) for action in actions)) + + _, history = functools.reduce(advance, events, (boot()[0], ())) + spawns: Final = tuple((event, state) for event, state, action in history if isinstance(action, Spawn)) + assert len(spawns) <= 1 + assert all( + isinstance(event, PythonListening) and isinstance(state.python, Spawned) and not state.stop_requested + for event, state in spawns + ) @pytest.mark.parametrize( ("events", "expected"), [ - ((), (Running((PYTHON,)), (Start(PYTHON),))), - ((Ready(RUST),), (Running((PYTHON,)), (Start(PYTHON),))), - ((PYTHON_READY,), (Running((PYTHON, RUST)), STARTED)), - ((PYTHON_READY, PYTHON_READY, Ready(RUST)), (Running((PYTHON, RUST)), STARTED)), - ((PYTHON_READY, Stop(TERM)), (Stopping(frozenset({RUST, PYTHON})), BOTH_TERM)), - ((PYTHON_READY, Stop(TERM), Exited(RUST, 0), Exited(PYTHON, 0)), (CLEAN, BOTH_TERM)), - ((PYTHON_READY, Stop(TERM), Exited(PYTHON, 0), Exited(RUST, 0)), (CLEAN, BOTH_TERM)), - ( - (PYTHON_READY, Stop(INT), Stop(INT), Exited(RUST, 0), Exited(PYTHON, 0)), - (CLEAN, (*STARTED, Signal(RUST, INT), Signal(PYTHON, INT))), - ), - ((PYTHON_READY, Stop(TERM), Exited(RUST, 0), Exited(RUST, 0)), (FAILED, BOTH_TERM)), - ((PYTHON_READY, Stop(TERM), Exited(RUST, 1)), (FAILED, BOTH_TERM)), - ((PYTHON_READY, Stop(TERM), Exited(RUST, 0), Exited(PYTHON, 1)), (FAILED, BOTH_TERM)), - ((PYTHON_READY, Exited(PYTHON, 0)), (FAILED, STARTED)), - ((PYTHON_READY, Exited(RUST, 1)), (FAILED, STARTED)), - ((PYTHON_READY, Exited(RUST, 1), Stop(TERM), PYTHON_READY, Exited(PYTHON, 0)), (FAILED, STARTED)), - ((Exited(PYTHON, 1),), (FAILED, (Start(PYTHON),))), - ((Stop(TERM), PYTHON_READY, Exited(PYTHON, 0)), (CLEAN, (Start(PYTHON), Signal(PYTHON, TERM)))), - ((Stop(TERM), Exited(PYTHON, 1)), (FAILED, (Start(PYTHON), Signal(PYTHON, TERM)))), + ((PythonListening(), Caught(TERM), ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, -15)), -15), + ((PythonListening(), Caught(INT), Caught(INT), ChildExited(Child.PYTHON, 1), ChildExited(Child.RUST, 0)), 1), + ((PythonListening(), Caught(signal.SIGHUP), ChildExited(Child.PYTHON, -1), ChildExited(Child.RUST, 0)), -1), + ((PythonListening(), ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, 0)), FAILURE), + ((ChildExited(Child.PYTHON, 3),), 3), ], ) -def test_children_start_in_order_once_ready_and_stop_together_in_reverse( - events: tuple[Event, ...], expected: Transition -) -> None: - assert _fold(events) == expected +def test_supervisor_exits_as_python_did_unless_rust_was_lost(events: tuple[Event, ...], expected: int) -> None: + final: Final = functools.reduce(lambda state, event: step(state, event, STOPS)[0], events, boot()[0]) + assert exit_status(final) == expected diff --git a/tests/unit/proxy/test_proxy_cli.py b/tests/unit/proxy/test_proxy_cli.py index c7d83a00188..7a30d78fee4 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -4349,6 +4349,15 @@ class TestRustForwardServer: assert mock_uvicorn_run.call_args.kwargs["ssl_keyfile"] == "key.pem" assert mock_uvicorn_run.call_args.kwargs["port"] == 45999 + def test_an_invalid_rust_switch_stops_startup_with_one_line_naming_it(self) -> None: + from click.testing import CliRunner + + with patch.dict(os.environ, {"LITELLM_RUST": "sometimes"}): + result: Final = CliRunner().invoke(run_server, ["--port", "45999"]) + lines: Final = result.output.splitlines() + assert (result.exit_code, len(lines)) == (1, 1) + assert lines[0].startswith("Error: LITELLM_RUST is invalid: ") + @pytest.mark.parametrize( ("argv", "environ"), [(["--keepalive_timeout", "2"], {}), ([], {"KEEPALIVE_TIMEOUT": "7"}), ([], {})], diff --git a/tests/unit/rust_bridge/test_configuration.py b/tests/unit/rust_bridge/test_configuration.py index 8116cf98076..6ade85c0367 100644 --- a/tests/unit/rust_bridge/test_configuration.py +++ b/tests/unit/rust_bridge/test_configuration.py @@ -8,6 +8,7 @@ from concurrent.futures import ThreadPoolExecutor from typing import Final import pytest +from pydantic import ValidationError from litellm.rust_bridge import configuration @@ -106,8 +107,8 @@ 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", ("", " ")) +def test_blank_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 +117,22 @@ def test_invalid_environment_value_is_ignored(monkeypatch: pytest.MonkeyPatch, v assert configuration.rust_enabled() is True +@pytest.mark.parametrize("value", ("sometimes", "2", "enabled")) +def test_invalid_environment_value_is_rejected(monkeypatch: pytest.MonkeyPatch, value: str) -> None: + monkeypatch.setenv("LITELLM_RUST", value) + configuration.rust(True) + + with pytest.raises(ValidationError): + configuration.rust_enabled() + + +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 From 9f6c9d2246329ca70e71e61ce629a78fd86405af Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 6 Nov 2026 05:28:40 +0900 Subject: [PATCH 04/10] fix(proxy): keep the strict lint and test quality gates from growing Signal the lifecycle test processes through the integration harness and give stop_with_supervisor a named default instead of a call. Co-Authored-By: Claude Opus 5.5 --- litellm/proxy/rust_forward/mode.py | 11 +++++------ tests/integration/_support/process.py | 5 +++++ .../configuration/test_rust_forward_lifecycle.py | 16 ++++++++-------- tests/unit/proxy/rust_forward/test_supervisor.py | 4 +++- 4 files changed, 21 insertions(+), 15 deletions(-) diff --git a/litellm/proxy/rust_forward/mode.py b/litellm/proxy/rust_forward/mode.py index 4b4cd77e05b..4f995d4a602 100644 --- a/litellm/proxy/rust_forward/mode.py +++ b/litellm/proxy/rust_forward/mode.py @@ -4,7 +4,6 @@ import signal import threading from collections.abc import Callable, Mapping from dataclasses import dataclass -from functools import partial from types import MappingProxyType from typing import Final, TypeAlias @@ -109,11 +108,11 @@ def _stop_when_closed(fd: int, stop: Callable[[], None]) -> None: stop() -def stop_with_supervisor( - mode: RustForwardMode, - fd: int = 0, - stop: Callable[[], None] = partial(signal.raise_signal, signal.SIGTERM), -) -> None: +def _terminate_self() -> None: + signal.raise_signal(signal.SIGTERM) + + +def stop_with_supervisor(mode: RustForwardMode, fd: int = 0, stop: Callable[[], None] = _terminate_self) -> None: """Shut the supervised server down gracefully once the supervisor's end of its stdin closes, so it never outlives a supervisor that was killed. Call it only in the process the supervisor started, not in its workers.""" match mode: diff --git a/tests/integration/_support/process.py b/tests/integration/_support/process.py index 262114ee32b..b8e0a262c41 100644 --- a/tests/integration/_support/process.py +++ b/tests/integration/_support/process.py @@ -145,6 +145,11 @@ def signal_group(group: int, action: int) -> None: pass +def send_signal(process: subprocess.Popen[bytes], signum: signal.Signals) -> None: + with suppress(ProcessLookupError): + process.send_signal(signum) + + def graceful_stop_seconds() -> float: return readiness_for(import_seconds(), 1) diff --git a/tests/integration/configuration/test_rust_forward_lifecycle.py b/tests/integration/configuration/test_rust_forward_lifecycle.py index cc186af98d7..3c639ce6d9a 100644 --- a/tests/integration/configuration/test_rust_forward_lifecycle.py +++ b/tests/integration/configuration/test_rust_forward_lifecycle.py @@ -16,6 +16,7 @@ import pytest from tests.integration._support.client import eventually from tests.integration._support.ports import allocate_port +from tests.integration._support.process import kill_stray_process, send_signal pytestmark = pytest.mark.skipif(sys.platform == "win32", reason="drives POSIX signals") @@ -186,22 +187,22 @@ def _signal_after_ready(process: subprocess.Popen[bytes], port: int, scenario: S started: Final = time.monotonic() match scenario: case Scenario.TERM | Scenario.INT: - process.send_signal(signal.SIGTERM if scenario is Scenario.TERM else signal.SIGINT) + send_signal(process, signal.SIGTERM if scenario is Scenario.TERM else signal.SIGINT) return Run(Outcome(_exit_status(process, EXIT_SECONDS), False), time.monotonic() - started) case Scenario.INT_TWICE_WHILE_STREAMING: with _streaming(port): - process.send_signal(signal.SIGINT) + send_signal(process, signal.SIGINT) time.sleep(0.5) - process.send_signal(signal.SIGINT) + send_signal(process, signal.SIGINT) status: Final = _exit_status(process, EXIT_SECONDS) return Run(Outcome(status, False), time.monotonic() - started) case Scenario.HUP: - process.send_signal(signal.SIGHUP) + send_signal(process, signal.SIGHUP) after_hup: Final = _exit_status(process, 3) if after_hup is not None: return Run(Outcome(after_hup, False), time.monotonic() - started) assert eventually(lambda: _serving(port), bool, seconds=30) - process.send_signal(signal.SIGTERM) + send_signal(process, signal.SIGTERM) return Run(Outcome(_exit_status(process, EXIT_SECONDS), True), time.monotonic() - started) @@ -212,7 +213,7 @@ def _run(argv: tuple[str, ...], environment: dict[str, str], act: Callable[[subp finally: if process.poll() is None: for leftover in (*psutil.Process(process.pid).children(recursive=True), psutil.Process(process.pid)): - leftover.kill() + kill_stray_process(leftover.pid) def _environment(markers: Path) -> dict[str, str]: @@ -257,8 +258,7 @@ def test_children_exit_when_the_supervisor_is_killed(tmp_path: Path, workers: in rust: Final = int((tmp_path / "rust.pid").read_text()) descendants: Final = (python, rust, *(child.pid for child in psutil.Process(python).children(recursive=True))) started: Final = time.monotonic() - process.kill() - process.wait() + kill_stray_process(process.pid) assert eventually(lambda: tuple(pid for pid in descendants if _alive(pid)), lambda alive: not alive) == () return Run(Outcome(None, False), time.monotonic() - started) diff --git a/tests/unit/proxy/rust_forward/test_supervisor.py b/tests/unit/proxy/rust_forward/test_supervisor.py index 970d654ebde..1931b11c7a9 100644 --- a/tests/unit/proxy/rust_forward/test_supervisor.py +++ b/tests/unit/proxy/rust_forward/test_supervisor.py @@ -8,7 +8,7 @@ from pathlib import Path from typing import Final import pytest -from hypothesis import given +from hypothesis import example, given from hypothesis import strategies as st from litellm.constants import GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS @@ -272,6 +272,8 @@ EVENTS: Final = st.lists( @given(EVENTS) +@example([Caught(TERM), Caught(signal.SIGHUP), PythonListening()]) +@example([PythonListening(), Caught(TERM), Caught(signal.SIGWINCH), ChildExited(Child.RUST, 0)]) def test_lifecycle_relays_every_signal_and_ends_with_pythons_status(events: list[Event]) -> None: state, actions = boot() world = World() From e74261edc17f5f5ecc4c274dbd7320edcd4a5044 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 6 Nov 2026 05:44:44 +0900 Subject: [PATCH 05/10] refactor(proxy): type the Rust forwarding ASGI boundary precisely Everything in rust_forward now speaks uvicorn's ASGI3Application, and served() is the one cast from Starlette's ASGIApp, since the two type the same callable with unassignable scope, receive and send types. Co-Authored-By: Claude Opus 5.5 --- litellm/proxy/rust_forward/peer.py | 26 ++++++++++++++++------ tests/unit/proxy/rust_forward/test_peer.py | 5 ++--- 2 files changed, 21 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/rust_forward/peer.py b/litellm/proxy/rust_forward/peer.py index 2d03d6ca35c..cf891495c2b 100644 --- a/litellm/proxy/rust_forward/peer.py +++ b/litellm/proxy/rust_forward/peer.py @@ -1,16 +1,27 @@ -from collections.abc import Awaitable, Callable -from typing import TYPE_CHECKING, Final, TypeAlias +from typing import ( + TYPE_CHECKING, + Final, + TypeAlias, + cast, # noqa: TID251 # the one Starlette-to-uvicorn ASGI conversion in served +) from typing_extensions import assert_never from litellm.proxy.rust_forward.mode import RUST_HOP, Disabled, RustForwardMode, Supervised, Supervisor if TYPE_CHECKING: - from uvicorn._types import ASGIReceiveCallable, ASGISendCallable, Scope, WWWScope + import starlette.types + from uvicorn._types import ASGI3Application, ASGIReceiveCallable, ASGISendCallable, Scope, WWWScope PEER_HEADER: Final = b"x-litellm-rust-peer" -ASGIApp: TypeAlias = Callable[..., Awaitable[None]] +ASGIApp: TypeAlias = "ASGI3Application" + + +def served(app: "starlette.types.ASGIApp") -> ASGIApp: + return cast( # cast-ok: Starlette and uvicorn type the same ASGI 3 callable with unassignable scope/receive/send + "ASGI3Application", app + ) def _restore_peer(scope: "WWWScope") -> None: @@ -47,11 +58,12 @@ class RustPeerMiddleware: await self.app(scope, receive, send) -def wrap(app: ASGIApp, mode: RustForwardMode) -> ASGIApp: +def wrap(app: "starlette.types.ASGIApp", mode: RustForwardMode) -> ASGIApp: + converted: Final = served(app) match mode: case Supervised(): - return RustPeerMiddleware(app, mode) + return RustPeerMiddleware(converted, mode) case Disabled() | Supervisor(): - return app + return converted case _: assert_never(mode) diff --git a/tests/unit/proxy/rust_forward/test_peer.py b/tests/unit/proxy/rust_forward/test_peer.py index ee625c4a4be..b0c59641e7f 100644 --- a/tests/unit/proxy/rust_forward/test_peer.py +++ b/tests/unit/proxy/rust_forward/test_peer.py @@ -21,8 +21,7 @@ if TYPE_CHECKING: from litellm.proxy.proxy_server import app as proxy_app from litellm.proxy.rust_forward.mode import InheritedFd, Supervised -from litellm.proxy.rust_forward.peer import wrap - +from litellm.proxy.rust_forward.peer import RustPeerMiddleware Headers = tuple[tuple[bytes, bytes], ...] @@ -47,7 +46,7 @@ async def _seen_by_app( seen.set_result((scope["client"], scope["scheme"], tuple(scope["headers"]))) direct: Final = ProxyHeadersMiddleware(probe, mode.trusted_proxies) if mode.proxy_headers else probe - app: Final = wrap(probe, mode) if behind_rust else direct + app: Final = RustPeerMiddleware(probe, mode) if behind_rust else direct scope: Final[HTTPScope] = { "type": "http", "asgi": {"version": "3.0", "spec_version": "2.3"}, From cc4bce9cbed53a643b72f23a688f9c20ebb91f2f Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 6 Nov 2026 13:34:34 +0900 Subject: [PATCH 06/10] refactor(proxy): lean on the stdlib, uvicorn and pydantic in Rust forwarding The supervisor runs its children through one contextlib.AsyncExitStack loop instead of recursive async with blocks, and payload-free states and modes are Enum members The Python child's env contract is the Supervised type itself, written and read through one pydantic TypeAdapter, so the trusted proxies it carries are always set. Each entry point resolves them from its server library, uvicorn's Config or gunicorn's, instead of a default copied into this package gateway.launch validates uvicorn's parsed CLI options with one pydantic model. TLS now means what each server serves TLS with: a key or cert for uvicorn's CLI, both for proxy_cli. A cipher list alone, or one half of the pair under proxy_cli, no longer refuses forwarding, since Python alone serves plain HTTP there --- gateway/launch.py | 46 ++----- litellm/proxy/proxy_cli.py | 10 +- litellm/proxy/rust_forward/AGENTS.md | 4 +- litellm/proxy/rust_forward/mode.py | 75 +++-------- litellm/proxy/rust_forward/peer.py | 4 +- litellm/proxy/rust_forward/server.py | 42 +++++- litellm/proxy/rust_forward/supervisor.py | 123 ++++++------------ .../test_rust_forward_lifecycle.py | 14 +- tests/unit/gateway/test_launch.py | 26 ++-- tests/unit/proxy/rust_forward/test_mode.py | 69 +++++----- tests/unit/proxy/rust_forward/test_peer.py | 6 +- .../proxy/rust_forward/test_supervisor.py | 45 ++++--- tests/unit/proxy/test_proxy_cli.py | 47 +++++-- 13 files changed, 252 insertions(+), 259 deletions(-) diff --git a/gateway/launch.py b/gateway/launch.py index 979c40984d9..1adb2ef1034 100644 --- a/gateway/launch.py +++ b/gateway/launch.py @@ -80,40 +80,23 @@ def _serve(argv: Sequence[str]) -> None: def python_server(argv: Sequence[str]) -> "PythonServer": """What uvicorn would serve the public port with, read from the same flags and environment it reads.""" - from pydantic import TypeAdapter - - from litellm.proxy.rust_forward.server import Listener, PythonServer, Server + from litellm.proxy.rust_forward.server import PythonServer, Server, UvicornCliParams, uvicorn_trusted_proxies from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager - params: Final[Mapping[str, object]] = uvicorn_main.make_context( - "uvicorn", list(uvicorn_argv(argv, os.environ)) - ).params - tls: Final = ( - params["ssl_keyfile"] is not None - or params["ssl_certfile"] is not None - or any(arg.split("=")[0] == "--ssl-ciphers" for arg in argv) + params: Final = UvicornCliParams.model_validate( + uvicorn_main.make_context("uvicorn", list(uvicorn_argv(argv, os.environ))).params ) - listener: Final = ( - Listener.TLS - if tls - else Listener.UNIX_SOCKET - if params["uds"] is not None - else Listener.INHERITED_FD - if params["fd"] is not None - else Listener.TCP - ) - graceful_timeout: Final = TypeAdapter[int | None](int | None).validate_python(params["timeout_graceful_shutdown"]) return PythonServer( server=Server.UVICORN, - listener=listener, - host=TypeAdapter(str).validate_python(params["host"]), - port=TypeAdapter(int).validate_python(params["port"]), - keepalive_seconds=TypeAdapter(int).validate_python(params["timeout_keep_alive"]), + listener=params.listener, + host=params.host, + port=params.port, + keepalive_seconds=params.timeout_keep_alive, graceful_timeout_seconds=GracefulShutdownManager.get_timeout() - if graceful_timeout is None - else graceful_timeout, - proxy_headers=TypeAdapter(bool).validate_python(params["proxy_headers"]), - trusted_proxies=TypeAdapter[str | None](str | None).validate_python(params["forwarded_allow_ips"]), + if params.timeout_graceful_shutdown is None + else params.timeout_graceful_shutdown, + proxy_headers=params.proxy_headers, + trusted_proxies=uvicorn_trusted_proxies(params.forwarded_allow_ips), ) @@ -132,9 +115,8 @@ def configure_database() -> None: def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -> None: from litellm.proxy.rust_forward.mode import ( - Disabled, + Mode, Supervised, - Supervisor, rust_forward_mode, stop_with_supervisor, uvicorn_flags, @@ -143,14 +125,14 @@ def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) - mode: Final = rust_forward_mode(os.environ) match mode: - case Supervisor() if not isinstance( + case Mode.SUPERVISOR if not isinstance( exit_code := supervise( python_server(argv), (sys.executable, "-m", "gateway.launch", *argv), before_start=configure_database ), StartupError, ): sys.exit(exit_code) - case Disabled() | Supervisor(): + case Mode.DISABLED | Mode.SUPERVISOR: configure_database() serve(uvicorn_argv(argv, os.environ)) case Supervised(): diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 0338b7261b9..632a71ec5c8 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -854,7 +854,7 @@ class ProxyInitializationHelpers: import uvicorn from gunicorn.config import Config as GunicornConfig - from litellm.proxy.rust_forward.server import Listener, PythonServer, Server + from litellm.proxy.rust_forward.server import Listener, PythonServer, Server, uvicorn_trusted_proxies from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager server: Final = ( @@ -883,7 +883,7 @@ class ProxyInitializationHelpers: proxy_headers=True, trusted_proxies=",".join(TypeAdapter(list[str]).validate_python(GunicornConfig().forwarded_allow_ips)) if server is Server.GUNICORN - else None, + else uvicorn_trusted_proxies(None), ) @staticmethod @@ -1359,7 +1359,7 @@ def run_server( from pydantic import ValidationError from litellm.proxy.rust_forward.mode import ( - Supervisor, + Mode, environment_error, gunicorn_bind, rust_forward_mode, @@ -1373,7 +1373,7 @@ def run_server( except ValidationError as error: raise click.ClickException(environment_error(error)) from None stop_with_supervisor(forward_mode) - if isinstance(forward_mode, Supervisor) and not any( + if forward_mode is Mode.SUPERVISOR and not any( (version, validate_config, health, test, skip_server_startup, check_migrations_against_existing_data) ): from litellm.proxy.rust_forward.supervisor import StartupError, supervise @@ -1385,7 +1385,7 @@ def run_server( run_gunicorn=run_gunicorn, run_hypercorn=run_hypercorn, run_granian=run_granian, - tls=any((ssl_keyfile_path, ssl_certfile_path, ciphers)), + tls=None not in (ssl_keyfile_path, ssl_certfile_path), ) supervised: Final = supervise(server, (sys.executable, "-m", "litellm.proxy.proxy_cli", *sys.argv[1:])) if not isinstance(supervised, StartupError): diff --git a/litellm/proxy/rust_forward/AGENTS.md b/litellm/proxy/rust_forward/AGENTS.md index 2c76bf73674..907b7dc1a3c 100644 --- a/litellm/proxy/rust_forward/AGENTS.md +++ b/litellm/proxy/rust_forward/AGENTS.md @@ -11,9 +11,9 @@ Every OS-specific operation lives behind `platform.Platform`, which `supervise` Windows support exists only for running the proxy locally from the CLI. Nobody deploys the proxy on Windows, and the Windows design leans on that. A port picked by the supervisor can be taken by another process before Python binds it, `CTRL_BREAK_EVENT` needs the console a CLI run always has, and gunicorn is refused. Revisit those choices before treating Windows as a deployment target -Each environment contract is written and read in one place. `mode.supervised_environment` and `mode.rust_forward_mode` are the two halves of the Python child's, and `supervisor.rust_spec` writes the `LITELLM_FORWARD_*` variables that `gateway-forward`'s `settings.rs` reads. Python owns every value Rust needs, so a new one gets a new variable in `rust_spec` +Each environment contract is written and read in one place. The Python child's is the `mode.Supervised` type itself: `mode.supervised_environment` serializes it into one variable and `mode.rust_forward_mode` validates it back, so a new field needs no parsing code. The trust it carries is resolved by the server library the entry point uses (`server.uvicorn_trusted_proxies` or gunicorn's config), never defaulted in this package. `supervisor.rust_spec` writes the `LITELLM_FORWARD_*` variables that `gateway-forward`'s `settings.rs` reads. Python owns every value Rust needs, so a new one gets a new variable in `rust_spec` -The supervisor is transparent to Python, so `python` with `LITELLM_RUST` behaves on signals and exits exactly as Python alone. Its whole lifecycle is `supervisor.step`, a pure function from a `Lifecycle` of facts and an event to the next facts and the actions to run, and `exit_status` reads the result from those facts. `_run` only executes actions, and the `async with` around each child owns its cleanup. Three rules decide everything there, and a change that breaks one of them breaks the equivalence `tests/integration/configuration/test_rust_forward_lifecycle.py` checks against Python alone +The supervisor is transparent to Python, so `python` with `LITELLM_RUST` behaves on signals and exits exactly as Python alone. Its whole lifecycle is `supervisor.step`, a pure function from a `Lifecycle` of facts and an event to the next facts and the actions to run, and `exit_status` reads the result from those facts. `_run` only executes actions, and the `contextlib.AsyncExitStack` it enters each child into owns their cleanup, last spawned first. Three rules decide everything there, and a change that breaks one of them breaks the equivalence `tests/integration/configuration/test_rust_forward_lifecycle.py` checks against Python alone - Signals pass through. Every signal in `Platform.relayed_signals` goes to Python unchanged, and the ones in `Platform.stop_signals` go to Rust first. The supervisor never reads meaning into a signal, so a repeated SIGINT reaches Python again and forces its exit as it would alone. Children run in their own session (`start_new_session`) only so that a terminal Ctrl-C reaches Python once, through the supervisor - Python's exit status is the result. When Python exits for any reason, Rust gets SIGTERM, and once Rust is gone the supervisor exits with Python's status, dying by the same signal when Python died by one. The one exception is Rust exiting while Python serves with no stop requested, which stops Python and exits 1 diff --git a/litellm/proxy/rust_forward/mode.py b/litellm/proxy/rust_forward/mode.py index 4f995d4a602..84c0a24d0fc 100644 --- a/litellm/proxy/rust_forward/mode.py +++ b/litellm/proxy/rust_forward/mode.py @@ -4,40 +4,28 @@ import signal import threading from collections.abc import Callable, Mapping from dataclasses import dataclass +from enum import Enum, auto from types import MappingProxyType from typing import Final, TypeAlias -from pydantic import BaseModel, ConfigDict, Field, ValidationError +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError from typing_extensions import assert_never from litellm.rust_bridge.configuration import ENV_NAME, RustSwitch -GATEWAY_FD: Final = "_LITELLM_GATEWAY_FD" -GATEWAY_PORT: Final = "_LITELLM_GATEWAY_PORT" -PROXY_HEADERS: Final = "_LITELLM_GATEWAY_PROXY_HEADERS" -TRUSTED_PROXIES: Final = "FORWARDED_ALLOW_IPS" -DEFAULT_TRUSTED_PROXIES: Final = "127.0.0.1" +SUPERVISED_ENV: Final = "_LITELLM_GATEWAY_SUPERVISED" RUST_HOP: Final = "127.0.0.1" -class _ForwardEnvironment(BaseModel): +class _Switch(BaseModel): model_config = ConfigDict(frozen=True) rust: RustSwitch = Field(None, validation_alias=ENV_NAME) - listener_fd: int | None = Field(None, validation_alias=GATEWAY_FD) - listener_port: int | None = Field(None, validation_alias=GATEWAY_PORT) - proxy_headers: bool = Field(True, validation_alias=PROXY_HEADERS) - trusted_proxies: str = Field(DEFAULT_TRUSTED_PROXIES, validation_alias=TRUSTED_PROXIES) -@dataclass(frozen=True, slots=True) -class Disabled: - pass - - -@dataclass(frozen=True, slots=True) -class Supervisor: - pass +class Mode(Enum): + DISABLED = auto() + SUPERVISOR = auto() @dataclass(frozen=True, slots=True) @@ -60,41 +48,20 @@ class Supervised: trusted_proxies: str -RustForwardMode: TypeAlias = Disabled | Supervisor | Supervised +RustForwardMode: TypeAlias = Mode | Supervised +_SUPERVISED: Final = TypeAdapter(Supervised) -def _listener_environment(listener: ChildListener) -> tuple[str, str]: - match listener: - case InheritedFd(fd): - return GATEWAY_FD, str(fd) - case AssignedPort(port): - return GATEWAY_PORT, str(port) - case _: - assert_never(listener) - - -def supervised_environment( - listener: ChildListener, proxy_headers: bool, trusted_proxies: str | None -) -> Mapping[str, str]: - name, value = _listener_environment(listener) - return MappingProxyType( - { - name: value, - PROXY_HEADERS: str(proxy_headers).lower(), - **({} if trusted_proxies is None else {TRUSTED_PROXIES: trusted_proxies}), - } - ) +def supervised_environment(supervised: Supervised) -> Mapping[str, str]: + return MappingProxyType({SUPERVISED_ENV: _SUPERVISED.dump_json(supervised).decode()}) def rust_forward_mode(environ: Mapping[str, str]) -> RustForwardMode: - parsed: Final = _ForwardEnvironment.model_validate(environ) - if parsed.listener_fd is not None: - return Supervised(InheritedFd(parsed.listener_fd), parsed.proxy_headers, parsed.trusted_proxies) - if parsed.listener_port is not None: - return Supervised(AssignedPort(parsed.listener_port), parsed.proxy_headers, parsed.trusted_proxies) - if parsed.rust: - return Supervisor() - return Disabled() + switch: Final = _Switch.model_validate(environ) + supervised: Final = environ.get(SUPERVISED_ENV) + if supervised is not None: + return _SUPERVISED.validate_json(supervised) + return Mode.SUPERVISOR if switch.rust else Mode.DISABLED def _wait_for_eof(fd: int) -> None: @@ -118,7 +85,7 @@ def stop_with_supervisor(mode: RustForwardMode, fd: int = 0, stop: Callable[[], match mode: case Supervised(): threading.Thread(target=_stop_when_closed, args=(fd, stop), name="supervisor-stdin", daemon=True).start() - case Disabled() | Supervisor(): + case Mode.DISABLED | Mode.SUPERVISOR: pass case _: assert_never(mode) @@ -151,7 +118,7 @@ def uvicorn_options(mode: RustForwardMode, factory: str) -> Mapping[str, object] return MappingProxyType( {"app": factory, "factory": True, **_uvicorn_listener(listener), "proxy_headers": False} ) - case Disabled() | Supervisor(): + case Mode.DISABLED | Mode.SUPERVISOR: return MappingProxyType({}) case _: assert_never(mode) @@ -171,7 +138,7 @@ def uvicorn_flags(mode: RustForwardMode) -> tuple[str, ...]: match mode: case Supervised(listener): return (*_uvicorn_listener_flags(listener), "--no-proxy-headers") - case Disabled() | Supervisor(): + case Mode.DISABLED | Mode.SUPERVISOR: return () case _: assert_never(mode) @@ -191,7 +158,7 @@ def gunicorn_bind(mode: RustForwardMode, public: str) -> str: match mode: case Supervised(listener): return _gunicorn_listener(listener) - case Disabled() | Supervisor(): + case Mode.DISABLED | Mode.SUPERVISOR: return public case _: assert_never(mode) @@ -201,7 +168,7 @@ def server_proxy_headers(mode: RustForwardMode) -> bool: match mode: case Supervised(): return False - case Disabled() | Supervisor(): + case Mode.DISABLED | Mode.SUPERVISOR: return True case _: assert_never(mode) diff --git a/litellm/proxy/rust_forward/peer.py b/litellm/proxy/rust_forward/peer.py index cf891495c2b..78445aad4e4 100644 --- a/litellm/proxy/rust_forward/peer.py +++ b/litellm/proxy/rust_forward/peer.py @@ -7,7 +7,7 @@ from typing import ( from typing_extensions import assert_never -from litellm.proxy.rust_forward.mode import RUST_HOP, Disabled, RustForwardMode, Supervised, Supervisor +from litellm.proxy.rust_forward.mode import RUST_HOP, Mode, RustForwardMode, Supervised if TYPE_CHECKING: import starlette.types @@ -63,7 +63,7 @@ def wrap(app: "starlette.types.ASGIApp", mode: RustForwardMode) -> ASGIApp: match mode: case Supervised(): return RustPeerMiddleware(converted, mode) - case Disabled() | Supervisor(): + case Mode.DISABLED | Mode.SUPERVISOR: return converted case _: assert_never(mode) diff --git a/litellm/proxy/rust_forward/server.py b/litellm/proxy/rust_forward/server.py index 8398cbda533..58d0ca1c2ec 100644 --- a/litellm/proxy/rust_forward/server.py +++ b/litellm/proxy/rust_forward/server.py @@ -1,5 +1,10 @@ from dataclasses import dataclass from enum import Enum +from typing import Final + +from pydantic import BaseModel, ConfigDict, TypeAdapter + +_TRUSTED_PROXIES: Final = TypeAdapter(str) class Server(Enum): @@ -27,4 +32,39 @@ class PythonServer: keepalive_seconds: int graceful_timeout_seconds: float proxy_headers: bool - trusted_proxies: str | None + trusted_proxies: str + + +class UvicornCliParams(BaseModel): + """The options uvicorn's own CLI parsed that decide how it serves the public port.""" + + model_config = ConfigDict(frozen=True, extra="ignore") + + host: str + port: int + uds: str | None + fd: int | None + ssl_keyfile: str | None + ssl_certfile: str | None + timeout_keep_alive: int + timeout_graceful_shutdown: int | None + proxy_headers: bool + forwarded_allow_ips: str | None + + @property + def listener(self) -> Listener: + if self.ssl_keyfile is not None or self.ssl_certfile is not None: + return Listener.TLS + if self.uds is not None: + return Listener.UNIX_SOCKET + if self.fd is not None: + return Listener.INHERITED_FD + return Listener.TCP + + +def uvicorn_trusted_proxies(configured: str | None) -> str: + import uvicorn + + return _TRUSTED_PROXIES.validate_python( + uvicorn.Config(app="", log_config=None, forwarded_allow_ips=configured).forwarded_allow_ips + ) diff --git a/litellm/proxy/rust_forward/supervisor.py b/litellm/proxy/rust_forward/supervisor.py index f02d6d4e1cd..d29db08dba0 100644 --- a/litellm/proxy/rust_forward/supervisor.py +++ b/litellm/proxy/rust_forward/supervisor.py @@ -7,15 +7,15 @@ import sys import sysconfig from collections.abc import AsyncGenerator, Callable, Coroutine, Mapping, Sequence from dataclasses import dataclass, replace -from enum import Enum +from enum import Enum, auto from functools import partial from types import MappingProxyType -from typing import Final, TypeAlias +from typing import Final, Literal, TypeAlias from typing_extensions import assert_never from litellm.constants import GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS -from litellm.proxy.rust_forward.mode import RUST_HOP, supervised_environment +from litellm.proxy.rust_forward.mode import RUST_HOP, Supervised, supervised_environment from litellm.proxy.rust_forward.platform import Platform, PrivateListener, SpawnOptions, current_platform from litellm.proxy.rust_forward.server import Listener, PythonServer, Server from litellm.rust_bridge.configuration import ENV_NAME @@ -52,19 +52,10 @@ class Child(Enum): FAILURE: Final = 1 -@dataclass(frozen=True, slots=True) -class NotStarted: - pass - - -@dataclass(frozen=True, slots=True) -class Spawned: - pass - - -@dataclass(frozen=True, slots=True) -class Listening: - pass +class Phase(Enum): + NOT_STARTED = auto() + SPAWNED = auto() + LISTENING = auto() @dataclass(frozen=True, slots=True) @@ -79,8 +70,8 @@ class Lost: status: int -PythonState: TypeAlias = Spawned | Listening | Exited -RustState: TypeAlias = NotStarted | Spawned | Exited | Lost +PythonState: TypeAlias = Literal[Phase.SPAWNED, Phase.LISTENING] | Exited +RustState: TypeAlias = Literal[Phase.NOT_STARTED, Phase.SPAWNED] | Exited | Lost @dataclass(frozen=True, slots=True) @@ -95,9 +86,8 @@ class Caught: signum: signal.Signals -@dataclass(frozen=True, slots=True) -class PythonListening: - pass +class Probe(Enum): + PYTHON_LISTENING = auto() @dataclass(frozen=True, slots=True) @@ -117,39 +107,39 @@ class Send: signum: signal.Signals -Event: TypeAlias = Caught | PythonListening | ChildExited +Event: TypeAlias = Caught | Literal[Probe.PYTHON_LISTENING] | ChildExited Action: TypeAlias = Spawn | Send Transition: TypeAlias = tuple[Lifecycle, tuple[Action, ...]] def boot() -> Transition: - return Lifecycle(Spawned(), NotStarted(), stop_requested=False), (Spawn(Child.PYTHON),) + return Lifecycle(Phase.SPAWNED, Phase.NOT_STARTED, stop_requested=False), (Spawn(Child.PYTHON),) def _relay(state: Lifecycle, signum: signal.Signals, stops: bool) -> Transition: - to_rust: Final = (Send(Child.RUST, signum),) if stops and isinstance(state.rust, Spawned) else () + to_rust: Final = (Send(Child.RUST, signum),) if stops and state.rust is Phase.SPAWNED else () to_python: Final = (Send(Child.PYTHON, signum),) if not isinstance(state.python, Exited) else () return replace(state, stop_requested=state.stop_requested or stops), (*to_rust, *to_python) # Python alone takes no connections until it has booted, so Rust only opens the public port once Python listens. def _python_listening(state: Lifecycle) -> Transition: - if not isinstance(state.python, Spawned): + if state.python is not Phase.SPAWNED: return state, () if state.stop_requested: - return replace(state, python=Listening()), () - return replace(state, python=Listening(), rust=Spawned()), (Spawn(Child.RUST),) + return replace(state, python=Phase.LISTENING), () + return replace(state, python=Phase.LISTENING, rust=Phase.SPAWNED), (Spawn(Child.RUST),) def _python_exited(state: Lifecycle, status: int) -> Transition: if isinstance(state.python, Exited): return state, () exited: Final = replace(state, python=Exited(status)) - return exited, ((Send(Child.RUST, signal.SIGTERM),) if isinstance(state.rust, Spawned) else ()) + return exited, ((Send(Child.RUST, signal.SIGTERM),) if state.rust is Phase.SPAWNED else ()) def _rust_exited(state: Lifecycle, status: int) -> Transition: - if not isinstance(state.rust, Spawned): + if state.rust is not Phase.SPAWNED: return state, () if state.stop_requested or isinstance(state.python, Exited): return replace(state, rust=Exited(status)), () @@ -162,7 +152,7 @@ def step(state: Lifecycle, event: Event, stop_signals: frozenset[signal.Signals] match event: case Caught(signum): return _relay(state, signum, signum in stop_signals) - case PythonListening(): + case Probe.PYTHON_LISTENING: return _python_listening(state) case ChildExited(child, status): match child: @@ -178,7 +168,7 @@ def step(state: Lifecycle, event: Event, stop_signals: frozenset[signal.Signals] def exit_status(state: Lifecycle) -> int | None: match state.python, state.rust: - case Exited(status), NotStarted() | Exited(): + case Exited(status), Phase.NOT_STARTED | Exited(): return status case Exited(), Lost(): return FAILURE @@ -198,7 +188,9 @@ def python_spec( ) -> ChildSpec: return ChildSpec( argv=(*command, "--port", str(private.port)), - env_overrides=supervised_environment(private.listener, server.proxy_headers, server.trusted_proxies), + env_overrides=supervised_environment( + Supervised(private.listener, server.proxy_headers, server.trusted_proxies) + ), spawn=platform.spawn_options(private.pass_fds), ) @@ -251,7 +243,7 @@ async def _accepts(port: int) -> bool: async def _report_once_listening(port: int, events: asyncio.Queue[Event]) -> None: while not await _accepts(port): await asyncio.sleep(PROBE_INTERVAL_SECONDS) - events.put_nowait(PythonListening()) + events.put_nowait(Probe.PYTHON_LISTENING) @contextlib.asynccontextmanager @@ -295,53 +287,22 @@ async def _launch_python( yield process -@dataclass(frozen=True, slots=True) -class _Launching: - state: Lifecycle - child: Child - rest: tuple[Action, ...] - - -async def _advance( - platform: Platform, - transition: Transition, - events: asyncio.Queue[Event], - children: Mapping[Child, asyncio.subprocess.Process], -) -> int | _Launching: - state, actions = transition - while True: - for index, action in enumerate(actions): - match action: - case Send(child, signum): - platform.signal_child(children[child], signum) - case Spawn(child): - return _Launching(state, child, actions[index + 1 :]) - case _: - assert_never(action) - status = exit_status(state) - if status is not None: - return status - state, actions = step(state, await events.get(), platform.stop_signals) - - -async def _run( - platform: Platform, - transition: Transition, - events: asyncio.Queue[Event], - children: Mapping[Child, asyncio.subprocess.Process], - launches: Mapping[Child, Launch], -) -> int: - advanced: Final = await _advance(platform, transition, events, children) - match advanced: - case int(): - return advanced - case _Launching(state, child, rest): - async with launches[child](events) as process: - return await _run( - platform, (state, rest), events, MappingProxyType({**children, child: process}), launches - ) - case _: - assert_never(advanced) +async def _run(platform: Platform, events: asyncio.Queue[Event], launches: Mapping[Child, Launch]) -> int: + async with contextlib.AsyncExitStack() as stack: + children: Final[dict[Child, asyncio.subprocess.Process]] = {} # mutable-ok: indexes the stack's children + state, actions = boot() + while True: + for action in actions: + match action: + case Spawn(child): + children[child] = await stack.enter_async_context(launches[child](events)) + case Send(child, signum): + platform.signal_child(children[child], signum) + case _: + assert_never(action) + if (status := exit_status(state)) is not None: + return status + state, actions = step(state, await events.get(), platform.stop_signals) async def _supervise(platform: Platform, binary: str, command: Sequence[str], server: PythonServer) -> int: @@ -356,7 +317,7 @@ async def _supervise(platform: Platform, binary: str, command: Sequence[str], se Child.RUST: partial(_launch, platform, Child.RUST, rust_spec(platform, binary, private.port, server)), } ) - return await _run(platform, boot(), events, MappingProxyType({}), launches) + return await _run(platform, events, launches) def _exit_like(status: int) -> int: diff --git a/tests/integration/configuration/test_rust_forward_lifecycle.py b/tests/integration/configuration/test_rust_forward_lifecycle.py index 3c639ce6d9a..7abafc6c6cd 100644 --- a/tests/integration/configuration/test_rust_forward_lifecycle.py +++ b/tests/integration/configuration/test_rust_forward_lifecycle.py @@ -46,14 +46,14 @@ import uvicorn def listener(): - fd = os.environ.get("_LITELLM_GATEWAY_FD") - if fd is None: - return {"host": "127.0.0.1", "port": int(sys.argv[2])} - from litellm.proxy.rust_forward.mode import rust_forward_mode, stop_with_supervisor + from litellm.proxy.rust_forward.mode import Supervised, rust_forward_mode, stop_with_supervisor - stop_with_supervisor(rust_forward_mode(os.environ)) + mode = rust_forward_mode(os.environ) + if not isinstance(mode, Supervised): + return {"host": "127.0.0.1", "port": int(sys.argv[2])} + stop_with_supervisor(mode) pathlib.Path(os.environ["MARKERS"], "python.pid").write_text(str(os.getpid())) - return {"fd": int(fd)} + return {"fd": mode.listener.fd} if __name__ == "__main__": @@ -123,7 +123,7 @@ from litellm.proxy.rust_forward.server import Listener, PythonServer, Server from litellm.proxy.rust_forward.supervisor import supervise workers, port, scripts = sys.argv[1], int(sys.argv[2]), Path(sys.argv[3]) -server = PythonServer(Server.UVICORN, Listener.TCP, "127.0.0.1", port, 5, 30.0, True, None) +server = PythonServer(Server.UVICORN, Listener.TCP, "127.0.0.1", port, 5, 30.0, True, "127.0.0.1") result = supervise(server, (sys.executable, str(scripts / "serve.py"), workers, str(port)), scripts_dir=str(scripts)) raise SystemExit(result if isinstance(result, int) else 3) """ diff --git a/tests/unit/gateway/test_launch.py b/tests/unit/gateway/test_launch.py index 8a916560c1b..becafe07d5d 100644 --- a/tests/unit/gateway/test_launch.py +++ b/tests/unit/gateway/test_launch.py @@ -126,7 +126,8 @@ class TestUvicornArgv: (("--fd", "0"), Listener.INHERITED_FD), (("--fd=17",), Listener.INHERITED_FD), (("--ssl-keyfile", "unused"), Listener.TLS), - (("--ssl-ciphers", "TLSv1.2"), Listener.TLS), + (("--ssl-certfile", "unused"), Listener.TLS), + (("--ssl-ciphers", "TLSv1.2"), Listener.TCP), ], ) def test_forwarding_sees_the_public_listener_uvicorn_would_open( @@ -135,18 +136,25 @@ class TestUvicornArgv: assert python_server(arguments).listener is listener @pytest.mark.parametrize( - ("argv", "trust"), + ("argv", "environ", "trust"), [ - ((), (True, None)), - (("--forwarded-allow-ips", "10.0.0.5"), (True, "10.0.0.5")), - (("--forwarded-allow-ips=*",), (True, "*")), - (("--no-proxy-headers",), (False, None)), + ((), {}, (True, "127.0.0.1")), + ((), {"FORWARDED_ALLOW_IPS": "10.9.9.9"}, (True, "10.9.9.9")), + (("--forwarded-allow-ips", "10.0.0.5"), {"FORWARDED_ALLOW_IPS": "10.9.9.9"}, (True, "10.0.0.5")), + (("--forwarded-allow-ips=*",), {}, (True, "*")), + (("--no-proxy-headers",), {}, (False, "127.0.0.1")), ], ) - def test_forwarding_carries_operator_proxy_header_options_to_the_child( - self, argv: tuple[str, ...], trust: tuple[bool, str | None], monkeypatch: pytest.MonkeyPatch + def test_forwarding_carries_the_trust_uvicorn_resolves_to_the_child( + self, + argv: tuple[str, ...], + environ: dict[str, str], + trust: tuple[bool, str], + monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.delenv("FORWARDED_ALLOW_IPS", raising=False) + for name, value in environ.items(): + monkeypatch.setenv(name, value) server: Final = python_server(("--host", "0.0.0.0", "--port", "4000", *argv)) assert (server.host, server.port) == ("0.0.0.0", 4000) assert (server.proxy_headers, server.trusted_proxies) == trust @@ -335,7 +343,7 @@ class TestMain: capsys: pytest.CaptureFixture[str], ) -> None: monkeypatch.setenv("LITELLM_RUST", rust) - monkeypatch.delenv("_LITELLM_GATEWAY_FD", raising=False) + monkeypatch.delenv("_LITELLM_GATEWAY_SUPERVISED", raising=False) monkeypatch.delenv("LITELLM_PGBOUNCER_ENABLED", raising=False) served: Final[list[tuple[str, ...]]] = [] main(("--uds", "/tmp/gateway.sock"), serve=lambda argv: served.append(tuple(argv))) diff --git a/tests/unit/proxy/rust_forward/test_mode.py b/tests/unit/proxy/rust_forward/test_mode.py index 34329bacbb2..5acd76860fb 100644 --- a/tests/unit/proxy/rust_forward/test_mode.py +++ b/tests/unit/proxy/rust_forward/test_mode.py @@ -8,12 +8,12 @@ from pydantic import ValidationError from litellm.proxy.proxy_cli import WORKER_APP, ProxyInitializationHelpers from litellm.proxy.rust_forward.mode import ( + SUPERVISED_ENV, AssignedPort, - Disabled, InheritedFd, + Mode, RustForwardMode, Supervised, - Supervisor, gunicorn_bind, rust_forward_mode, server_proxy_headers, @@ -26,21 +26,17 @@ from litellm.proxy.rust_forward.mode import ( @pytest.mark.parametrize( ("environ", "expected"), [ - ({}, Disabled()), - ({"LITELLM_RUST": "false"}, Disabled()), - ({"LITELLM_RUST": "off"}, Disabled()), - ({"LITELLM_RUST": "0"}, Disabled()), - ({"LITELLM_RUST": "true"}, Supervisor()), - ({"LITELLM_RUST": "1"}, Supervisor()), - ({"LITELLM_RUST": "yes"}, Supervisor()), - ({"LITELLM_RUST": "on"}, Supervisor()), - ({"LITELLM_RUST": " 1 "}, Supervisor()), - ({"LITELLM_RUST": ""}, Disabled()), - ({"LITELLM_RUST": " "}, Disabled()), - ( - {"LITELLM_RUST": "true", "_LITELLM_GATEWAY_FD": "17", "FORWARDED_ALLOW_IPS": "127.0.0.1,10.0.0.5"}, - Supervised(InheritedFd(17), True, "127.0.0.1,10.0.0.5"), - ), + ({}, Mode.DISABLED), + ({"LITELLM_RUST": "false"}, Mode.DISABLED), + ({"LITELLM_RUST": "off"}, Mode.DISABLED), + ({"LITELLM_RUST": "0"}, Mode.DISABLED), + ({"LITELLM_RUST": "true"}, Mode.SUPERVISOR), + ({"LITELLM_RUST": "1"}, Mode.SUPERVISOR), + ({"LITELLM_RUST": "yes"}, Mode.SUPERVISOR), + ({"LITELLM_RUST": "on"}, Mode.SUPERVISOR), + ({"LITELLM_RUST": " 1 "}, Mode.SUPERVISOR), + ({"LITELLM_RUST": ""}, Mode.DISABLED), + ({"LITELLM_RUST": " "}, Mode.DISABLED), ], ) def test_mode_is_resolved_from_environment(environ: dict[str, str], expected: RustForwardMode) -> None: @@ -48,35 +44,36 @@ def test_mode_is_resolved_from_environment(environ: dict[str, str], expected: Ru @pytest.mark.parametrize( - ("environ", "variable"), + ("environ", "invalid"), [ ({"LITELLM_RUST": "enabled"}, "LITELLM_RUST"), ({"LITELLM_RUST": "2"}, "LITELLM_RUST"), - ({"_LITELLM_GATEWAY_FD": "seventeen"}, "_LITELLM_GATEWAY_FD"), - ({"_LITELLM_GATEWAY_FD": "17", "_LITELLM_GATEWAY_PROXY_HEADERS": "maybe"}, "_LITELLM_GATEWAY_PROXY_HEADERS"), + ({SUPERVISED_ENV: '{"listener": {"fd": "seventeen"}, "proxy_headers": true, "trusted_proxies": "*"}'}, "fd"), + ( + {SUPERVISED_ENV: '{"listener": {"fd": 17}, "proxy_headers": "maybe", "trusted_proxies": "*"}'}, + "proxy_headers", + ), + ({SUPERVISED_ENV: '{"listener": {"fd": 17}, "proxy_headers": true}'}, "trusted_proxies"), ], ) -def test_invalid_forwarding_value_stops_startup_naming_the_variable(environ: dict[str, str], variable: str) -> None: - with pytest.raises(ValidationError, match=variable): +def test_invalid_forwarding_value_stops_startup_naming_it(environ: dict[str, str], invalid: str) -> None: + with pytest.raises(ValidationError, match=invalid): rust_forward_mode(environ) @pytest.mark.parametrize( - ("proxy_headers", "trusted_proxies", "inherited", "expected"), + "written", [ - (True, None, {}, Supervised(InheritedFd(17), True, "127.0.0.1")), - (True, None, {"FORWARDED_ALLOW_IPS": "10.0.0.5"}, Supervised(InheritedFd(17), True, "10.0.0.5")), - (True, "*", {"FORWARDED_ALLOW_IPS": "10.0.0.5"}, Supervised(InheritedFd(17), True, "*")), - (False, None, {}, Supervised(InheritedFd(17), False, "127.0.0.1")), - (True, None, {}, Supervised(AssignedPort(45555), True, "127.0.0.1")), - (False, "*", {}, Supervised(AssignedPort(45555), False, "*")), + Supervised(InheritedFd(17), True, "127.0.0.1"), + Supervised(InheritedFd(17), True, "127.0.0.1,10.0.0.5"), + Supervised(InheritedFd(17), False, "127.0.0.1"), + Supervised(AssignedPort(45555), True, "*"), + Supervised(AssignedPort(45555), False, "*"), ], ) -def test_child_reads_back_the_listener_and_trust_the_supervisor_wrote( - proxy_headers: bool, trusted_proxies: str | None, inherited: dict[str, str], expected: Supervised -) -> None: - written: Final = supervised_environment(expected.listener, proxy_headers, trusted_proxies) - assert rust_forward_mode({**inherited, **written}) == expected +def test_child_reads_back_exactly_what_the_supervisor_wrote(written: Supervised) -> None: + inherited: Final = {"LITELLM_RUST": "true", "FORWARDED_ALLOW_IPS": "10.9.9.9"} + assert rust_forward_mode({**inherited, **supervised_environment(written)}) == written def _served_uvicorn_config(mode: RustForwardMode) -> uvicorn.Config: @@ -93,7 +90,7 @@ def test_supervised_uvicorn_serves_the_inherited_socket_without_trusting_the_rus assert (config.app, config.factory) == (WORKER_APP, True) -@pytest.mark.parametrize("mode", [Disabled(), Supervisor()]) +@pytest.mark.parametrize("mode", list(Mode)) def test_unsupervised_servers_keep_their_own_socket_and_proxy_headers(mode: RustForwardMode) -> None: config: Final = _served_uvicorn_config(mode) assert (config.fd, config.proxy_headers, config.port) == (None, True, 45555) @@ -117,7 +114,7 @@ def test_assigned_port_serves_only_the_loopback_hop_without_trusting_it() -> Non @pytest.mark.parametrize( ("mode", "stops"), - [(Supervised(InheritedFd(17), False, "127.0.0.1"), True), (Disabled(), False), (Supervisor(), False)], + [(Supervised(InheritedFd(17), False, "127.0.0.1"), True), (Mode.DISABLED, False), (Mode.SUPERVISOR, False)], ) def test_only_a_supervised_server_stops_once_its_supervisor_is_gone(mode: RustForwardMode, stops: bool) -> None: read_end, write_end = os.pipe() diff --git a/tests/unit/proxy/rust_forward/test_peer.py b/tests/unit/proxy/rust_forward/test_peer.py index b0c59641e7f..01331e2c69e 100644 --- a/tests/unit/proxy/rust_forward/test_peer.py +++ b/tests/unit/proxy/rust_forward/test_peer.py @@ -20,7 +20,7 @@ if TYPE_CHECKING: ) from litellm.proxy.proxy_server import app as proxy_app -from litellm.proxy.rust_forward.mode import InheritedFd, Supervised +from litellm.proxy.rust_forward.mode import SUPERVISED_ENV, InheritedFd, Supervised, supervised_environment from litellm.proxy.rust_forward.peer import RustPeerMiddleware Headers = tuple[tuple[bytes, bytes], ...] @@ -124,7 +124,9 @@ async def _peer_probe(request: Request) -> dict[str, object]: async def test_every_entry_point_hands_routes_the_peer_rust_saw( entry_point: str, monkeypatch: pytest.MonkeyPatch ) -> None: - monkeypatch.setenv("_LITELLM_GATEWAY_FD", "17") + monkeypatch.setenv( + SUPERVISED_ENV, supervised_environment(Supervised(InheritedFd(17), True, "127.0.0.1"))[SUPERVISED_ENV] + ) monkeypatch.setattr(proxy_app.router, "lifespan_context", proxy_app.router.lifespan_context) monkeypatch.setattr( proxy_app.router, "routes", [APIRoute("/__rust_forward_peer_probe", _peer_probe), *proxy_app.router.routes] diff --git a/tests/unit/proxy/rust_forward/test_supervisor.py b/tests/unit/proxy/rust_forward/test_supervisor.py index 1931b11c7a9..bb14443c96e 100644 --- a/tests/unit/proxy/rust_forward/test_supervisor.py +++ b/tests/unit/proxy/rust_forward/test_supervisor.py @@ -31,10 +31,10 @@ from litellm.proxy.rust_forward.supervisor import ( ChildExited, Event, Lifecycle, - PythonListening, + Phase, + Probe, Send, Spawn, - Spawned, StartupError, boot, exit_status, @@ -54,7 +54,7 @@ POSIX_ONLY: Final = pytest.mark.skipif(sys.platform == "win32", reason="runs a P def _server( - server: Server = Server.UVICORN, listener: Listener = Listener.TCP, trusted_proxies: str | None = None + server: Server = Server.UVICORN, listener: Listener = Listener.TCP, trusted_proxies: str = "127.0.0.1" ) -> PythonServer: return PythonServer(server, listener, "0.0.0.0", 4000, 2, 2.5, False, trusted_proxies) @@ -168,12 +168,11 @@ def test_binary_is_found_only_in_the_environment_scripts_dir(tmp_path: Path, mon PYTHON_CHILD: Final = """ -import os, pathlib, socket, time +import json, os, pathlib, socket, time time.sleep(0.5) pathlib.Path(os.environ["MARKERS"], "listening").touch() -fd = os.environ.get("_LITELLM_GATEWAY_FD") -port = os.environ.get("_LITELLM_GATEWAY_PORT") -listener = socket.socket(fileno=int(fd)) if fd else socket.create_server(("127.0.0.1", int(port))) +private = json.loads(os.environ["_LITELLM_GATEWAY_SUPERVISED"])["listener"] +listener = socket.socket(fileno=private["fd"]) if "fd" in private else socket.create_server(("127.0.0.1", private["port"])) listener.listen() time.sleep(10) """ @@ -240,7 +239,7 @@ def _after(world: World, event: Event) -> World: match event: case Caught(signum): return replace(world, stop_requested=world.stop_requested or signum in STOPS) - case PythonListening(): + case Probe.PYTHON_LISTENING: return replace(world, reported_listening=world.reported_listening or Child.PYTHON in world.alive) case ChildExited(child, status) if child in world.alive: lost: Final = child is Child.RUST and Child.PYTHON in world.alive and not world.stop_requested @@ -264,7 +263,7 @@ STOPS: Final = POSIX.stop_signals EVENTS: Final = st.lists( st.one_of( st.sampled_from(SIGNALS).map(Caught), - st.just(PythonListening()), + st.just(Probe.PYTHON_LISTENING), st.builds(ChildExited, st.sampled_from(Child), st.integers(-15, 3)), ), max_size=40, @@ -272,8 +271,8 @@ EVENTS: Final = st.lists( @given(EVENTS) -@example([Caught(TERM), Caught(signal.SIGHUP), PythonListening()]) -@example([PythonListening(), Caught(TERM), Caught(signal.SIGWINCH), ChildExited(Child.RUST, 0)]) +@example([Caught(TERM), Caught(signal.SIGHUP), Probe.PYTHON_LISTENING]) +@example([Probe.PYTHON_LISTENING, Caught(TERM), Caught(signal.SIGWINCH), ChildExited(Child.RUST, 0)]) def test_lifecycle_relays_every_signal_and_ends_with_pythons_status(events: list[Event]) -> None: state, actions = boot() world = World() @@ -284,7 +283,7 @@ def test_lifecycle_relays_every_signal_and_ends_with_pythons_status(events: list match event: case Caught(signum): assert actions == _expected_relay(before, signum) - case PythonListening(): + case Probe.PYTHON_LISTENING: starts_rust = Child.PYTHON in world.alive and not world.reported_listening and not world.stop_requested assert actions == ((Spawn(Child.RUST),) if starts_rust else ()) case ChildExited(child) if child in world.alive: @@ -308,7 +307,7 @@ def test_rust_starts_at_most_once_and_only_while_python_listens_unstopped(events spawns: Final = tuple((event, state) for event, state, action in history if isinstance(action, Spawn)) assert len(spawns) <= 1 assert all( - isinstance(event, PythonListening) and isinstance(state.python, Spawned) and not state.stop_requested + event is Probe.PYTHON_LISTENING and state.python is Phase.SPAWNED and not state.stop_requested for event, state in spawns ) @@ -316,10 +315,22 @@ def test_rust_starts_at_most_once_and_only_while_python_listens_unstopped(events @pytest.mark.parametrize( ("events", "expected"), [ - ((PythonListening(), Caught(TERM), ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, -15)), -15), - ((PythonListening(), Caught(INT), Caught(INT), ChildExited(Child.PYTHON, 1), ChildExited(Child.RUST, 0)), 1), - ((PythonListening(), Caught(signal.SIGHUP), ChildExited(Child.PYTHON, -1), ChildExited(Child.RUST, 0)), -1), - ((PythonListening(), ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, 0)), FAILURE), + ((Probe.PYTHON_LISTENING, Caught(TERM), ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, -15)), -15), + ( + ( + Probe.PYTHON_LISTENING, + Caught(INT), + Caught(INT), + ChildExited(Child.PYTHON, 1), + ChildExited(Child.RUST, 0), + ), + 1, + ), + ( + (Probe.PYTHON_LISTENING, Caught(signal.SIGHUP), ChildExited(Child.PYTHON, -1), ChildExited(Child.RUST, 0)), + -1, + ), + ((Probe.PYTHON_LISTENING, ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, 0)), FAILURE), ((ChildExited(Child.PYTHON, 3),), 3), ], ) diff --git a/tests/unit/proxy/test_proxy_cli.py b/tests/unit/proxy/test_proxy_cli.py index 7a30d78fee4..5ad21572605 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -4297,15 +4297,16 @@ class TestRustForwardServer: run_gunicorn=params["run_gunicorn"] is True, run_hypercorn=params["run_hypercorn"] is True, run_granian=params["run_granian"] is True, - tls=any((params["ssl_keyfile_path"], params["ssl_certfile_path"], params["ciphers"])), + tls=None not in (params["ssl_keyfile_path"], params["ssl_certfile_path"]), ) @pytest.mark.parametrize( ("argv", "reason"), [ - (["--ssl_keyfile_path", "unused"], "does not support TLS on the public port"), - (["--ssl_certfile_path", "unused"], "does not support TLS on the public port"), - (["--ciphers", "unused"], "does not support TLS on the public port"), + ( + ["--ssl_keyfile_path", "key.pem", "--ssl_certfile_path", "cert.pem"], + "does not support TLS on the public port", + ), (["--run_hypercorn"], "does not support Hypercorn"), (["--run_granian"], "does not support Granian"), ], @@ -4321,8 +4322,28 @@ class TestRustForwardServer: @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) + @pytest.mark.parametrize( + ("argv", "reason", "served_keyfile"), + [ + ( + ["--ssl_keyfile_path", "key.pem", "--ssl_certfile_path", "cert.pem"], + "does not support TLS on the public port", + "key.pem", + ), + (["--ssl_keyfile_path", "key.pem"], "requires the litellm-gateway binary", None), + (["--ssl_certfile_path", "cert.pem"], "requires the litellm-gateway binary", None), + (["--ciphers", "ECDHE"], "requires the litellm-gateway binary", None), + ], + ) def test_a_refused_forwarding_start_serves_from_python_alone( - self, mock_should_update, mock_setup_db, mock_atexit_register, mock_uvicorn_run + self, + mock_should_update, + mock_setup_db, + mock_atexit_register, + mock_uvicorn_run, + argv: list[str], + reason: str, + served_keyfile: str | None, ) -> None: from click.testing import CliRunner @@ -4333,20 +4354,24 @@ class TestRustForwardServer: save_worker_config=MagicMock(), ) environ: Final = { - **{k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL", "_LITELLM_GATEWAY_FD")}, + **{ + k: v + for k, v in os.environ.items() + if k not in ("DATABASE_URL", "DIRECT_URL", "_LITELLM_GATEWAY_SUPERVISED") + }, "LITELLM_RUST": "true", } with ( patch.dict(os.environ, environ, clear=True), patch.dict("sys.modules", {"proxy_server": proxy_module, "litellm.proxy.proxy_server": proxy_module}), + patch("litellm.proxy.rust_forward.supervisor.gateway_binary", return_value=None), ): - result: Final = CliRunner().invoke( - run_server, ["--port", "45999", "--ssl_keyfile_path", "key.pem", "--ssl_certfile_path", "cert.pem"] - ) + result: Final = CliRunner().invoke(run_server, ["--port", "45999", *argv]) assert result.exit_code == 0, result.output - assert "LITELLM_RUST does not support TLS on the public port, so Python serves the public port" in result.output + assert f"LITELLM_RUST {reason}" in result.output + assert "so Python serves the public port without Rust forwarding" in result.output mock_uvicorn_run.assert_called_once() - assert mock_uvicorn_run.call_args.kwargs["ssl_keyfile"] == "key.pem" + assert mock_uvicorn_run.call_args.kwargs.get("ssl_keyfile") == served_keyfile assert mock_uvicorn_run.call_args.kwargs["port"] == 45999 def test_an_invalid_rust_switch_stops_startup_with_one_line_naming_it(self) -> None: From 5f97ff0e60fbc68f0c59927dc6056ca3b87e63af Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 6 Nov 2026 14:32:44 +0900 Subject: [PATCH 07/10] refactor(proxy): split the Rust forwarding supervisor by job supervisor.py held the preflight checks, the pure lifecycle, both child specs and the asyncio runtime in one file. It is now a package: preflight decides whether forwarding can start, children says what each child runs with, lifecycle decides, runtime acts, and supervise in __init__ is the one entry point the CLIs call. Code moves unchanged and the tests split to mirror the modules --- gateway/launch.py | 3 +- litellm/proxy/proxy_cli.py | 3 +- litellm/proxy/rust_forward/AGENTS.md | 4 +- litellm/proxy/rust_forward/supervisor.py | 360 ------------------ .../proxy/rust_forward/supervisor/__init__.py | 41 ++ .../proxy/rust_forward/supervisor/children.py | 45 +++ .../rust_forward/supervisor/lifecycle.py | 138 +++++++ .../rust_forward/supervisor/preflight.py | 27 ++ .../proxy/rust_forward/supervisor/runtime.py | 142 +++++++ .../proxy/rust_forward/platform/__init__.py | 0 .../rust_forward/platform/test_windows.py | 13 + .../proxy/rust_forward/supervisor/__init__.py | 0 .../proxy/rust_forward/supervisor/_support.py | 7 + .../rust_forward/supervisor/test_children.py | 53 +++ .../rust_forward/supervisor/test_lifecycle.py | 167 ++++++++ .../rust_forward/supervisor/test_preflight.py | 86 +++++ .../rust_forward/supervisor/test_runtime.py | 45 +++ .../proxy/rust_forward/test_supervisor.py | 339 ----------------- tests/unit/proxy/test_proxy_cli.py | 4 +- 19 files changed, 772 insertions(+), 705 deletions(-) delete mode 100644 litellm/proxy/rust_forward/supervisor.py create mode 100644 litellm/proxy/rust_forward/supervisor/__init__.py create mode 100644 litellm/proxy/rust_forward/supervisor/children.py create mode 100644 litellm/proxy/rust_forward/supervisor/lifecycle.py create mode 100644 litellm/proxy/rust_forward/supervisor/preflight.py create mode 100644 litellm/proxy/rust_forward/supervisor/runtime.py create mode 100644 tests/unit/proxy/rust_forward/platform/__init__.py create mode 100644 tests/unit/proxy/rust_forward/platform/test_windows.py create mode 100644 tests/unit/proxy/rust_forward/supervisor/__init__.py create mode 100644 tests/unit/proxy/rust_forward/supervisor/_support.py create mode 100644 tests/unit/proxy/rust_forward/supervisor/test_children.py create mode 100644 tests/unit/proxy/rust_forward/supervisor/test_lifecycle.py create mode 100644 tests/unit/proxy/rust_forward/supervisor/test_preflight.py create mode 100644 tests/unit/proxy/rust_forward/supervisor/test_runtime.py delete mode 100644 tests/unit/proxy/rust_forward/test_supervisor.py diff --git a/gateway/launch.py b/gateway/launch.py index 1adb2ef1034..e6b0c9dbeef 100644 --- a/gateway/launch.py +++ b/gateway/launch.py @@ -121,7 +121,8 @@ def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) - stop_with_supervisor, uvicorn_flags, ) - from litellm.proxy.rust_forward.supervisor import StartupError, supervise + from litellm.proxy.rust_forward.supervisor import supervise + from litellm.proxy.rust_forward.supervisor.preflight import StartupError mode: Final = rust_forward_mode(os.environ) match mode: diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 632a71ec5c8..afee156981b 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -1376,7 +1376,8 @@ def run_server( if forward_mode is Mode.SUPERVISOR and not any( (version, validate_config, health, test, skip_server_startup, check_migrations_against_existing_data) ): - from litellm.proxy.rust_forward.supervisor import StartupError, supervise + from litellm.proxy.rust_forward.supervisor import supervise + from litellm.proxy.rust_forward.supervisor.preflight import StartupError server: Final = ProxyInitializationHelpers.python_server( host, diff --git a/litellm/proxy/rust_forward/AGENTS.md b/litellm/proxy/rust_forward/AGENTS.md index 907b7dc1a3c..2084a19a66d 100644 --- a/litellm/proxy/rust_forward/AGENTS.md +++ b/litellm/proxy/rust_forward/AGENTS.md @@ -11,9 +11,9 @@ Every OS-specific operation lives behind `platform.Platform`, which `supervise` Windows support exists only for running the proxy locally from the CLI. Nobody deploys the proxy on Windows, and the Windows design leans on that. A port picked by the supervisor can be taken by another process before Python binds it, `CTRL_BREAK_EVENT` needs the console a CLI run always has, and gunicorn is refused. Revisit those choices before treating Windows as a deployment target -Each environment contract is written and read in one place. The Python child's is the `mode.Supervised` type itself: `mode.supervised_environment` serializes it into one variable and `mode.rust_forward_mode` validates it back, so a new field needs no parsing code. The trust it carries is resolved by the server library the entry point uses (`server.uvicorn_trusted_proxies` or gunicorn's config), never defaulted in this package. `supervisor.rust_spec` writes the `LITELLM_FORWARD_*` variables that `gateway-forward`'s `settings.rs` reads. Python owns every value Rust needs, so a new one gets a new variable in `rust_spec` +Each environment contract is written and read in one place. The Python child's is the `mode.Supervised` type itself: `mode.supervised_environment` serializes it into one variable and `mode.rust_forward_mode` validates it back, so a new field needs no parsing code. The trust it carries is resolved by the server library the entry point uses (`server.uvicorn_trusted_proxies` or gunicorn's config), never defaulted in this package. `supervisor.children.rust_spec` writes the `LITELLM_FORWARD_*` variables that `gateway-forward`'s `settings.rs` reads. Python owns every value Rust needs, so a new one gets a new variable in `children.rust_spec` -The supervisor is transparent to Python, so `python` with `LITELLM_RUST` behaves on signals and exits exactly as Python alone. Its whole lifecycle is `supervisor.step`, a pure function from a `Lifecycle` of facts and an event to the next facts and the actions to run, and `exit_status` reads the result from those facts. `_run` only executes actions, and the `contextlib.AsyncExitStack` it enters each child into owns their cleanup, last spawned first. Three rules decide everything there, and a change that breaks one of them breaks the equivalence `tests/integration/configuration/test_rust_forward_lifecycle.py` checks against Python alone +The supervisor is transparent to Python, so `python` with `LITELLM_RUST` behaves on signals and exits exactly as Python alone. The `supervisor` package splits by job: `preflight` decides whether forwarding can start, `children` says what each child runs with, `lifecycle` decides, and `runtime` acts. The whole lifecycle is `lifecycle.step`, a pure function from a `Lifecycle` of facts and an event to the next facts and the actions to run, and `exit_status` reads the result from those facts. `runtime` only executes actions, and the `contextlib.AsyncExitStack` it enters each child into owns their cleanup, last spawned first. Three rules decide everything there, and a change that breaks one of them breaks the equivalence `tests/integration/configuration/test_rust_forward_lifecycle.py` checks against Python alone - Signals pass through. Every signal in `Platform.relayed_signals` goes to Python unchanged, and the ones in `Platform.stop_signals` go to Rust first. The supervisor never reads meaning into a signal, so a repeated SIGINT reaches Python again and forces its exit as it would alone. Children run in their own session (`start_new_session`) only so that a terminal Ctrl-C reaches Python once, through the supervisor - Python's exit status is the result. When Python exits for any reason, Rust gets SIGTERM, and once Rust is gone the supervisor exits with Python's status, dying by the same signal when Python died by one. The one exception is Rust exiting while Python serves with no stop requested, which stops Python and exits 1 diff --git a/litellm/proxy/rust_forward/supervisor.py b/litellm/proxy/rust_forward/supervisor.py deleted file mode 100644 index d29db08dba0..00000000000 --- a/litellm/proxy/rust_forward/supervisor.py +++ /dev/null @@ -1,360 +0,0 @@ -import asyncio -import contextlib -import os -import shutil -import signal -import sys -import sysconfig -from collections.abc import AsyncGenerator, Callable, Coroutine, Mapping, Sequence -from dataclasses import dataclass, replace -from enum import Enum, auto -from functools import partial -from types import MappingProxyType -from typing import Final, Literal, TypeAlias - -from typing_extensions import assert_never - -from litellm.constants import GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS -from litellm.proxy.rust_forward.mode import RUST_HOP, Supervised, supervised_environment -from litellm.proxy.rust_forward.platform import Platform, PrivateListener, SpawnOptions, current_platform -from litellm.proxy.rust_forward.server import Listener, PythonServer, Server -from litellm.rust_bridge.configuration import ENV_NAME - -PROBE_INTERVAL_SECONDS: Final = 0.05 - - -def gateway_binary(platform: Platform, scripts_dir: str) -> str | None: - return shutil.which(platform.gateway_name, path=scripts_dir) - - -@dataclass(frozen=True, slots=True) -class StartupError: - reason: str - - -def startup_check(platform: Platform, server: PythonServer, binary: str | None) -> str | StartupError: - if server.server in (Server.HYPERCORN, Server.GRANIAN): - return StartupError(f"{ENV_NAME} does not support {server.server.value}") - if server.server not in platform.servers: - return StartupError(f"{ENV_NAME} does not support {server.server.value} on {platform.name}") - if server.listener is not Listener.TCP: - return StartupError(f"{ENV_NAME} does not support {server.listener.value} on the public port") - if binary is None: - return StartupError(f"{ENV_NAME} requires the litellm-gateway binary from the litellm-proxy package") - return binary - - -class Child(Enum): - RUST = "rust" - PYTHON = "python" - - -FAILURE: Final = 1 - - -class Phase(Enum): - NOT_STARTED = auto() - SPAWNED = auto() - LISTENING = auto() - - -@dataclass(frozen=True, slots=True) -class Exited: - status: int - - -@dataclass(frozen=True, slots=True) -class Lost: - """Rust exited while Python was serving and nobody had asked either to stop.""" - - status: int - - -PythonState: TypeAlias = Literal[Phase.SPAWNED, Phase.LISTENING] | Exited -RustState: TypeAlias = Literal[Phase.NOT_STARTED, Phase.SPAWNED] | Exited | Lost - - -@dataclass(frozen=True, slots=True) -class Lifecycle: - python: PythonState - rust: RustState - stop_requested: bool - - -@dataclass(frozen=True, slots=True) -class Caught: - signum: signal.Signals - - -class Probe(Enum): - PYTHON_LISTENING = auto() - - -@dataclass(frozen=True, slots=True) -class ChildExited: - child: Child - status: int - - -@dataclass(frozen=True, slots=True) -class Spawn: - child: Child - - -@dataclass(frozen=True, slots=True) -class Send: - child: Child - signum: signal.Signals - - -Event: TypeAlias = Caught | Literal[Probe.PYTHON_LISTENING] | ChildExited -Action: TypeAlias = Spawn | Send -Transition: TypeAlias = tuple[Lifecycle, tuple[Action, ...]] - - -def boot() -> Transition: - return Lifecycle(Phase.SPAWNED, Phase.NOT_STARTED, stop_requested=False), (Spawn(Child.PYTHON),) - - -def _relay(state: Lifecycle, signum: signal.Signals, stops: bool) -> Transition: - to_rust: Final = (Send(Child.RUST, signum),) if stops and state.rust is Phase.SPAWNED else () - to_python: Final = (Send(Child.PYTHON, signum),) if not isinstance(state.python, Exited) else () - return replace(state, stop_requested=state.stop_requested or stops), (*to_rust, *to_python) - - -# Python alone takes no connections until it has booted, so Rust only opens the public port once Python listens. -def _python_listening(state: Lifecycle) -> Transition: - if state.python is not Phase.SPAWNED: - return state, () - if state.stop_requested: - return replace(state, python=Phase.LISTENING), () - return replace(state, python=Phase.LISTENING, rust=Phase.SPAWNED), (Spawn(Child.RUST),) - - -def _python_exited(state: Lifecycle, status: int) -> Transition: - if isinstance(state.python, Exited): - return state, () - exited: Final = replace(state, python=Exited(status)) - return exited, ((Send(Child.RUST, signal.SIGTERM),) if state.rust is Phase.SPAWNED else ()) - - -def _rust_exited(state: Lifecycle, status: int) -> Transition: - if state.rust is not Phase.SPAWNED: - return state, () - if state.stop_requested or isinstance(state.python, Exited): - return replace(state, rust=Exited(status)), () - return replace(state, rust=Lost(status)), (Send(Child.PYTHON, signal.SIGTERM),) - - -def step(state: Lifecycle, event: Event, stop_signals: frozenset[signal.Signals]) -> Transition: - """The supervisor never gives a signal or an exit a meaning of its own. Every signal reaches Python as sent, - Python's exit status is the supervisor's, and Rust only follows Python, except that losing Rust stops Python.""" - match event: - case Caught(signum): - return _relay(state, signum, signum in stop_signals) - case Probe.PYTHON_LISTENING: - return _python_listening(state) - case ChildExited(child, status): - match child: - case Child.PYTHON: - return _python_exited(state, status) - case Child.RUST: - return _rust_exited(state, status) - case _: - assert_never(child) - case _: - assert_never(event) - - -def exit_status(state: Lifecycle) -> int | None: - match state.python, state.rust: - case Exited(status), Phase.NOT_STARTED | Exited(): - return status - case Exited(), Lost(): - return FAILURE - case _: - return None - - -@dataclass(frozen=True, slots=True) -class ChildSpec: - argv: tuple[str, ...] - env_overrides: Mapping[str, str] - spawn: SpawnOptions - - -def python_spec( - platform: Platform, command: Sequence[str], private: PrivateListener, server: PythonServer -) -> ChildSpec: - return ChildSpec( - argv=(*command, "--port", str(private.port)), - env_overrides=supervised_environment( - Supervised(private.listener, server.proxy_headers, server.trusted_proxies) - ), - spawn=platform.spawn_options(private.pass_fds), - ) - - -def rust_spec(platform: Platform, binary: str, upstream_port: int, server: PythonServer) -> ChildSpec: - return ChildSpec( - argv=(binary,), - env_overrides=MappingProxyType( - { - "LITELLM_FORWARD_UPSTREAM": f"http://{RUST_HOP}:{upstream_port}", - "LITELLM_FORWARD_HOST": server.host, - "LITELLM_FORWARD_PORT": str(server.port), - "LITELLM_FORWARD_UPSTREAM_KEEPALIVE": str(server.keepalive_seconds), - "LITELLM_FORWARD_SHUTDOWN_DRAIN": str( - server.graceful_timeout_seconds + GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS - ), - } - ), - spawn=platform.spawn_options(()), - ) - - -async def _spawn(spec: ChildSpec) -> asyncio.subprocess.Process: - return await asyncio.create_subprocess_exec( - *spec.argv, - stdin=asyncio.subprocess.PIPE, - env={**os.environ, **spec.env_overrides}, - pass_fds=spec.spawn.pass_fds, - start_new_session=spec.spawn.start_new_session, - creationflags=spec.spawn.creationflags, - ) - - -Launch: TypeAlias = Callable[[asyncio.Queue[Event]], contextlib.AbstractAsyncContextManager[asyncio.subprocess.Process]] - - -async def _report_exit(child: Child, process: asyncio.subprocess.Process, events: asyncio.Queue[Event]) -> None: - events.put_nowait(ChildExited(child, await process.wait())) - - -async def _accepts(port: int) -> bool: - try: - _, writer = await asyncio.wait_for(asyncio.open_connection(RUST_HOP, port), PROBE_INTERVAL_SECONDS) - except (OSError, asyncio.TimeoutError): - return False - writer.close() - return True - - -async def _report_once_listening(port: int, events: asyncio.Queue[Event]) -> None: - while not await _accepts(port): - await asyncio.sleep(PROBE_INTERVAL_SECONDS) - events.put_nowait(Probe.PYTHON_LISTENING) - - -@contextlib.asynccontextmanager -async def _in_background(work: Coroutine[None, None, None]) -> AsyncGenerator[None]: - task: Final = asyncio.create_task(work) - try: - yield - finally: - task.cancel() - await asyncio.gather(task, return_exceptions=True) - - -# A child's stdin is a pipe only the supervisor holds open, so the child sees EOF as soon as the supervisor is gone, -# however it went, and shuts down instead of outliving it. -@contextlib.asynccontextmanager -async def _launch( - platform: Platform, - child: Child, - spec: ChildSpec, - events: asyncio.Queue[Event], -) -> AsyncGenerator[asyncio.subprocess.Process]: - process: Final = await _spawn(spec) - waiter: Final = asyncio.create_task(_report_exit(child, process, events)) - try: - yield process - finally: - platform.stop_group(process) - await waiter - if process.stdin is not None: - process.stdin.close() - - -@contextlib.asynccontextmanager -async def _launch_python( - platform: Platform, spec: ChildSpec, port: int, events: asyncio.Queue[Event] -) -> AsyncGenerator[asyncio.subprocess.Process]: - async with ( - _launch(platform, Child.PYTHON, spec, events) as process, - _in_background(_report_once_listening(port, events)), - ): - yield process - - -async def _run(platform: Platform, events: asyncio.Queue[Event], launches: Mapping[Child, Launch]) -> int: - async with contextlib.AsyncExitStack() as stack: - children: Final[dict[Child, asyncio.subprocess.Process]] = {} # mutable-ok: indexes the stack's children - state, actions = boot() - while True: - for action in actions: - match action: - case Spawn(child): - children[child] = await stack.enter_async_context(launches[child](events)) - case Send(child, signum): - platform.signal_child(children[child], signum) - case _: - assert_never(action) - if (status := exit_status(state)) is not None: - return status - state, actions = step(state, await events.get(), platform.stop_signals) - - -async def _supervise(platform: Platform, binary: str, command: Sequence[str], server: PythonServer) -> int: - events: Final[asyncio.Queue[Event]] = asyncio.Queue() - platform.catch(asyncio.get_running_loop(), lambda signum: events.put_nowait(Caught(signum))) - with platform.private_listener() as private: - launches: Final[Mapping[Child, Launch]] = MappingProxyType( - { - Child.PYTHON: partial( - _launch_python, platform, python_spec(platform, command, private, server), private.port - ), - Child.RUST: partial(_launch, platform, Child.RUST, rust_spec(platform, binary, private.port, server)), - } - ) - return await _run(platform, events, launches) - - -def _exit_like(status: int) -> int: - """A child killed by a signal kills the supervisor with the same signal, as uvicorn re-raises the one it caught.""" - if status >= 0: - return status - signum: Final = signal.Signals(-status) - signal.signal(signum, signal.SIG_DFL) - signal.raise_signal(signum) - return 128 + signum - - -def _nothing_to_prepare() -> None: - pass - - -def supervise( - server: PythonServer, - command: Sequence[str], - before_start: Callable[[], None] = _nothing_to_prepare, - scripts_dir: str = sysconfig.get_path("scripts"), - platform: Platform | None = None, -) -> int | StartupError: - """Serve ``server``'s public port from Rust and run ``command`` behind it, or say why that cannot start. - - A refusal is returned for the caller to serve from Python alone, and ``before_start`` runs only once forwarding - is known to be possible, so a refused start leaves nothing behind. - """ - chosen: Final = current_platform() if platform is None else platform - checked: Final = startup_check(chosen, server, gateway_binary(chosen, scripts_dir)) - if isinstance(checked, StartupError): - sys.stderr.write(f"{checked.reason}, so Python serves the public port without Rust forwarding\n") - return checked - before_start() - try: - status: Final = asyncio.run(_supervise(chosen, checked, command, server)) - except OSError as error: - sys.stderr.write(f"Rust gateway startup failed: {error}\n") - return FAILURE - return _exit_like(status) diff --git a/litellm/proxy/rust_forward/supervisor/__init__.py b/litellm/proxy/rust_forward/supervisor/__init__.py new file mode 100644 index 00000000000..2819f837da3 --- /dev/null +++ b/litellm/proxy/rust_forward/supervisor/__init__.py @@ -0,0 +1,41 @@ +import asyncio +import sys +import sysconfig +from collections.abc import Callable, Sequence +from typing import Final + +from litellm.proxy.rust_forward.platform import Platform, current_platform +from litellm.proxy.rust_forward.server import PythonServer +from litellm.proxy.rust_forward.supervisor.lifecycle import FAILURE +from litellm.proxy.rust_forward.supervisor.preflight import StartupError, gateway_binary, startup_check +from litellm.proxy.rust_forward.supervisor.runtime import exit_like, run + + +def _nothing_to_prepare() -> None: + pass + + +def supervise( + server: PythonServer, + command: Sequence[str], + before_start: Callable[[], None] = _nothing_to_prepare, + scripts_dir: str = sysconfig.get_path("scripts"), + platform: Platform | None = None, +) -> int | StartupError: + """Serve ``server``'s public port from Rust and run ``command`` behind it, or say why that cannot start. + + A refusal is returned for the caller to serve from Python alone, and ``before_start`` runs only once forwarding + is known to be possible, so a refused start leaves nothing behind. + """ + chosen: Final = current_platform() if platform is None else platform + checked: Final = startup_check(chosen, server, gateway_binary(chosen, scripts_dir)) + if isinstance(checked, StartupError): + sys.stderr.write(f"{checked.reason}, so Python serves the public port without Rust forwarding\n") + return checked + before_start() + try: + status: Final = asyncio.run(run(chosen, checked, command, server)) + except OSError as error: + sys.stderr.write(f"Rust gateway startup failed: {error}\n") + return FAILURE + return exit_like(status) diff --git a/litellm/proxy/rust_forward/supervisor/children.py b/litellm/proxy/rust_forward/supervisor/children.py new file mode 100644 index 00000000000..c1da60519ed --- /dev/null +++ b/litellm/proxy/rust_forward/supervisor/children.py @@ -0,0 +1,45 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType + +from litellm.constants import GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS +from litellm.proxy.rust_forward.mode import RUST_HOP, Supervised, supervised_environment +from litellm.proxy.rust_forward.platform import Platform, PrivateListener, SpawnOptions +from litellm.proxy.rust_forward.server import PythonServer + + +@dataclass(frozen=True, slots=True) +class ChildSpec: + argv: tuple[str, ...] + env_overrides: Mapping[str, str] + spawn: SpawnOptions + + +def python_spec( + platform: Platform, command: Sequence[str], private: PrivateListener, server: PythonServer +) -> ChildSpec: + return ChildSpec( + argv=(*command, "--port", str(private.port)), + env_overrides=supervised_environment( + Supervised(private.listener, server.proxy_headers, server.trusted_proxies) + ), + spawn=platform.spawn_options(private.pass_fds), + ) + + +def rust_spec(platform: Platform, binary: str, upstream_port: int, server: PythonServer) -> ChildSpec: + return ChildSpec( + argv=(binary,), + env_overrides=MappingProxyType( + { + "LITELLM_FORWARD_UPSTREAM": f"http://{RUST_HOP}:{upstream_port}", + "LITELLM_FORWARD_HOST": server.host, + "LITELLM_FORWARD_PORT": str(server.port), + "LITELLM_FORWARD_UPSTREAM_KEEPALIVE": str(server.keepalive_seconds), + "LITELLM_FORWARD_SHUTDOWN_DRAIN": str( + server.graceful_timeout_seconds + GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS + ), + } + ), + spawn=platform.spawn_options(()), + ) diff --git a/litellm/proxy/rust_forward/supervisor/lifecycle.py b/litellm/proxy/rust_forward/supervisor/lifecycle.py new file mode 100644 index 00000000000..e462c7df08a --- /dev/null +++ b/litellm/proxy/rust_forward/supervisor/lifecycle.py @@ -0,0 +1,138 @@ +import signal +from dataclasses import dataclass, replace +from enum import Enum, auto +from typing import Final, Literal, TypeAlias + +from typing_extensions import assert_never + + +class Child(Enum): + RUST = "rust" + PYTHON = "python" + + +FAILURE: Final = 1 + + +class Phase(Enum): + NOT_STARTED = auto() + SPAWNED = auto() + LISTENING = auto() + + +@dataclass(frozen=True, slots=True) +class Exited: + status: int + + +@dataclass(frozen=True, slots=True) +class Lost: + """Rust exited while Python was serving and nobody had asked either to stop.""" + + status: int + + +PythonState: TypeAlias = Literal[Phase.SPAWNED, Phase.LISTENING] | Exited +RustState: TypeAlias = Literal[Phase.NOT_STARTED, Phase.SPAWNED] | Exited | Lost + + +@dataclass(frozen=True, slots=True) +class Lifecycle: + python: PythonState + rust: RustState + stop_requested: bool + + +@dataclass(frozen=True, slots=True) +class Caught: + signum: signal.Signals + + +class Probe(Enum): + PYTHON_LISTENING = auto() + + +@dataclass(frozen=True, slots=True) +class ChildExited: + child: Child + status: int + + +@dataclass(frozen=True, slots=True) +class Spawn: + child: Child + + +@dataclass(frozen=True, slots=True) +class Send: + child: Child + signum: signal.Signals + + +Event: TypeAlias = Caught | Literal[Probe.PYTHON_LISTENING] | ChildExited +Action: TypeAlias = Spawn | Send +Transition: TypeAlias = tuple[Lifecycle, tuple[Action, ...]] + + +def boot() -> Transition: + return Lifecycle(Phase.SPAWNED, Phase.NOT_STARTED, stop_requested=False), (Spawn(Child.PYTHON),) + + +def _relay(state: Lifecycle, signum: signal.Signals, stops: bool) -> Transition: + to_rust: Final = (Send(Child.RUST, signum),) if stops and state.rust is Phase.SPAWNED else () + to_python: Final = (Send(Child.PYTHON, signum),) if not isinstance(state.python, Exited) else () + return replace(state, stop_requested=state.stop_requested or stops), (*to_rust, *to_python) + + +# Python alone takes no connections until it has booted, so Rust only opens the public port once Python listens. +def _python_listening(state: Lifecycle) -> Transition: + if state.python is not Phase.SPAWNED: + return state, () + if state.stop_requested: + return replace(state, python=Phase.LISTENING), () + return replace(state, python=Phase.LISTENING, rust=Phase.SPAWNED), (Spawn(Child.RUST),) + + +def _python_exited(state: Lifecycle, status: int) -> Transition: + if isinstance(state.python, Exited): + return state, () + exited: Final = replace(state, python=Exited(status)) + return exited, ((Send(Child.RUST, signal.SIGTERM),) if state.rust is Phase.SPAWNED else ()) + + +def _rust_exited(state: Lifecycle, status: int) -> Transition: + if state.rust is not Phase.SPAWNED: + return state, () + if state.stop_requested or isinstance(state.python, Exited): + return replace(state, rust=Exited(status)), () + return replace(state, rust=Lost(status)), (Send(Child.PYTHON, signal.SIGTERM),) + + +def step(state: Lifecycle, event: Event, stop_signals: frozenset[signal.Signals]) -> Transition: + """The supervisor never gives a signal or an exit a meaning of its own. Every signal reaches Python as sent, + Python's exit status is the supervisor's, and Rust only follows Python, except that losing Rust stops Python.""" + match event: + case Caught(signum): + return _relay(state, signum, signum in stop_signals) + case Probe.PYTHON_LISTENING: + return _python_listening(state) + case ChildExited(child, status): + match child: + case Child.PYTHON: + return _python_exited(state, status) + case Child.RUST: + return _rust_exited(state, status) + case _: + assert_never(child) + case _: + assert_never(event) + + +def exit_status(state: Lifecycle) -> int | None: + match state.python, state.rust: + case Exited(status), Phase.NOT_STARTED | Exited(): + return status + case Exited(), Lost(): + return FAILURE + case _: + return None diff --git a/litellm/proxy/rust_forward/supervisor/preflight.py b/litellm/proxy/rust_forward/supervisor/preflight.py new file mode 100644 index 00000000000..cfd5dd948e6 --- /dev/null +++ b/litellm/proxy/rust_forward/supervisor/preflight.py @@ -0,0 +1,27 @@ +import shutil +from dataclasses import dataclass + +from litellm.proxy.rust_forward.platform import Platform +from litellm.proxy.rust_forward.server import Listener, PythonServer, Server +from litellm.rust_bridge.configuration import ENV_NAME + + +@dataclass(frozen=True, slots=True) +class StartupError: + reason: str + + +def gateway_binary(platform: Platform, scripts_dir: str) -> str | None: + return shutil.which(platform.gateway_name, path=scripts_dir) + + +def startup_check(platform: Platform, server: PythonServer, binary: str | None) -> str | StartupError: + if server.server in (Server.HYPERCORN, Server.GRANIAN): + return StartupError(f"{ENV_NAME} does not support {server.server.value}") + if server.server not in platform.servers: + return StartupError(f"{ENV_NAME} does not support {server.server.value} on {platform.name}") + if server.listener is not Listener.TCP: + return StartupError(f"{ENV_NAME} does not support {server.listener.value} on the public port") + if binary is None: + return StartupError(f"{ENV_NAME} requires the litellm-gateway binary from the litellm-proxy package") + return binary diff --git a/litellm/proxy/rust_forward/supervisor/runtime.py b/litellm/proxy/rust_forward/supervisor/runtime.py new file mode 100644 index 00000000000..1e8aced9077 --- /dev/null +++ b/litellm/proxy/rust_forward/supervisor/runtime.py @@ -0,0 +1,142 @@ +import asyncio +import contextlib +import os +import signal +from collections.abc import AsyncGenerator, Callable, Coroutine, Mapping, Sequence +from functools import partial +from types import MappingProxyType +from typing import Final, TypeAlias + +from typing_extensions import assert_never + +from litellm.proxy.rust_forward.mode import RUST_HOP +from litellm.proxy.rust_forward.platform import Platform +from litellm.proxy.rust_forward.server import PythonServer +from litellm.proxy.rust_forward.supervisor.children import ChildSpec, python_spec, rust_spec +from litellm.proxy.rust_forward.supervisor.lifecycle import ( + Caught, + Child, + ChildExited, + Event, + Probe, + Send, + Spawn, + boot, + exit_status, + step, +) + +PROBE_INTERVAL_SECONDS: Final = 0.05 + +Launch: TypeAlias = Callable[[asyncio.Queue[Event]], contextlib.AbstractAsyncContextManager[asyncio.subprocess.Process]] + + +async def run(platform: Platform, binary: str, command: Sequence[str], server: PythonServer) -> int: + """Run Python and Rust under ``lifecycle.step`` until it decides the exit status.""" + events: Final[asyncio.Queue[Event]] = asyncio.Queue() + platform.catch(asyncio.get_running_loop(), lambda signum: events.put_nowait(Caught(signum))) + with platform.private_listener() as private: + launches: Final[Mapping[Child, Launch]] = MappingProxyType( + { + Child.PYTHON: partial( + _launch_python, platform, python_spec(platform, command, private, server), private.port + ), + Child.RUST: partial(_launch, platform, Child.RUST, rust_spec(platform, binary, private.port, server)), + } + ) + return await _drive(platform, events, launches) + + +def exit_like(status: int) -> int: + """A child killed by a signal kills the supervisor with the same signal, as uvicorn re-raises the one it caught.""" + if status >= 0: + return status + signum: Final = signal.Signals(-status) + signal.signal(signum, signal.SIG_DFL) + signal.raise_signal(signum) + return 128 + signum + + +async def _drive(platform: Platform, events: asyncio.Queue[Event], launches: Mapping[Child, Launch]) -> int: + async with contextlib.AsyncExitStack() as stack: + children: Final[dict[Child, asyncio.subprocess.Process]] = {} # mutable-ok: indexes the stack's children + state, actions = boot() + while True: + for action in actions: + match action: + case Spawn(child): + children[child] = await stack.enter_async_context(launches[child](events)) + case Send(child, signum): + platform.signal_child(children[child], signum) + case _: + assert_never(action) + if (status := exit_status(state)) is not None: + return status + state, actions = step(state, await events.get(), platform.stop_signals) + + +# A child's stdin is a pipe only the supervisor holds open, so the child sees EOF as soon as the supervisor is gone, +# however it went, and shuts down instead of outliving it. +@contextlib.asynccontextmanager +async def _launch( + platform: Platform, + child: Child, + spec: ChildSpec, + events: asyncio.Queue[Event], +) -> AsyncGenerator[asyncio.subprocess.Process]: + process: Final = await asyncio.create_subprocess_exec( + *spec.argv, + stdin=asyncio.subprocess.PIPE, + env={**os.environ, **spec.env_overrides}, + pass_fds=spec.spawn.pass_fds, + start_new_session=spec.spawn.start_new_session, + creationflags=spec.spawn.creationflags, + ) + waiter: Final = asyncio.create_task(_report_exit(child, process, events)) + try: + yield process + finally: + platform.stop_group(process) + await waiter + if process.stdin is not None: + process.stdin.close() + + +@contextlib.asynccontextmanager +async def _launch_python( + platform: Platform, spec: ChildSpec, port: int, events: asyncio.Queue[Event] +) -> AsyncGenerator[asyncio.subprocess.Process]: + async with ( + _launch(platform, Child.PYTHON, spec, events) as process, + _in_background(_report_once_listening(port, events)), + ): + yield process + + +async def _report_exit(child: Child, process: asyncio.subprocess.Process, events: asyncio.Queue[Event]) -> None: + events.put_nowait(ChildExited(child, await process.wait())) + + +async def _report_once_listening(port: int, events: asyncio.Queue[Event]) -> None: + while not await _accepts(port): + await asyncio.sleep(PROBE_INTERVAL_SECONDS) + events.put_nowait(Probe.PYTHON_LISTENING) + + +async def _accepts(port: int) -> bool: + try: + _, writer = await asyncio.wait_for(asyncio.open_connection(RUST_HOP, port), PROBE_INTERVAL_SECONDS) + except (OSError, asyncio.TimeoutError): + return False + writer.close() + return True + + +@contextlib.asynccontextmanager +async def _in_background(work: Coroutine[None, None, None]) -> AsyncGenerator[None]: + task: Final = asyncio.create_task(work) + try: + yield + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) diff --git a/tests/unit/proxy/rust_forward/platform/__init__.py b/tests/unit/proxy/rust_forward/platform/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/proxy/rust_forward/platform/test_windows.py b/tests/unit/proxy/rust_forward/platform/test_windows.py new file mode 100644 index 00000000000..986fb79ed53 --- /dev/null +++ b/tests/unit/proxy/rust_forward/platform/test_windows.py @@ -0,0 +1,13 @@ +import signal +import subprocess +import sys + +import pytest + +from litellm.proxy.rust_forward.platform.windows import CREATE_NEW_PROCESS_GROUP, CTRL_BREAK_EVENT, SIGBREAK + + +@pytest.mark.skipif(sys.platform != "win32", reason="these names exist only on Windows") +def test_windows_constants_match_the_stdlib() -> None: + assert (CTRL_BREAK_EVENT, SIGBREAK) == (signal.CTRL_BREAK_EVENT, signal.SIGBREAK) + assert CREATE_NEW_PROCESS_GROUP == subprocess.CREATE_NEW_PROCESS_GROUP diff --git a/tests/unit/proxy/rust_forward/supervisor/__init__.py b/tests/unit/proxy/rust_forward/supervisor/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/proxy/rust_forward/supervisor/_support.py b/tests/unit/proxy/rust_forward/supervisor/_support.py new file mode 100644 index 00000000000..cc4287f5ea7 --- /dev/null +++ b/tests/unit/proxy/rust_forward/supervisor/_support.py @@ -0,0 +1,7 @@ +from litellm.proxy.rust_forward.server import Listener, PythonServer, Server + + +def python_server( + server: Server = Server.UVICORN, listener: Listener = Listener.TCP, trusted_proxies: str = "127.0.0.1" +) -> PythonServer: + return PythonServer(server, listener, "0.0.0.0", 4000, 2, 2.5, False, trusted_proxies) diff --git a/tests/unit/proxy/rust_forward/supervisor/test_children.py b/tests/unit/proxy/rust_forward/supervisor/test_children.py new file mode 100644 index 00000000000..fc5d8cf703a --- /dev/null +++ b/tests/unit/proxy/rust_forward/supervisor/test_children.py @@ -0,0 +1,53 @@ +import socket +from typing import Final + +from litellm.constants import GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS +from litellm.proxy.proxy_cli import run_server +from litellm.proxy.rust_forward.mode import AssignedPort, InheritedFd, Supervised, rust_forward_mode +from litellm.proxy.rust_forward.platform import PrivateListener +from litellm.proxy.rust_forward.platform.posix import PosixPlatform +from litellm.proxy.rust_forward.platform.windows import CREATE_NEW_PROCESS_GROUP, WindowsPlatform +from litellm.proxy.rust_forward.supervisor.children import python_spec, rust_spec +from tests.unit.proxy.rust_forward.supervisor._support import python_server + +POSIX: Final = PosixPlatform() +WINDOWS: Final = WindowsPlatform() + + +def test_python_child_serves_the_private_socket_with_the_operators_trust() -> None: + private: Final = PrivateListener(InheritedFd(17), 45555, (17,)) + spec: Final = python_spec(POSIX, ("python",), private, python_server(trusted_proxies="10.0.0.5")) + assert rust_forward_mode(spec.env_overrides) == Supervised(InheritedFd(17), False, "10.0.0.5") + assert (spec.spawn.pass_fds, spec.spawn.start_new_session) == ((17,), True) + + +def test_windows_python_child_binds_the_assigned_port_in_its_own_process_group() -> None: + with WINDOWS.private_listener() as private: + pass + spec: Final = python_spec(WINDOWS, ("python",), private, python_server()) + assert rust_forward_mode(spec.env_overrides) == Supervised(AssignedPort(private.port), False, "127.0.0.1") + assert spec.argv[-2:] == ("--port", str(private.port)) + assert (spec.spawn.pass_fds, spec.spawn.creationflags) == ((), CREATE_NEW_PROCESS_GROUP) + with socket.create_server(("127.0.0.1", private.port)): + pass + + +def test_python_child_keeps_the_public_host_and_its_options() -> None: + arguments: Final = ("--host", "0.0.0.0", "--port", "4000", "--num_workers", "2", "--reload") + spec: Final = python_spec(POSIX, arguments, PrivateListener(InheritedFd(17), 45555, (17,)), python_server()) + params: Final = run_server.make_context("litellm", list(spec.argv)).params + assert (params["host"], params["port"]) == ("0.0.0.0", 45555) + assert (params["num_workers"], params["reload"]) == (2, True) + + +def test_rust_child_serves_public_port_and_forwards_to_private_listener() -> None: + spec: Final = rust_spec(POSIX, "gateway", 45555, python_server()) + assert spec.argv == ("gateway",) + assert spec.spawn.pass_fds == () + assert dict(spec.env_overrides) == { + "LITELLM_FORWARD_UPSTREAM": "http://127.0.0.1:45555", + "LITELLM_FORWARD_HOST": "0.0.0.0", + "LITELLM_FORWARD_PORT": "4000", + "LITELLM_FORWARD_UPSTREAM_KEEPALIVE": "2", + "LITELLM_FORWARD_SHUTDOWN_DRAIN": str(2.5 + GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS), + } diff --git a/tests/unit/proxy/rust_forward/supervisor/test_lifecycle.py b/tests/unit/proxy/rust_forward/supervisor/test_lifecycle.py new file mode 100644 index 00000000000..a71fb6b335e --- /dev/null +++ b/tests/unit/proxy/rust_forward/supervisor/test_lifecycle.py @@ -0,0 +1,167 @@ +import functools +import signal +from dataclasses import dataclass, replace +from typing import Final + +import pytest +from hypothesis import example, given +from hypothesis import strategies as st + +from litellm.proxy.rust_forward.platform.posix import PosixPlatform +from litellm.proxy.rust_forward.supervisor.lifecycle import ( + FAILURE, + Action, + Caught, + Child, + ChildExited, + Event, + Lifecycle, + Phase, + Probe, + Send, + Spawn, + boot, + exit_status, + step, +) + +TERM: Final = signal.SIGTERM +INT: Final = signal.SIGINT +POSIX: Final = PosixPlatform() + + +@dataclass(frozen=True, slots=True) +class World: + """What the operating system knows about the children, kept independently of ``Lifecycle``.""" + + alive: frozenset[Child] = frozenset({Child.PYTHON}) + spawned: frozenset[Child] = frozenset({Child.PYTHON}) + reported_listening: bool = False + stop_requested: bool = False + python_status: int | None = None + rust_lost: bool = False + + +def _expected_exit(world: World) -> int | None: + if world.alive or world.python_status is None: + return None + return FAILURE if world.rust_lost else world.python_status + + +def _expected_relay(world: World, signum: signal.Signals) -> tuple[Action, ...]: + to_rust: Final = (Send(Child.RUST, signum),) if signum in STOPS and Child.RUST in world.alive else () + to_python: Final = (Send(Child.PYTHON, signum),) if Child.PYTHON in world.alive else () + return (*to_rust, *to_python) + + +def _expected_on_exit(world: World, child: Child) -> tuple[Action, ...]: + match child: + case Child.PYTHON: + return (Send(Child.RUST, TERM),) if Child.RUST in world.alive else () + case Child.RUST: + serving: Final = Child.PYTHON in world.alive and not world.stop_requested + return (Send(Child.PYTHON, TERM),) if serving else () + + +def _after(world: World, event: Event) -> World: + match event: + case Caught(signum): + return replace(world, stop_requested=world.stop_requested or signum in STOPS) + case Probe.PYTHON_LISTENING: + return replace(world, reported_listening=world.reported_listening or Child.PYTHON in world.alive) + case ChildExited(child, status) if child in world.alive: + lost: Final = child is Child.RUST and Child.PYTHON in world.alive and not world.stop_requested + return replace( + world, + alive=world.alive - {child}, + python_status=status if child is Child.PYTHON else world.python_status, + rust_lost=world.rust_lost or lost, + ) + case ChildExited(): + return world + + +def _spawned(world: World, actions: tuple[Action, ...]) -> World: + started: Final = frozenset(action.child for action in actions if isinstance(action, Spawn)) + return replace(world, alive=world.alive | started, spawned=world.spawned | started) + + +SIGNALS: Final = POSIX.relayed_signals +STOPS: Final = POSIX.stop_signals +EVENTS: Final = st.lists( + st.one_of( + st.sampled_from(SIGNALS).map(Caught), + st.just(Probe.PYTHON_LISTENING), + st.builds(ChildExited, st.sampled_from(Child), st.integers(-15, 3)), + ), + max_size=40, +) + + +@given(EVENTS) +@example([Caught(TERM), Caught(signal.SIGHUP), Probe.PYTHON_LISTENING]) +@example([Probe.PYTHON_LISTENING, Caught(TERM), Caught(signal.SIGWINCH), ChildExited(Child.RUST, 0)]) +def test_lifecycle_relays_every_signal_and_ends_with_pythons_status(events: list[Event]) -> None: + state, actions = boot() + world = World() + assert actions == (Spawn(Child.PYTHON),) + for event in events: + before = _after(world, event) + state, actions = step(state, event, STOPS) + match event: + case Caught(signum): + assert actions == _expected_relay(before, signum) + case Probe.PYTHON_LISTENING: + starts_rust = Child.PYTHON in world.alive and not world.reported_listening and not world.stop_requested + assert actions == ((Spawn(Child.RUST),) if starts_rust else ()) + case ChildExited(child) if child in world.alive: + assert actions == _expected_on_exit(world, child) + case ChildExited(): + assert actions == () + assert all(action.child in before.alive for action in actions if isinstance(action, Send)) + assert all(action.child not in before.spawned for action in actions if isinstance(action, Spawn)) + world = _spawned(before, actions) + assert exit_status(state) == _expected_exit(world) + + +@given(EVENTS) +def test_rust_starts_at_most_once_and_only_while_python_listens_unstopped(events: list[Event]) -> None: + def advance(acc: tuple[Lifecycle, tuple[tuple[Event, Lifecycle, Action], ...]], event: Event): + state, seen = acc + following, actions = step(state, event, STOPS) + return following, (*seen, *((event, state, action) for action in actions)) + + _, history = functools.reduce(advance, events, (boot()[0], ())) + spawns: Final = tuple((event, state) for event, state, action in history if isinstance(action, Spawn)) + assert len(spawns) <= 1 + assert all( + event is Probe.PYTHON_LISTENING and state.python is Phase.SPAWNED and not state.stop_requested + for event, state in spawns + ) + + +@pytest.mark.parametrize( + ("events", "expected"), + [ + ((Probe.PYTHON_LISTENING, Caught(TERM), ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, -15)), -15), + ( + ( + Probe.PYTHON_LISTENING, + Caught(INT), + Caught(INT), + ChildExited(Child.PYTHON, 1), + ChildExited(Child.RUST, 0), + ), + 1, + ), + ( + (Probe.PYTHON_LISTENING, Caught(signal.SIGHUP), ChildExited(Child.PYTHON, -1), ChildExited(Child.RUST, 0)), + -1, + ), + ((Probe.PYTHON_LISTENING, ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, 0)), FAILURE), + ((ChildExited(Child.PYTHON, 3),), 3), + ], +) +def test_supervisor_exits_as_python_did_unless_rust_was_lost(events: tuple[Event, ...], expected: int) -> None: + final: Final = functools.reduce(lambda state, event: step(state, event, STOPS)[0], events, boot()[0]) + assert exit_status(final) == expected diff --git a/tests/unit/proxy/rust_forward/supervisor/test_preflight.py b/tests/unit/proxy/rust_forward/supervisor/test_preflight.py new file mode 100644 index 00000000000..71899262c30 --- /dev/null +++ b/tests/unit/proxy/rust_forward/supervisor/test_preflight.py @@ -0,0 +1,86 @@ +import sys +from pathlib import Path +from typing import Final + +import pytest + +from litellm.proxy.rust_forward.platform.posix import PosixPlatform +from litellm.proxy.rust_forward.platform.windows import WindowsPlatform +from litellm.proxy.rust_forward.server import Listener, PythonServer, Server +from litellm.proxy.rust_forward.supervisor import supervise +from litellm.proxy.rust_forward.supervisor.preflight import StartupError, gateway_binary, startup_check +from tests.unit.proxy.rust_forward.supervisor._support import python_server + +POSIX: Final = PosixPlatform() +WINDOWS: Final = WindowsPlatform() +POSIX_ONLY: Final = pytest.mark.skipif(sys.platform == "win32", reason="runs a POSIX shebang gateway") + + +@pytest.mark.parametrize( + ("server", "reason"), + [ + (python_server(Server.HYPERCORN), "LITELLM_RUST does not support Hypercorn"), + (python_server(Server.GRANIAN), "LITELLM_RUST does not support Granian"), + (python_server(listener=Listener.TLS), "LITELLM_RUST does not support TLS on the public port"), + ( + python_server(listener=Listener.UNIX_SOCKET), + "LITELLM_RUST does not support a Unix socket on the public port", + ), + ( + python_server(Server.GUNICORN, Listener.INHERITED_FD), + "LITELLM_RUST does not support an inherited file descriptor on the public port", + ), + ], +) +def test_unsupported_servers_are_refused_before_anything_starts( + server: PythonServer, reason: str, capsys: pytest.CaptureFixture[str] +) -> None: + assert startup_check(POSIX, server, "gateway") == StartupError(reason) + prepared: Final[list[bool]] = [] + refused: Final = supervise(server, ("python",), before_start=lambda: prepared.append(True), platform=POSIX) + assert refused == StartupError(reason) + assert (capsys.readouterr().err, prepared) == ( + f"{reason}, so Python serves the public port without Rust forwarding\n", + [], + ) + + +@pytest.mark.parametrize("server", [Server.UVICORN, Server.GUNICORN]) +def test_plain_tcp_uvicorn_and_gunicorn_can_be_forwarded(server: Server) -> None: + assert startup_check(POSIX, python_server(server), "gateway") == "gateway" + + +def test_windows_forwards_uvicorn_and_refuses_gunicorn() -> None: + assert startup_check(WINDOWS, python_server(), "gateway") == "gateway" + assert startup_check(WINDOWS, python_server(Server.GUNICORN), "gateway") == StartupError( + "LITELLM_RUST does not support gunicorn on Windows" + ) + + +def test_missing_binary_is_refused_before_anything_starts(tmp_path: Path) -> None: + prepared: Final[list[bool]] = [] + refused: Final = supervise( + python_server(), + ("python",), + before_start=lambda: prepared.append(True), + scripts_dir=str(tmp_path), + platform=POSIX, + ) + assert (refused, prepared) == ( + StartupError("LITELLM_RUST requires the litellm-gateway binary from the litellm-proxy package"), + [], + ) + + +@POSIX_ONLY +def test_binary_is_found_only_in_the_environment_scripts_dir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + scripts: Final = tmp_path / "scripts" + elsewhere: Final = tmp_path / "elsewhere" + scripts.mkdir() + elsewhere.mkdir() + binary: Final = elsewhere / "litellm-gateway" + binary.write_text("#!/bin/sh\nexit 0\n") + binary.chmod(0o755) + monkeypatch.setenv("PATH", str(elsewhere)) + assert gateway_binary(POSIX, str(scripts)) is None + assert gateway_binary(POSIX, str(elsewhere)) == str(binary) diff --git a/tests/unit/proxy/rust_forward/supervisor/test_runtime.py b/tests/unit/proxy/rust_forward/supervisor/test_runtime.py new file mode 100644 index 00000000000..ac47b3975ce --- /dev/null +++ b/tests/unit/proxy/rust_forward/supervisor/test_runtime.py @@ -0,0 +1,45 @@ +import sys +from pathlib import Path +from typing import Final + +import pytest + +from litellm.proxy.rust_forward.platform import Platform, current_platform +from litellm.proxy.rust_forward.platform.windows import WindowsPlatform +from litellm.proxy.rust_forward.supervisor import supervise +from tests.unit.proxy.rust_forward.supervisor._support import python_server + + +PYTHON_CHILD: Final = """ +import json, os, pathlib, socket, time +time.sleep(0.5) +pathlib.Path(os.environ["MARKERS"], "listening").touch() +private = json.loads(os.environ["_LITELLM_GATEWAY_SUPERVISED"])["listener"] +listener = socket.socket(fileno=private["fd"]) if "fd" in private else socket.create_server(("127.0.0.1", private["port"])) +listener.listen() +time.sleep(10) +""" +RUST_CHILD: Final = """ +import os, pathlib +markers = pathlib.Path(os.environ["MARKERS"]) +(markers / "rust_saw_python_listening").write_text(str((markers / "listening").exists())) +""" + + +def _fake_gateway(scripts: Path) -> Platform: + if sys.platform != "win32": + gateway: Final = scripts / "litellm-gateway" + gateway.write_text(f"#!{sys.executable} -I\n{RUST_CHILD}") + gateway.chmod(0o755) + return current_platform() + script: Final = scripts / "rust_child.py" + script.write_text(RUST_CHILD) + (scripts / "litellm-gateway.cmd").write_text(f'@"{sys.executable}" -I "{script}"\r\n') + return WindowsPlatform(gateway_name="litellm-gateway.cmd") + + +def test_rust_opens_the_public_port_only_once_python_listens(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("MARKERS", str(tmp_path)) + command: Final = (sys.executable, "-I", "-c", PYTHON_CHILD) + supervise(python_server(), command, scripts_dir=str(tmp_path), platform=_fake_gateway(tmp_path)) + assert (tmp_path / "rust_saw_python_listening").read_text() == "True" diff --git a/tests/unit/proxy/rust_forward/test_supervisor.py b/tests/unit/proxy/rust_forward/test_supervisor.py deleted file mode 100644 index bb14443c96e..00000000000 --- a/tests/unit/proxy/rust_forward/test_supervisor.py +++ /dev/null @@ -1,339 +0,0 @@ -import functools -import signal -import socket -import subprocess -import sys -from dataclasses import dataclass, replace -from pathlib import Path -from typing import Final - -import pytest -from hypothesis import example, given -from hypothesis import strategies as st - -from litellm.constants import GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS -from litellm.proxy.proxy_cli import run_server -from litellm.proxy.rust_forward.mode import AssignedPort, InheritedFd, Supervised, rust_forward_mode -from litellm.proxy.rust_forward.platform import Platform, PrivateListener, current_platform -from litellm.proxy.rust_forward.platform.posix import PosixPlatform -from litellm.proxy.rust_forward.platform.windows import ( - CREATE_NEW_PROCESS_GROUP, - CTRL_BREAK_EVENT, - SIGBREAK, - WindowsPlatform, -) -from litellm.proxy.rust_forward.server import Listener, PythonServer, Server -from litellm.proxy.rust_forward.supervisor import ( - FAILURE, - Action, - Caught, - Child, - ChildExited, - Event, - Lifecycle, - Phase, - Probe, - Send, - Spawn, - StartupError, - boot, - exit_status, - gateway_binary, - python_spec, - rust_spec, - startup_check, - step, - supervise, -) - -TERM: Final = signal.SIGTERM -INT: Final = signal.SIGINT -POSIX: Final = PosixPlatform() -WINDOWS: Final = WindowsPlatform() -POSIX_ONLY: Final = pytest.mark.skipif(sys.platform == "win32", reason="runs a POSIX shebang gateway") - - -def _server( - server: Server = Server.UVICORN, listener: Listener = Listener.TCP, trusted_proxies: str = "127.0.0.1" -) -> PythonServer: - return PythonServer(server, listener, "0.0.0.0", 4000, 2, 2.5, False, trusted_proxies) - - -def test_python_child_serves_the_private_socket_with_the_operators_trust() -> None: - private: Final = PrivateListener(InheritedFd(17), 45555, (17,)) - spec: Final = python_spec(POSIX, ("python",), private, _server(trusted_proxies="10.0.0.5")) - assert rust_forward_mode(spec.env_overrides) == Supervised(InheritedFd(17), False, "10.0.0.5") - assert (spec.spawn.pass_fds, spec.spawn.start_new_session) == ((17,), True) - - -def test_windows_python_child_binds_the_assigned_port_in_its_own_process_group() -> None: - with WINDOWS.private_listener() as private: - pass - spec: Final = python_spec(WINDOWS, ("python",), private, _server()) - assert rust_forward_mode(spec.env_overrides) == Supervised(AssignedPort(private.port), False, "127.0.0.1") - assert spec.argv[-2:] == ("--port", str(private.port)) - assert (spec.spawn.pass_fds, spec.spawn.creationflags) == ((), CREATE_NEW_PROCESS_GROUP) - with socket.create_server(("127.0.0.1", private.port)): - pass - - -@pytest.mark.skipif(sys.platform != "win32", reason="these names exist only on Windows") -def test_windows_constants_match_the_stdlib() -> None: - assert (CTRL_BREAK_EVENT, SIGBREAK) == (signal.CTRL_BREAK_EVENT, signal.SIGBREAK) - assert CREATE_NEW_PROCESS_GROUP == subprocess.CREATE_NEW_PROCESS_GROUP - - -def test_python_child_keeps_the_public_host_and_its_options() -> None: - arguments: Final = ("--host", "0.0.0.0", "--port", "4000", "--num_workers", "2", "--reload") - spec: Final = python_spec(POSIX, arguments, PrivateListener(InheritedFd(17), 45555, (17,)), _server()) - params: Final = run_server.make_context("litellm", list(spec.argv)).params - assert (params["host"], params["port"]) == ("0.0.0.0", 45555) - assert (params["num_workers"], params["reload"]) == (2, True) - - -def test_rust_child_serves_public_port_and_forwards_to_private_listener() -> None: - spec: Final = rust_spec(POSIX, "gateway", 45555, _server()) - assert spec.argv == ("gateway",) - assert spec.spawn.pass_fds == () - assert dict(spec.env_overrides) == { - "LITELLM_FORWARD_UPSTREAM": "http://127.0.0.1:45555", - "LITELLM_FORWARD_HOST": "0.0.0.0", - "LITELLM_FORWARD_PORT": "4000", - "LITELLM_FORWARD_UPSTREAM_KEEPALIVE": "2", - "LITELLM_FORWARD_SHUTDOWN_DRAIN": str(2.5 + GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS), - } - - -@pytest.mark.parametrize( - ("server", "reason"), - [ - (_server(Server.HYPERCORN), "LITELLM_RUST does not support Hypercorn"), - (_server(Server.GRANIAN), "LITELLM_RUST does not support Granian"), - (_server(listener=Listener.TLS), "LITELLM_RUST does not support TLS on the public port"), - (_server(listener=Listener.UNIX_SOCKET), "LITELLM_RUST does not support a Unix socket on the public port"), - ( - _server(Server.GUNICORN, Listener.INHERITED_FD), - "LITELLM_RUST does not support an inherited file descriptor on the public port", - ), - ], -) -def test_unsupported_servers_are_refused_before_anything_starts( - server: PythonServer, reason: str, capsys: pytest.CaptureFixture[str] -) -> None: - assert startup_check(POSIX, server, "gateway") == StartupError(reason) - prepared: Final[list[bool]] = [] - refused: Final = supervise(server, ("python",), before_start=lambda: prepared.append(True), platform=POSIX) - assert refused == StartupError(reason) - assert (capsys.readouterr().err, prepared) == ( - f"{reason}, so Python serves the public port without Rust forwarding\n", - [], - ) - - -@pytest.mark.parametrize("server", [Server.UVICORN, Server.GUNICORN]) -def test_plain_tcp_uvicorn_and_gunicorn_can_be_forwarded(server: Server) -> None: - assert startup_check(POSIX, _server(server), "gateway") == "gateway" - - -def test_windows_forwards_uvicorn_and_refuses_gunicorn() -> None: - assert startup_check(WINDOWS, _server(), "gateway") == "gateway" - assert startup_check(WINDOWS, _server(Server.GUNICORN), "gateway") == StartupError( - "LITELLM_RUST does not support gunicorn on Windows" - ) - - -def test_missing_binary_is_refused_before_anything_starts(tmp_path: Path) -> None: - prepared: Final[list[bool]] = [] - refused: Final = supervise( - _server(), ("python",), before_start=lambda: prepared.append(True), scripts_dir=str(tmp_path), platform=POSIX - ) - assert (refused, prepared) == ( - StartupError("LITELLM_RUST requires the litellm-gateway binary from the litellm-proxy package"), - [], - ) - - -@POSIX_ONLY -def test_binary_is_found_only_in_the_environment_scripts_dir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - scripts: Final = tmp_path / "scripts" - elsewhere: Final = tmp_path / "elsewhere" - scripts.mkdir() - elsewhere.mkdir() - binary: Final = elsewhere / "litellm-gateway" - binary.write_text("#!/bin/sh\nexit 0\n") - binary.chmod(0o755) - monkeypatch.setenv("PATH", str(elsewhere)) - assert gateway_binary(POSIX, str(scripts)) is None - assert gateway_binary(POSIX, str(elsewhere)) == str(binary) - - -PYTHON_CHILD: Final = """ -import json, os, pathlib, socket, time -time.sleep(0.5) -pathlib.Path(os.environ["MARKERS"], "listening").touch() -private = json.loads(os.environ["_LITELLM_GATEWAY_SUPERVISED"])["listener"] -listener = socket.socket(fileno=private["fd"]) if "fd" in private else socket.create_server(("127.0.0.1", private["port"])) -listener.listen() -time.sleep(10) -""" -RUST_CHILD: Final = """ -import os, pathlib -markers = pathlib.Path(os.environ["MARKERS"]) -(markers / "rust_saw_python_listening").write_text(str((markers / "listening").exists())) -""" - - -def _fake_gateway(scripts: Path) -> Platform: - if sys.platform != "win32": - gateway: Final = scripts / "litellm-gateway" - gateway.write_text(f"#!{sys.executable} -I\n{RUST_CHILD}") - gateway.chmod(0o755) - return current_platform() - script: Final = scripts / "rust_child.py" - script.write_text(RUST_CHILD) - (scripts / "litellm-gateway.cmd").write_text(f'@"{sys.executable}" -I "{script}"\r\n') - return WindowsPlatform(gateway_name="litellm-gateway.cmd") - - -def test_rust_opens_the_public_port_only_once_python_listens(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("MARKERS", str(tmp_path)) - command: Final = (sys.executable, "-I", "-c", PYTHON_CHILD) - supervise(_server(), command, scripts_dir=str(tmp_path), platform=_fake_gateway(tmp_path)) - assert (tmp_path / "rust_saw_python_listening").read_text() == "True" - - -@dataclass(frozen=True, slots=True) -class World: - """What the operating system knows about the children, kept independently of ``Lifecycle``.""" - - alive: frozenset[Child] = frozenset({Child.PYTHON}) - spawned: frozenset[Child] = frozenset({Child.PYTHON}) - reported_listening: bool = False - stop_requested: bool = False - python_status: int | None = None - rust_lost: bool = False - - -def _expected_exit(world: World) -> int | None: - if world.alive or world.python_status is None: - return None - return FAILURE if world.rust_lost else world.python_status - - -def _expected_relay(world: World, signum: signal.Signals) -> tuple[Action, ...]: - to_rust: Final = (Send(Child.RUST, signum),) if signum in STOPS and Child.RUST in world.alive else () - to_python: Final = (Send(Child.PYTHON, signum),) if Child.PYTHON in world.alive else () - return (*to_rust, *to_python) - - -def _expected_on_exit(world: World, child: Child) -> tuple[Action, ...]: - match child: - case Child.PYTHON: - return (Send(Child.RUST, TERM),) if Child.RUST in world.alive else () - case Child.RUST: - serving: Final = Child.PYTHON in world.alive and not world.stop_requested - return (Send(Child.PYTHON, TERM),) if serving else () - - -def _after(world: World, event: Event) -> World: - match event: - case Caught(signum): - return replace(world, stop_requested=world.stop_requested or signum in STOPS) - case Probe.PYTHON_LISTENING: - return replace(world, reported_listening=world.reported_listening or Child.PYTHON in world.alive) - case ChildExited(child, status) if child in world.alive: - lost: Final = child is Child.RUST and Child.PYTHON in world.alive and not world.stop_requested - return replace( - world, - alive=world.alive - {child}, - python_status=status if child is Child.PYTHON else world.python_status, - rust_lost=world.rust_lost or lost, - ) - case ChildExited(): - return world - - -def _spawned(world: World, actions: tuple[Action, ...]) -> World: - started: Final = frozenset(action.child for action in actions if isinstance(action, Spawn)) - return replace(world, alive=world.alive | started, spawned=world.spawned | started) - - -SIGNALS: Final = POSIX.relayed_signals -STOPS: Final = POSIX.stop_signals -EVENTS: Final = st.lists( - st.one_of( - st.sampled_from(SIGNALS).map(Caught), - st.just(Probe.PYTHON_LISTENING), - st.builds(ChildExited, st.sampled_from(Child), st.integers(-15, 3)), - ), - max_size=40, -) - - -@given(EVENTS) -@example([Caught(TERM), Caught(signal.SIGHUP), Probe.PYTHON_LISTENING]) -@example([Probe.PYTHON_LISTENING, Caught(TERM), Caught(signal.SIGWINCH), ChildExited(Child.RUST, 0)]) -def test_lifecycle_relays_every_signal_and_ends_with_pythons_status(events: list[Event]) -> None: - state, actions = boot() - world = World() - assert actions == (Spawn(Child.PYTHON),) - for event in events: - before = _after(world, event) - state, actions = step(state, event, STOPS) - match event: - case Caught(signum): - assert actions == _expected_relay(before, signum) - case Probe.PYTHON_LISTENING: - starts_rust = Child.PYTHON in world.alive and not world.reported_listening and not world.stop_requested - assert actions == ((Spawn(Child.RUST),) if starts_rust else ()) - case ChildExited(child) if child in world.alive: - assert actions == _expected_on_exit(world, child) - case ChildExited(): - assert actions == () - assert all(action.child in before.alive for action in actions if isinstance(action, Send)) - assert all(action.child not in before.spawned for action in actions if isinstance(action, Spawn)) - world = _spawned(before, actions) - assert exit_status(state) == _expected_exit(world) - - -@given(EVENTS) -def test_rust_starts_at_most_once_and_only_while_python_listens_unstopped(events: list[Event]) -> None: - def advance(acc: tuple[Lifecycle, tuple[tuple[Event, Lifecycle, Action], ...]], event: Event): - state, seen = acc - following, actions = step(state, event, STOPS) - return following, (*seen, *((event, state, action) for action in actions)) - - _, history = functools.reduce(advance, events, (boot()[0], ())) - spawns: Final = tuple((event, state) for event, state, action in history if isinstance(action, Spawn)) - assert len(spawns) <= 1 - assert all( - event is Probe.PYTHON_LISTENING and state.python is Phase.SPAWNED and not state.stop_requested - for event, state in spawns - ) - - -@pytest.mark.parametrize( - ("events", "expected"), - [ - ((Probe.PYTHON_LISTENING, Caught(TERM), ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, -15)), -15), - ( - ( - Probe.PYTHON_LISTENING, - Caught(INT), - Caught(INT), - ChildExited(Child.PYTHON, 1), - ChildExited(Child.RUST, 0), - ), - 1, - ), - ( - (Probe.PYTHON_LISTENING, Caught(signal.SIGHUP), ChildExited(Child.PYTHON, -1), ChildExited(Child.RUST, 0)), - -1, - ), - ((Probe.PYTHON_LISTENING, ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, 0)), FAILURE), - ((ChildExited(Child.PYTHON, 3),), 3), - ], -) -def test_supervisor_exits_as_python_did_unless_rust_was_lost(events: tuple[Event, ...], expected: int) -> None: - final: Final = functools.reduce(lambda state, event: step(state, event, STOPS)[0], events, boot()[0]) - assert exit_status(final) == expected diff --git a/tests/unit/proxy/test_proxy_cli.py b/tests/unit/proxy/test_proxy_cli.py index 5ad21572605..df7bbd43324 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -4313,7 +4313,7 @@ class TestRustForwardServer: ) def test_unsupported_servers_are_refused_before_starting_children(self, argv: list[str], reason: str) -> None: from litellm.proxy.rust_forward.platform.posix import PosixPlatform - from litellm.proxy.rust_forward.supervisor import StartupError, startup_check + from litellm.proxy.rust_forward.supervisor.preflight import StartupError, startup_check _, server = self._server(argv) assert startup_check(PosixPlatform(), server, "gateway") == StartupError(f"LITELLM_RUST {reason}") @@ -4417,7 +4417,7 @@ class TestRustForwardServer: from litellm.proxy.rust_forward.mode import InheritedFd, Supervised, rust_forward_mode from litellm.proxy.rust_forward.platform import PrivateListener from litellm.proxy.rust_forward.platform.posix import PosixPlatform - from litellm.proxy.rust_forward.supervisor import python_spec + from litellm.proxy.rust_forward.supervisor.children import python_spec monkeypatch.delenv("FORWARDED_ALLOW_IPS", raising=False) _, server = self._server(["--run_gunicorn"] if gunicorn else []) From 6ae722d049ff35d3606fa9bbf86877195a7792d9 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 6 Nov 2026 14:59:11 +0900 Subject: [PATCH 08/10] refactor(proxy): run the forwarding supervisor's children on anyio runtime.py hand-wired an asyncio.Queue, a map of partials and per-child background tasks. The driver loop now takes an Executor protocol, and open_children implements it with an AsyncExitStack, an anyio task group and a memory object stream, so the event plumbing lives in one place and the loop has fast unit tests with a scripted executor Stopping a child is shielded from cancellation, since anyio cancels every await in a cancelled scope and an unshielded wait would leave a child still stopping. A child that cannot be spawned comes back as SpawnFailed instead of an OSError, which an anyio task group would wrap in an ExceptionGroup --- .../proxy/rust_forward/supervisor/__init__.py | 22 ++- .../proxy/rust_forward/supervisor/runtime.py | 187 +++++++++++------- .../rust_forward/supervisor/test_runtime.py | 109 ++++++++++ 3 files changed, 240 insertions(+), 78 deletions(-) diff --git a/litellm/proxy/rust_forward/supervisor/__init__.py b/litellm/proxy/rust_forward/supervisor/__init__.py index 2819f837da3..bfb09ddabe2 100644 --- a/litellm/proxy/rust_forward/supervisor/__init__.py +++ b/litellm/proxy/rust_forward/supervisor/__init__.py @@ -4,11 +4,13 @@ import sysconfig from collections.abc import Callable, Sequence from typing import Final +from typing_extensions import assert_never + from litellm.proxy.rust_forward.platform import Platform, current_platform from litellm.proxy.rust_forward.server import PythonServer from litellm.proxy.rust_forward.supervisor.lifecycle import FAILURE from litellm.proxy.rust_forward.supervisor.preflight import StartupError, gateway_binary, startup_check -from litellm.proxy.rust_forward.supervisor.runtime import exit_like, run +from litellm.proxy.rust_forward.supervisor.runtime import SpawnFailed, exit_like, run def _nothing_to_prepare() -> None: @@ -34,8 +36,18 @@ def supervise( return checked before_start() try: - status: Final = asyncio.run(run(chosen, checked, command, server)) + outcome: Final = asyncio.run(run(chosen, checked, command, server)) except OSError as error: - sys.stderr.write(f"Rust gateway startup failed: {error}\n") - return FAILURE - return exit_like(status) + return _startup_failed(error) + match outcome: + case SpawnFailed(error): + return _startup_failed(error) + case int(): + return exit_like(outcome) + case _: + assert_never(outcome) + + +def _startup_failed(error: OSError) -> int: + sys.stderr.write(f"Rust gateway startup failed: {error}\n") + return FAILURE diff --git a/litellm/proxy/rust_forward/supervisor/runtime.py b/litellm/proxy/rust_forward/supervisor/runtime.py index 1e8aced9077..8fbb2af4133 100644 --- a/litellm/proxy/rust_forward/supervisor/runtime.py +++ b/litellm/proxy/rust_forward/supervisor/runtime.py @@ -1,12 +1,16 @@ import asyncio import contextlib +import math import os import signal -from collections.abc import AsyncGenerator, Callable, Coroutine, Mapping, Sequence -from functools import partial +from collections.abc import AsyncGenerator, Mapping, Sequence +from dataclasses import dataclass from types import MappingProxyType -from typing import Final, TypeAlias +from typing import Final, Protocol +import anyio +from anyio.abc import TaskGroup +from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream from typing_extensions import assert_never from litellm.proxy.rust_forward.mode import RUST_HOP @@ -28,23 +32,49 @@ from litellm.proxy.rust_forward.supervisor.lifecycle import ( PROBE_INTERVAL_SECONDS: Final = 0.05 -Launch: TypeAlias = Callable[[asyncio.Queue[Event]], contextlib.AbstractAsyncContextManager[asyncio.subprocess.Process]] + +@dataclass(frozen=True, slots=True) +class SpawnFailed: + error: OSError -async def run(platform: Platform, binary: str, command: Sequence[str], server: PythonServer) -> int: - """Run Python and Rust under ``lifecycle.step`` until it decides the exit status.""" - events: Final[asyncio.Queue[Event]] = asyncio.Queue() - platform.catch(asyncio.get_running_loop(), lambda signum: events.put_nowait(Caught(signum))) +class Executor(Protocol): + async def spawn(self, child: Child) -> SpawnFailed | None: ... + + def send(self, child: Child, signum: signal.Signals) -> None: ... + + async def next_event(self) -> Event: ... + + +async def run(platform: Platform, binary: str, command: Sequence[str], server: PythonServer) -> int | SpawnFailed: with platform.private_listener() as private: - launches: Final[Mapping[Child, Launch]] = MappingProxyType( + specs: Final = MappingProxyType( { - Child.PYTHON: partial( - _launch_python, platform, python_spec(platform, command, private, server), private.port - ), - Child.RUST: partial(_launch, platform, Child.RUST, rust_spec(platform, binary, private.port, server)), + Child.PYTHON: python_spec(platform, command, private, server), + Child.RUST: rust_spec(platform, binary, private.port, server), } ) - return await _drive(platform, events, launches) + async with open_children(platform, specs, private.port) as children: + platform.catch(asyncio.get_running_loop(), children.caught) + return await drive(children, platform.stop_signals) + + +async def drive(executor: Executor, stop_signals: frozenset[signal.Signals]) -> int | SpawnFailed: + """Carry out what ``lifecycle.step`` decides until it settles the exit status.""" + state, actions = boot() + while True: + for action in actions: + match action: + case Spawn(child): + if (failed := await executor.spawn(child)) is not None: + return failed + case Send(child, signum): + executor.send(child, signum) + case _: + assert_never(action) + if (status := exit_status(state)) is not None: + return status + state, actions = step(state, await executor.next_event(), stop_signals) def exit_like(status: int) -> int: @@ -57,33 +87,59 @@ def exit_like(status: int) -> int: return 128 + signum -async def _drive(platform: Platform, events: asyncio.Queue[Event], launches: Mapping[Child, Launch]) -> int: - async with contextlib.AsyncExitStack() as stack: - children: Final[dict[Child, asyncio.subprocess.Process]] = {} # mutable-ok: indexes the stack's children - state, actions = boot() - while True: - for action in actions: - match action: - case Spawn(child): - children[child] = await stack.enter_async_context(launches[child](events)) - case Send(child, signum): - platform.signal_child(children[child], signum) - case _: - assert_never(action) - if (status := exit_status(state)) is not None: - return status - state, actions = step(state, await events.get(), platform.stop_signals) +@dataclass(frozen=True, slots=True) +class Children: + platform: Platform + specs: Mapping[Child, ChildSpec] + private_port: int + stack: contextlib.AsyncExitStack + tasks: TaskGroup + events_in: MemoryObjectSendStream[Event] + events_out: MemoryObjectReceiveStream[Event] + processes: dict[Child, asyncio.subprocess.Process] # mutable-ok: the running children the stack owns + + def caught(self, signum: signal.Signals) -> None: + with contextlib.suppress(anyio.ClosedResourceError): + self.events_in.send_nowait(Caught(signum)) + + async def spawn(self, child: Child) -> SpawnFailed | None: + try: + process: Final = await self.stack.enter_async_context(_process(self.platform, self.specs[child])) + except OSError as error: + return SpawnFailed(error) + self.processes[child] = process + self.tasks.start_soon(_report_exit, child, process, self.events_in.clone()) + if child is Child.PYTHON: + probe: Final = self.tasks.start_soon(_report_once_listening, self.private_port, self.events_in.clone()) + self.stack.callback(probe.cancel) + return None + + def send(self, child: Child, signum: signal.Signals) -> None: + self.platform.signal_child(self.processes[child], signum) + + async def next_event(self) -> Event: + return await self.events_out.receive() + + +# The task group is entered before any child, so on the way out every child is stopped, last spawned first, before +# the task group waits for the tasks reporting their exits. Signal handlers cannot wait for room, hence the +# unbounded event buffer. +@contextlib.asynccontextmanager +async def open_children( + platform: Platform, specs: Mapping[Child, ChildSpec], private_port: int +) -> AsyncGenerator[Children]: + events_in, events_out = anyio.create_memory_object_stream[Event](math.inf) + with events_in, events_out: + async with contextlib.AsyncExitStack() as stack: + tasks: Final = await stack.enter_async_context(anyio.create_task_group()) + yield Children(platform, specs, private_port, stack, tasks, events_in, events_out, {}) # A child's stdin is a pipe only the supervisor holds open, so the child sees EOF as soon as the supervisor is gone, -# however it went, and shuts down instead of outliving it. +# however it went, and shuts down instead of outliving it. Stopping is shielded so a cancelled supervisor still waits +# for the child it is stopping. @contextlib.asynccontextmanager -async def _launch( - platform: Platform, - child: Child, - spec: ChildSpec, - events: asyncio.Queue[Event], -) -> AsyncGenerator[asyncio.subprocess.Process]: +async def _process(platform: Platform, spec: ChildSpec) -> AsyncGenerator[asyncio.subprocess.Process]: process: Final = await asyncio.create_subprocess_exec( *spec.argv, stdin=asyncio.subprocess.PIPE, @@ -92,51 +148,36 @@ async def _launch( start_new_session=spec.spawn.start_new_session, creationflags=spec.spawn.creationflags, ) - waiter: Final = asyncio.create_task(_report_exit(child, process, events)) try: yield process finally: - platform.stop_group(process) - await waiter + with anyio.CancelScope(shield=True): + platform.stop_group(process) + await process.wait() if process.stdin is not None: process.stdin.close() -@contextlib.asynccontextmanager -async def _launch_python( - platform: Platform, spec: ChildSpec, port: int, events: asyncio.Queue[Event] -) -> AsyncGenerator[asyncio.subprocess.Process]: - async with ( - _launch(platform, Child.PYTHON, spec, events) as process, - _in_background(_report_once_listening(port, events)), - ): - yield process +async def _report_exit( + child: Child, process: asyncio.subprocess.Process, events: MemoryObjectSendStream[Event] +) -> None: + with events: + events.send_nowait(ChildExited(child, await process.wait())) -async def _report_exit(child: Child, process: asyncio.subprocess.Process, events: asyncio.Queue[Event]) -> None: - events.put_nowait(ChildExited(child, await process.wait())) - - -async def _report_once_listening(port: int, events: asyncio.Queue[Event]) -> None: - while not await _accepts(port): - await asyncio.sleep(PROBE_INTERVAL_SECONDS) - events.put_nowait(Probe.PYTHON_LISTENING) +async def _report_once_listening(port: int, events: MemoryObjectSendStream[Event]) -> None: + with events: + while not await _accepts(port): + await anyio.sleep(PROBE_INTERVAL_SECONDS) + events.send_nowait(Probe.PYTHON_LISTENING) async def _accepts(port: int) -> bool: - try: - _, writer = await asyncio.wait_for(asyncio.open_connection(RUST_HOP, port), PROBE_INTERVAL_SECONDS) - except (OSError, asyncio.TimeoutError): - return False - writer.close() - return True - - -@contextlib.asynccontextmanager -async def _in_background(work: Coroutine[None, None, None]) -> AsyncGenerator[None]: - task: Final = asyncio.create_task(work) - try: - yield - finally: - task.cancel() - await asyncio.gather(task, return_exceptions=True) + with anyio.move_on_after(PROBE_INTERVAL_SECONDS): + try: + stream: Final = await anyio.connect_tcp(RUST_HOP, port) + except OSError: + return False + await stream.aclose() + return True + return False diff --git a/tests/unit/proxy/rust_forward/supervisor/test_runtime.py b/tests/unit/proxy/rust_forward/supervisor/test_runtime.py index ac47b3975ce..824ef9e3945 100644 --- a/tests/unit/proxy/rust_forward/supervisor/test_runtime.py +++ b/tests/unit/proxy/rust_forward/supervisor/test_runtime.py @@ -1,12 +1,29 @@ +import signal import sys +from dataclasses import dataclass, field from pathlib import Path +from types import MappingProxyType from typing import Final +import anyio import pytest from litellm.proxy.rust_forward.platform import Platform, current_platform from litellm.proxy.rust_forward.platform.windows import WindowsPlatform from litellm.proxy.rust_forward.supervisor import supervise +from litellm.proxy.rust_forward.supervisor.children import ChildSpec +from litellm.proxy.rust_forward.supervisor.lifecycle import ( + FAILURE, + Action, + Caught, + Child, + ChildExited, + Event, + Probe, + Send, + Spawn, +) +from litellm.proxy.rust_forward.supervisor.runtime import SpawnFailed, drive, open_children from tests.unit.proxy.rust_forward.supervisor._support import python_server @@ -43,3 +60,95 @@ def test_rust_opens_the_public_port_only_once_python_listens(tmp_path: Path, mon command: Final = (sys.executable, "-I", "-c", PYTHON_CHILD) supervise(python_server(), command, scripts_dir=str(tmp_path), platform=_fake_gateway(tmp_path)) assert (tmp_path / "rust_saw_python_listening").read_text() == "True" + + +TERM: Final = signal.SIGTERM +POSIX_ONLY: Final = pytest.mark.skipif(sys.platform == "win32", reason="the child traps SIGTERM") +STOPS: Final = frozenset({signal.SIGTERM, signal.SIGINT, signal.SIGQUIT}) + + +@dataclass +class ScriptedExecutor: + """Replays a fixed event script and records every action the driver carries out.""" + + events: list[Event] + failing: frozenset[Child] = frozenset() + done: list[Action] = field(default_factory=list) + + async def spawn(self, child: Child) -> SpawnFailed | None: + self.done.append(Spawn(child)) + return SpawnFailed(FileNotFoundError(child.value)) if child in self.failing else None + + def send(self, child: Child, signum: signal.Signals) -> None: + self.done.append(Send(child, signum)) + + async def next_event(self) -> Event: + return self.events.pop(0) + + +async def test_driver_carries_out_each_action_in_order_and_ends_with_pythons_status() -> None: + executor: Final = ScriptedExecutor( + [Probe.PYTHON_LISTENING, Caught(TERM), ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, -15)] + ) + assert await drive(executor, STOPS) == -15 + assert executor.done == [Spawn(Child.PYTHON), Spawn(Child.RUST), Send(Child.RUST, TERM), Send(Child.PYTHON, TERM)] + assert executor.events == [] + + +async def test_driver_stops_python_and_fails_when_rust_is_lost() -> None: + executor: Final = ScriptedExecutor( + [Probe.PYTHON_LISTENING, ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, -15)] + ) + assert await drive(executor, STOPS) == FAILURE + assert executor.done == [Spawn(Child.PYTHON), Spawn(Child.RUST), Send(Child.PYTHON, TERM)] + + +async def test_driver_asks_for_no_event_once_the_status_is_settled() -> None: + executor: Final = ScriptedExecutor([ChildExited(Child.PYTHON, 3), Caught(TERM)]) + assert await drive(executor, STOPS) == 3 + assert executor.events == [Caught(TERM)] + + +async def test_driver_returns_a_failed_spawn_without_waiting_for_events() -> None: + executor: Final = ScriptedExecutor([Probe.PYTHON_LISTENING, Caught(TERM)], failing=frozenset({Child.RUST})) + outcome: Final = await drive(executor, STOPS) + assert isinstance(outcome, SpawnFailed) + assert executor.events == [Caught(TERM)] + + +def _children(argv: tuple[str, ...]) -> MappingProxyType[Child, ChildSpec]: + spec: Final = ChildSpec(argv, MappingProxyType({}), current_platform().spawn_options(())) + return MappingProxyType({Child.PYTHON: spec, Child.RUST: spec}) + + +SLOW_TO_STOP: Final = ( + sys.executable, + "-c", + "import signal, time; signal.signal(signal.SIGTERM, lambda *_: (time.sleep(0.5), exit(0))); time.sleep(30)", +) + + +@POSIX_ONLY +async def test_a_cancelled_supervisor_still_waits_for_the_child_it_stops() -> None: + with anyio.CancelScope() as scope: + async with open_children(current_platform(), _children(SLOW_TO_STOP), 1) as children: + assert await children.spawn(Child.RUST) is None + process: Final = children.processes[Child.RUST] + await anyio.sleep(0.3) + scope.cancel() + await anyio.sleep(1) + assert process.returncode == 0 + + +async def test_exiting_children_are_reported_as_events() -> None: + async with open_children(current_platform(), _children((sys.executable, "-c", "exit(7)")), 1) as children: + assert await children.spawn(Child.RUST) is None + with anyio.fail_after(5): + assert await children.next_event() == ChildExited(Child.RUST, 7) + + +async def test_a_child_that_cannot_start_is_returned_as_a_failure(tmp_path: Path) -> None: + async with open_children(current_platform(), _children((str(tmp_path / "missing"),)), 1) as children: + outcome: Final = await children.spawn(Child.RUST) + assert isinstance(outcome, SpawnFailed) + assert isinstance(outcome.error, FileNotFoundError) From 702e26af42e822d8ffa96fd8e7eae40afda86bba Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 6 Nov 2026 15:31:09 +0900 Subject: [PATCH 09/10] test(proxy): state Rust forwarding behavior as tables instead of mirrors The lifecycle test no longer reimplements the state machine. A trace table shows what each event makes the supervisor do, and one Hypothesis property checks it from what an outside observer of the two processes can see. Server option, preflight, driver and CLI tests collapse into one parametrized table each, and cases already covered by the peer equivalence matrix are dropped Co-Authored-By: Claude Opus 5.5 --- tests/unit/gateway/test_launch.py | 57 ++--- .../rust_forward/supervisor/test_lifecycle.py | 233 +++++++++--------- .../rust_forward/supervisor/test_preflight.py | 70 +++--- .../rust_forward/supervisor/test_runtime.py | 69 +++--- tests/unit/proxy/rust_forward/test_mode.py | 67 +++-- tests/unit/proxy/rust_forward/test_peer.py | 10 - tests/unit/proxy/test_proxy_cli.py | 104 +++----- 7 files changed, 280 insertions(+), 330 deletions(-) diff --git a/tests/unit/gateway/test_launch.py b/tests/unit/gateway/test_launch.py index becafe07d5d..a08478d3c2d 100644 --- a/tests/unit/gateway/test_launch.py +++ b/tests/unit/gateway/test_launch.py @@ -12,7 +12,6 @@ from typing import Final, cast from unittest.mock import MagicMock, patch import pytest -import uvicorn from uvicorn.importer import import_from_string from uvicorn.main import main as uvicorn_main @@ -20,7 +19,14 @@ from gateway.launch import GATEWAY_APP, main, pool_database_url, python_server, from litellm.proxy.db.db_url_settings import DatabaseURLSettings from litellm.proxy.db.pgbouncer import PGBOUNCER_POOLED_ENV_VAR, PgBouncerError, PgBouncerSettings from litellm.proxy.proxy_server import app as proxy_app -from litellm.proxy.rust_forward.mode import AssignedPort, InheritedFd, Supervised, uvicorn_flags +from litellm.proxy.rust_forward.mode import ( + AssignedPort, + InheritedFd, + Mode, + RustForwardMode, + Supervised, + uvicorn_flags, +) from litellm.proxy.rust_forward.server import Listener from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager @@ -172,10 +178,6 @@ class TestUvicornArgv: served: Final = _uvicorn_params(uvicorn_argv(argv, os.environ))["timeout_keep_alive"] assert python_server(argv).keepalive_seconds == served - def test_without_a_keepalive_rust_follows_uvicorns_own_default(self, monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.delenv("KEEPALIVE_TIMEOUT", raising=False) - assert python_server(()).keepalive_seconds == uvicorn.Config(app="", log_config=None).timeout_keep_alive - @pytest.mark.parametrize("argv", [(), ("--timeout-graceful-shutdown", "12")]) def test_rust_drains_for_pythons_graceful_timeout_plus_teardown( self, argv: tuple[str, ...], monkeypatch: pytest.MonkeyPatch @@ -185,28 +187,27 @@ class TestUvicornArgv: graceful: Final = GracefulShutdownManager.get_timeout() if flag is None else float(str(flag)) assert python_server(argv).graceful_timeout_seconds == graceful - def test_supervised_listener_uses_the_inherited_loopback_socket(self): - supervised: Final = uvicorn_flags(Supervised(InheritedFd(17), True, "127.0.0.1")) - params: Final = _uvicorn_params((*uvicorn_argv(("--workers", "2"), {}), *supervised)) - assert params["fd"] == 17 - assert params["workers"] == 2 - - def test_assigned_port_listener_binds_only_loopback(self): - supervised: Final = uvicorn_flags(Supervised(AssignedPort(45556), True, "127.0.0.1")) - params: Final = _uvicorn_params((*uvicorn_argv(("--host", "0.0.0.0", "--port", "4000"), {}), *supervised)) - assert (params["host"], params["port"], params["fd"]) == ("127.0.0.1", 45556, None) - - @pytest.mark.parametrize("argv", [(), ("--proxy-headers",), ("--forwarded-allow-ips=*",)]) - def test_private_hop_leaves_proxy_headers_to_the_app(self, argv: tuple[str, ...]): - params: Final = _uvicorn_params( - (*uvicorn_argv(argv, {}), *uvicorn_flags(Supervised(InheritedFd(17), True, "*"))) - ) - assert params["proxy_headers"] is False - - def test_direct_listener_keeps_explicit_proxy_header_options(self): - params: Final = _uvicorn_params(uvicorn_argv(("--no-proxy-headers", "--forwarded-allow-ips=*"), {})) - assert params["proxy_headers"] is False - assert params["forwarded_allow_ips"] == "*" + @pytest.mark.parametrize( + ("mode", "served"), + [ + (Mode.DISABLED, {"host": "0.0.0.0", "port": 4000, "fd": None, "proxy_headers": True}), + ( + Supervised(InheritedFd(17), True, "*"), + {"host": "0.0.0.0", "port": 4000, "fd": 17, "proxy_headers": False}, + ), + ( + Supervised(AssignedPort(45556), True, "*"), + {"host": "127.0.0.1", "port": 45556, "fd": None, "proxy_headers": False}, + ), + ], + ) + def test_supervised_uvicorn_serves_the_private_hop_and_never_trusts_it( + self, mode: RustForwardMode, served: dict[str, object] + ) -> None: + argv: Final = ("--host", "0.0.0.0", "--port", "4000", "--workers", "2", "--proxy-headers") + params: Final = _uvicorn_params((*uvicorn_argv((*argv, "--forwarded-allow-ips=*"), {}), *uvicorn_flags(mode))) + assert {name: params[name] for name in served} == served + assert (params["workers"], params["forwarded_allow_ips"]) == (2, "*") def test_keepalive_env_reaches_uvicorn(self): params: Final = _uvicorn_params(uvicorn_argv(("--workers", "4"), {"KEEPALIVE_TIMEOUT": "75"})) diff --git a/tests/unit/proxy/rust_forward/supervisor/test_lifecycle.py b/tests/unit/proxy/rust_forward/supervisor/test_lifecycle.py index a71fb6b335e..4b479f3d22d 100644 --- a/tests/unit/proxy/rust_forward/supervisor/test_lifecycle.py +++ b/tests/unit/proxy/rust_forward/supervisor/test_lifecycle.py @@ -1,10 +1,11 @@ import functools +import itertools import signal from dataclasses import dataclass, replace from typing import Final import pytest -from hypothesis import example, given +from hypothesis import given from hypothesis import strategies as st from litellm.proxy.rust_forward.platform.posix import PosixPlatform @@ -16,7 +17,6 @@ from litellm.proxy.rust_forward.supervisor.lifecycle import ( ChildExited, Event, Lifecycle, - Phase, Probe, Send, Spawn, @@ -25,73 +25,127 @@ from litellm.proxy.rust_forward.supervisor.lifecycle import ( step, ) +PY: Final = Child.PYTHON +RUST: Final = Child.RUST +LISTENING: Final = Probe.PYTHON_LISTENING TERM: Final = signal.SIGTERM INT: Final = signal.SIGINT -POSIX: Final = PosixPlatform() +HUP: Final = signal.SIGHUP +STOPS: Final = PosixPlatform().stop_signals + + +def _replay(events: tuple[Event, ...]) -> tuple[tuple[tuple[Action, ...], ...], int | None]: + transitions: Final = tuple( + itertools.accumulate(events, lambda transition, event: step(transition[0], event, STOPS), initial=boot()) + ) + return tuple(actions for _, actions in transitions[1:]), exit_status(transitions[-1][0]) + + +@pytest.mark.parametrize( + ("trace", "status"), + [ + pytest.param( + ( + (LISTENING, (Spawn(RUST),)), + (Caught(TERM), (Send(RUST, TERM), Send(PY, TERM))), + (ChildExited(RUST, 0), ()), + (ChildExited(PY, -15), ()), + ), + -15, + id="TERM reaches Rust first, then Python, and the supervisor dies as Python did", + ), + pytest.param( + ( + (LISTENING, (Spawn(RUST),)), + (Caught(INT), (Send(RUST, INT), Send(PY, INT))), + (Caught(INT), (Send(RUST, INT), Send(PY, INT))), + (ChildExited(PY, 1), (Send(RUST, TERM),)), + (ChildExited(RUST, 0), ()), + ), + 1, + id="a second Ctrl-C reaches Python again to force the exit", + ), + pytest.param( + ( + (LISTENING, (Spawn(RUST),)), + (Caught(HUP), (Send(PY, HUP),)), + (Caught(signal.SIGTTIN), (Send(PY, signal.SIGTTIN),)), + (ChildExited(PY, 0), (Send(RUST, TERM),)), + (ChildExited(RUST, 0), ()), + ), + 0, + id="gunicorn's HUP and TTIN reach only Python and Rust keeps serving", + ), + pytest.param( + ( + (LISTENING, (Spawn(RUST),)), + (ChildExited(RUST, 0), (Send(PY, TERM),)), + (ChildExited(PY, -15), ()), + ), + FAILURE, + id="losing Rust while Python serves stops Python and fails", + ), + pytest.param( + ( + (Caught(TERM), (Send(PY, TERM),)), + (LISTENING, ()), + (ChildExited(PY, -15), ()), + ), + -15, + id="a stop before Python listens never starts Rust", + ), + pytest.param(((ChildExited(PY, 3), ()),), 3, id="Python failing to boot exits with its status"), + ], +) +def test_supervisor_relays_signals_and_exits_as_python_did( + trace: tuple[tuple[Event, tuple[Action, ...]], ...], status: int +) -> None: + events: Final = tuple(event for event, _ in trace) + assert _replay(events) == (tuple(actions for _, actions in trace), status) @dataclass(frozen=True, slots=True) -class World: - """What the operating system knows about the children, kept independently of ``Lifecycle``.""" +class Observed: + """What an outside observer of the two processes knows, without looking inside ``Lifecycle``.""" - alive: frozenset[Child] = frozenset({Child.PYTHON}) - spawned: frozenset[Child] = frozenset({Child.PYTHON}) - reported_listening: bool = False + state: Lifecycle + alive: frozenset[Child] = frozenset({PY}) + rust_spawns: int = 0 stop_requested: bool = False - python_status: int | None = None rust_lost: bool = False + python_status: int | None = None -def _expected_exit(world: World) -> int | None: - if world.alive or world.python_status is None: - return None - return FAILURE if world.rust_lost else world.python_status +def _observe(seen: Observed, event: Event) -> Observed: + state, actions = step(seen.state, event, STOPS) + death: Final = event if isinstance(event, ChildExited) and event.child in seen.alive else None + exited: Final = None if death is None else death.child + running: Final = seen.alive - {exited} + assert all(action.child in running for action in actions if isinstance(action, Send)), "signalled a dead child" + if isinstance(event, Caught): + to_rust: Final = (Send(RUST, event.signum),) if event.signum in STOPS and RUST in running else () + assert actions == (*to_rust, *((Send(PY, event.signum),) if PY in running else ())) + spawned: Final = Spawn(RUST) in actions + assert not spawned or (event is LISTENING and not seen.stop_requested), "started Rust out of turn" + stop_requested: Final = seen.stop_requested or (isinstance(event, Caught) and event.signum in STOPS) + after: Final = replace( + seen, + state=state, + alive=running | ({RUST} if spawned else frozenset()), + rust_spawns=seen.rust_spawns + int(spawned), + stop_requested=stop_requested, + rust_lost=seen.rust_lost or (exited is RUST and PY in running and not stop_requested), + python_status=death.status if death is not None and exited is PY else seen.python_status, + ) + settled: Final = None if after.alive or after.python_status is None else after.python_status + assert exit_status(state) == (FAILURE if after.rust_lost and settled is not None else settled) + return after -def _expected_relay(world: World, signum: signal.Signals) -> tuple[Action, ...]: - to_rust: Final = (Send(Child.RUST, signum),) if signum in STOPS and Child.RUST in world.alive else () - to_python: Final = (Send(Child.PYTHON, signum),) if Child.PYTHON in world.alive else () - return (*to_rust, *to_python) - - -def _expected_on_exit(world: World, child: Child) -> tuple[Action, ...]: - match child: - case Child.PYTHON: - return (Send(Child.RUST, TERM),) if Child.RUST in world.alive else () - case Child.RUST: - serving: Final = Child.PYTHON in world.alive and not world.stop_requested - return (Send(Child.PYTHON, TERM),) if serving else () - - -def _after(world: World, event: Event) -> World: - match event: - case Caught(signum): - return replace(world, stop_requested=world.stop_requested or signum in STOPS) - case Probe.PYTHON_LISTENING: - return replace(world, reported_listening=world.reported_listening or Child.PYTHON in world.alive) - case ChildExited(child, status) if child in world.alive: - lost: Final = child is Child.RUST and Child.PYTHON in world.alive and not world.stop_requested - return replace( - world, - alive=world.alive - {child}, - python_status=status if child is Child.PYTHON else world.python_status, - rust_lost=world.rust_lost or lost, - ) - case ChildExited(): - return world - - -def _spawned(world: World, actions: tuple[Action, ...]) -> World: - started: Final = frozenset(action.child for action in actions if isinstance(action, Spawn)) - return replace(world, alive=world.alive | started, spawned=world.spawned | started) - - -SIGNALS: Final = POSIX.relayed_signals -STOPS: Final = POSIX.stop_signals EVENTS: Final = st.lists( st.one_of( - st.sampled_from(SIGNALS).map(Caught), - st.just(Probe.PYTHON_LISTENING), + st.sampled_from(PosixPlatform().relayed_signals).map(Caught), + st.just(LISTENING), st.builds(ChildExited, st.sampled_from(Child), st.integers(-15, 3)), ), max_size=40, @@ -99,69 +153,6 @@ EVENTS: Final = st.lists( @given(EVENTS) -@example([Caught(TERM), Caught(signal.SIGHUP), Probe.PYTHON_LISTENING]) -@example([Probe.PYTHON_LISTENING, Caught(TERM), Caught(signal.SIGWINCH), ChildExited(Child.RUST, 0)]) -def test_lifecycle_relays_every_signal_and_ends_with_pythons_status(events: list[Event]) -> None: - state, actions = boot() - world = World() - assert actions == (Spawn(Child.PYTHON),) - for event in events: - before = _after(world, event) - state, actions = step(state, event, STOPS) - match event: - case Caught(signum): - assert actions == _expected_relay(before, signum) - case Probe.PYTHON_LISTENING: - starts_rust = Child.PYTHON in world.alive and not world.reported_listening and not world.stop_requested - assert actions == ((Spawn(Child.RUST),) if starts_rust else ()) - case ChildExited(child) if child in world.alive: - assert actions == _expected_on_exit(world, child) - case ChildExited(): - assert actions == () - assert all(action.child in before.alive for action in actions if isinstance(action, Send)) - assert all(action.child not in before.spawned for action in actions if isinstance(action, Spawn)) - world = _spawned(before, actions) - assert exit_status(state) == _expected_exit(world) - - -@given(EVENTS) -def test_rust_starts_at_most_once_and_only_while_python_listens_unstopped(events: list[Event]) -> None: - def advance(acc: tuple[Lifecycle, tuple[tuple[Event, Lifecycle, Action], ...]], event: Event): - state, seen = acc - following, actions = step(state, event, STOPS) - return following, (*seen, *((event, state, action) for action in actions)) - - _, history = functools.reduce(advance, events, (boot()[0], ())) - spawns: Final = tuple((event, state) for event, state, action in history if isinstance(action, Spawn)) - assert len(spawns) <= 1 - assert all( - event is Probe.PYTHON_LISTENING and state.python is Phase.SPAWNED and not state.stop_requested - for event, state in spawns - ) - - -@pytest.mark.parametrize( - ("events", "expected"), - [ - ((Probe.PYTHON_LISTENING, Caught(TERM), ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, -15)), -15), - ( - ( - Probe.PYTHON_LISTENING, - Caught(INT), - Caught(INT), - ChildExited(Child.PYTHON, 1), - ChildExited(Child.RUST, 0), - ), - 1, - ), - ( - (Probe.PYTHON_LISTENING, Caught(signal.SIGHUP), ChildExited(Child.PYTHON, -1), ChildExited(Child.RUST, 0)), - -1, - ), - ((Probe.PYTHON_LISTENING, ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, 0)), FAILURE), - ((ChildExited(Child.PYTHON, 3),), 3), - ], -) -def test_supervisor_exits_as_python_did_unless_rust_was_lost(events: tuple[Event, ...], expected: int) -> None: - final: Final = functools.reduce(lambda state, event: step(state, event, STOPS)[0], events, boot()[0]) - assert exit_status(final) == expected +def test_any_event_order_keeps_the_supervisor_transparent_to_python(events: list[Event]) -> None: + final: Final = functools.reduce(_observe, events, Observed(boot()[0])) + assert final.rust_spawns <= 1 diff --git a/tests/unit/proxy/rust_forward/supervisor/test_preflight.py b/tests/unit/proxy/rust_forward/supervisor/test_preflight.py index 71899262c30..423ee4df20a 100644 --- a/tests/unit/proxy/rust_forward/supervisor/test_preflight.py +++ b/tests/unit/proxy/rust_forward/supervisor/test_preflight.py @@ -4,6 +4,7 @@ from typing import Final import pytest +from litellm.proxy.rust_forward.platform import Platform from litellm.proxy.rust_forward.platform.posix import PosixPlatform from litellm.proxy.rust_forward.platform.windows import WindowsPlatform from litellm.proxy.rust_forward.server import Listener, PythonServer, Server @@ -17,58 +18,47 @@ POSIX_ONLY: Final = pytest.mark.skipif(sys.platform == "win32", reason="runs a P @pytest.mark.parametrize( - ("server", "reason"), + ("platform", "server", "refusal"), [ - (python_server(Server.HYPERCORN), "LITELLM_RUST does not support Hypercorn"), - (python_server(Server.GRANIAN), "LITELLM_RUST does not support Granian"), - (python_server(listener=Listener.TLS), "LITELLM_RUST does not support TLS on the public port"), - ( - python_server(listener=Listener.UNIX_SOCKET), - "LITELLM_RUST does not support a Unix socket on the public port", - ), + (POSIX, python_server(), None), + (POSIX, python_server(Server.GUNICORN), None), + (WINDOWS, python_server(), None), + (WINDOWS, python_server(Server.GUNICORN), "does not support gunicorn on Windows"), + (POSIX, python_server(Server.HYPERCORN), "does not support Hypercorn"), + (POSIX, python_server(Server.GRANIAN), "does not support Granian"), + (POSIX, python_server(listener=Listener.TLS), "does not support TLS on the public port"), + (POSIX, python_server(listener=Listener.UNIX_SOCKET), "does not support a Unix socket on the public port"), ( + POSIX, python_server(Server.GUNICORN, Listener.INHERITED_FD), - "LITELLM_RUST does not support an inherited file descriptor on the public port", + "does not support an inherited file descriptor on the public port", ), ], ) -def test_unsupported_servers_are_refused_before_anything_starts( - server: PythonServer, reason: str, capsys: pytest.CaptureFixture[str] +def test_only_plain_tcp_on_a_supported_server_can_be_forwarded( + platform: Platform, server: PythonServer, refusal: str | None ) -> None: - assert startup_check(POSIX, server, "gateway") == StartupError(reason) - prepared: Final[list[bool]] = [] - refused: Final = supervise(server, ("python",), before_start=lambda: prepared.append(True), platform=POSIX) - assert refused == StartupError(reason) - assert (capsys.readouterr().err, prepared) == ( - f"{reason}, so Python serves the public port without Rust forwarding\n", - [], - ) + expected: Final = "gateway" if refusal is None else StartupError(f"LITELLM_RUST {refusal}") + assert startup_check(platform, server, "gateway") == expected -@pytest.mark.parametrize("server", [Server.UVICORN, Server.GUNICORN]) -def test_plain_tcp_uvicorn_and_gunicorn_can_be_forwarded(server: Server) -> None: - assert startup_check(POSIX, python_server(server), "gateway") == "gateway" - - -def test_windows_forwards_uvicorn_and_refuses_gunicorn() -> None: - assert startup_check(WINDOWS, python_server(), "gateway") == "gateway" - assert startup_check(WINDOWS, python_server(Server.GUNICORN), "gateway") == StartupError( - "LITELLM_RUST does not support gunicorn on Windows" - ) - - -def test_missing_binary_is_refused_before_anything_starts(tmp_path: Path) -> None: +@pytest.mark.parametrize( + ("server", "refusal"), + [ + (python_server(Server.HYPERCORN), "does not support Hypercorn"), + (python_server(), "requires the litellm-gateway binary from the litellm-proxy package"), + ], +) +def test_a_refused_start_says_why_in_one_line_and_starts_nothing( + server: PythonServer, refusal: str, tmp_path: Path, capsys: pytest.CaptureFixture[str] +) -> None: prepared: Final[list[bool]] = [] refused: Final = supervise( - python_server(), - ("python",), - before_start=lambda: prepared.append(True), - scripts_dir=str(tmp_path), - platform=POSIX, + server, ("python",), before_start=lambda: prepared.append(True), scripts_dir=str(tmp_path), platform=POSIX ) - assert (refused, prepared) == ( - StartupError("LITELLM_RUST requires the litellm-gateway binary from the litellm-proxy package"), - [], + assert (refused, prepared) == (StartupError(f"LITELLM_RUST {refusal}"), []) + assert ( + capsys.readouterr().err == f"LITELLM_RUST {refusal}, so Python serves the public port without Rust forwarding\n" ) diff --git a/tests/unit/proxy/rust_forward/supervisor/test_runtime.py b/tests/unit/proxy/rust_forward/supervisor/test_runtime.py index 824ef9e3945..e3cd55d1130 100644 --- a/tests/unit/proxy/rust_forward/supervisor/test_runtime.py +++ b/tests/unit/proxy/rust_forward/supervisor/test_runtime.py @@ -13,7 +13,6 @@ from litellm.proxy.rust_forward.platform.windows import WindowsPlatform from litellm.proxy.rust_forward.supervisor import supervise from litellm.proxy.rust_forward.supervisor.children import ChildSpec from litellm.proxy.rust_forward.supervisor.lifecycle import ( - FAILURE, Action, Caught, Child, @@ -65,6 +64,7 @@ def test_rust_opens_the_public_port_only_once_python_listens(tmp_path: Path, mon TERM: Final = signal.SIGTERM POSIX_ONLY: Final = pytest.mark.skipif(sys.platform == "win32", reason="the child traps SIGTERM") STOPS: Final = frozenset({signal.SIGTERM, signal.SIGINT, signal.SIGQUIT}) +RUST_MISSING: Final = SpawnFailed(FileNotFoundError("litellm-gateway")) @dataclass @@ -72,12 +72,12 @@ class ScriptedExecutor: """Replays a fixed event script and records every action the driver carries out.""" events: list[Event] - failing: frozenset[Child] = frozenset() + failing: frozenset[Child] done: list[Action] = field(default_factory=list) async def spawn(self, child: Child) -> SpawnFailed | None: self.done.append(Spawn(child)) - return SpawnFailed(FileNotFoundError(child.value)) if child in self.failing else None + return RUST_MISSING if child in self.failing else None def send(self, child: Child, signum: signal.Signals) -> None: self.done.append(Send(child, signum)) @@ -86,34 +86,41 @@ class ScriptedExecutor: return self.events.pop(0) -async def test_driver_carries_out_each_action_in_order_and_ends_with_pythons_status() -> None: - executor: Final = ScriptedExecutor( - [Probe.PYTHON_LISTENING, Caught(TERM), ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, -15)] - ) - assert await drive(executor, STOPS) == -15 - assert executor.done == [Spawn(Child.PYTHON), Spawn(Child.RUST), Send(Child.RUST, TERM), Send(Child.PYTHON, TERM)] - assert executor.events == [] - - -async def test_driver_stops_python_and_fails_when_rust_is_lost() -> None: - executor: Final = ScriptedExecutor( - [Probe.PYTHON_LISTENING, ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, -15)] - ) - assert await drive(executor, STOPS) == FAILURE - assert executor.done == [Spawn(Child.PYTHON), Spawn(Child.RUST), Send(Child.PYTHON, TERM)] - - -async def test_driver_asks_for_no_event_once_the_status_is_settled() -> None: - executor: Final = ScriptedExecutor([ChildExited(Child.PYTHON, 3), Caught(TERM)]) - assert await drive(executor, STOPS) == 3 - assert executor.events == [Caught(TERM)] - - -async def test_driver_returns_a_failed_spawn_without_waiting_for_events() -> None: - executor: Final = ScriptedExecutor([Probe.PYTHON_LISTENING, Caught(TERM)], failing=frozenset({Child.RUST})) - outcome: Final = await drive(executor, STOPS) - assert isinstance(outcome, SpawnFailed) - assert executor.events == [Caught(TERM)] +@pytest.mark.parametrize( + ("events", "failing", "done", "outcome", "unread"), + [ + pytest.param( + [Probe.PYTHON_LISTENING, Caught(TERM), ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, -15)], + frozenset(), + [Spawn(Child.PYTHON), Spawn(Child.RUST), Send(Child.RUST, TERM), Send(Child.PYTHON, TERM)], + -15, + [], + id="carries out each action in order and ends with Python's status", + ), + pytest.param( + [ChildExited(Child.PYTHON, 3), Caught(TERM)], + frozenset(), + [Spawn(Child.PYTHON)], + 3, + [Caught(TERM)], + id="reads no event once the status is settled", + ), + pytest.param( + [Probe.PYTHON_LISTENING, Caught(TERM)], + frozenset({Child.RUST}), + [Spawn(Child.PYTHON), Spawn(Child.RUST)], + RUST_MISSING, + [Caught(TERM)], + id="returns a failed spawn without waiting for more events", + ), + ], +) +async def test_driver_executes_lifecycle_decisions( + events: list[Event], failing: frozenset[Child], done: list[Action], outcome: int | SpawnFailed, unread: list[Event] +) -> None: + executor: Final = ScriptedExecutor(events, failing) + assert await drive(executor, STOPS) == outcome + assert (executor.done, executor.events) == (done, unread) def _children(argv: tuple[str, ...]) -> MappingProxyType[Child, ChildSpec]: diff --git a/tests/unit/proxy/rust_forward/test_mode.py b/tests/unit/proxy/rust_forward/test_mode.py index 5acd76860fb..ffecd1948ae 100644 --- a/tests/unit/proxy/rust_forward/test_mode.py +++ b/tests/unit/proxy/rust_forward/test_mode.py @@ -24,23 +24,14 @@ from litellm.proxy.rust_forward.mode import ( @pytest.mark.parametrize( - ("environ", "expected"), + ("value", "expected"), [ - ({}, Mode.DISABLED), - ({"LITELLM_RUST": "false"}, Mode.DISABLED), - ({"LITELLM_RUST": "off"}, Mode.DISABLED), - ({"LITELLM_RUST": "0"}, Mode.DISABLED), - ({"LITELLM_RUST": "true"}, Mode.SUPERVISOR), - ({"LITELLM_RUST": "1"}, Mode.SUPERVISOR), - ({"LITELLM_RUST": "yes"}, Mode.SUPERVISOR), - ({"LITELLM_RUST": "on"}, Mode.SUPERVISOR), - ({"LITELLM_RUST": " 1 "}, Mode.SUPERVISOR), - ({"LITELLM_RUST": ""}, Mode.DISABLED), - ({"LITELLM_RUST": " "}, Mode.DISABLED), + *((value, Mode.DISABLED) for value in (None, "", " ", "0", "false", "off")), + *((value, Mode.SUPERVISOR) for value in ("1", "true", "yes", "on", " 1 ")), ], ) -def test_mode_is_resolved_from_environment(environ: dict[str, str], expected: RustForwardMode) -> None: - assert rust_forward_mode(environ) == expected +def test_mode_is_resolved_from_environment(value: str | None, expected: RustForwardMode) -> None: + assert rust_forward_mode({} if value is None else {"LITELLM_RUST": value}) == expected @pytest.mark.parametrize( @@ -84,32 +75,34 @@ def _served_uvicorn_config(mode: RustForwardMode) -> uvicorn.Config: return uvicorn.Config(**options) -def test_supervised_uvicorn_serves_the_inherited_socket_without_trusting_the_rust_hop() -> None: - config: Final = _served_uvicorn_config(Supervised(InheritedFd(17), True, "*")) - assert (config.fd, config.proxy_headers) == (17, False) - assert (config.app, config.factory) == (WORKER_APP, True) +PUBLIC: Final = ("0.0.0.0", 45555, None) -@pytest.mark.parametrize("mode", list(Mode)) -def test_unsupervised_servers_keep_their_own_socket_and_proxy_headers(mode: RustForwardMode) -> None: +@pytest.mark.parametrize( + ("mode", "app", "uvicorn_listener", "gunicorn_listener", "proxy_headers"), + [ + (Mode.DISABLED, "litellm.proxy.proxy_server:app", PUBLIC, "0.0.0.0:4000", True), + (Mode.SUPERVISOR, "litellm.proxy.proxy_server:app", PUBLIC, "0.0.0.0:4000", True), + (Supervised(InheritedFd(17), True, "*"), WORKER_APP, ("0.0.0.0", 45555, 17), "fd://17", False), + (Supervised(AssignedPort(45556), True, "*"), WORKER_APP, ("127.0.0.1", 45556, None), "127.0.0.1:45556", False), + ], +) +def test_supervised_server_serves_the_private_hop_and_never_trusts_it( + mode: RustForwardMode, + app: str, + uvicorn_listener: tuple[str, int, int | None], + gunicorn_listener: str, + proxy_headers: bool, +) -> None: config: Final = _served_uvicorn_config(mode) - assert (config.fd, config.proxy_headers, config.port) == (None, True, 45555) - assert gunicorn_bind(mode, "0.0.0.0:4000") == "0.0.0.0:4000" - assert server_proxy_headers(mode) is True - - -def test_supervised_gunicorn_binds_the_inherited_socket_without_trusting_the_rust_hop() -> None: - mode: Final = Supervised(InheritedFd(17), True, "*") - assert gunicorn_bind(mode, "0.0.0.0:4000") == "fd://17" - assert server_proxy_headers(mode) is False - - -def test_assigned_port_serves_only_the_loopback_hop_without_trusting_it() -> None: - mode: Final = Supervised(AssignedPort(45556), True, "*") - config: Final = _served_uvicorn_config(mode) - assert (config.host, config.port, config.fd, config.proxy_headers) == ("127.0.0.1", 45556, None, False) - assert gunicorn_bind(mode, "0.0.0.0:4000") == "127.0.0.1:45556" - assert server_proxy_headers(mode) is False + assert (config.app, (config.host, config.port, config.fd), config.proxy_headers) == ( + app, + uvicorn_listener, + proxy_headers, + ) + assert config.factory is (app == WORKER_APP) + assert gunicorn_bind(mode, "0.0.0.0:4000") == gunicorn_listener + assert server_proxy_headers(mode) is proxy_headers @pytest.mark.parametrize( diff --git a/tests/unit/proxy/rust_forward/test_peer.py b/tests/unit/proxy/rust_forward/test_peer.py index 01331e2c69e..fdc9695d0d7 100644 --- a/tests/unit/proxy/rust_forward/test_peer.py +++ b/tests/unit/proxy/rust_forward/test_peer.py @@ -95,16 +95,6 @@ async def test_python_behind_rust_sees_what_a_direct_connection_would( assert behind_rust == direct -async def test_untrusted_peer_cannot_set_the_scheme_behind_rust() -> None: - client, scheme, _ = await _seen_by_app( - True, - ("127.0.0.1", 60000), - ((b"x-forwarded-proto", b"https"), (b"x-litellm-rust-peer", b"203.0.113.7:51234")), - Supervised(InheritedFd(17), True, "127.0.0.1"), - ) - assert (client, scheme) == (("203.0.113.7", 51234), "http") - - async def test_peer_header_from_anyone_but_the_rust_hop_is_dropped_and_ignored() -> None: client, _, headers = await _seen_by_app( True, diff --git a/tests/unit/proxy/test_proxy_cli.py b/tests/unit/proxy/test_proxy_cli.py index df7bbd43324..4627fb625e6 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -26,6 +26,7 @@ import urllib.parse as urlparse import uvicorn import yaml +from gunicorn.config import Config as GunicornConfig from uvicorn.config import LOOP_FACTORIES from uvicorn.importer import import_from_string @@ -34,7 +35,8 @@ 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.shutdown.graceful_shutdown_manager import DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT +from litellm.proxy.rust_forward.server import Listener, Server +from litellm.proxy.shutdown.graceful_shutdown_manager import DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT, GracefulShutdownManager @pytest.fixture(autouse=True) @@ -4301,22 +4303,20 @@ class TestRustForwardServer: ) @pytest.mark.parametrize( - ("argv", "reason"), + ("argv", "server", "listener"), [ - ( - ["--ssl_keyfile_path", "key.pem", "--ssl_certfile_path", "cert.pem"], - "does not support TLS on the public port", - ), - (["--run_hypercorn"], "does not support Hypercorn"), - (["--run_granian"], "does not support Granian"), + ([], Server.UVICORN, Listener.TCP), + (["--run_gunicorn"], Server.GUNICORN, Listener.TCP), + (["--run_hypercorn"], Server.HYPERCORN, Listener.TCP), + (["--run_granian"], Server.GRANIAN, Listener.TCP), + (["--ssl_keyfile_path", "key.pem", "--ssl_certfile_path", "cert.pem"], Server.UVICORN, Listener.TLS), ], ) - def test_unsupported_servers_are_refused_before_starting_children(self, argv: list[str], reason: str) -> None: - from litellm.proxy.rust_forward.platform.posix import PosixPlatform - from litellm.proxy.rust_forward.supervisor.preflight import StartupError, startup_check - - _, server = self._server(argv) - assert startup_check(PosixPlatform(), server, "gateway") == StartupError(f"LITELLM_RUST {reason}") + def test_the_supervisor_sees_the_server_and_listener_python_would_use( + self, argv: list[str], server: Server, listener: Listener + ) -> None: + _, described = self._server(argv) + assert (described.server, described.listener) == (server, listener) @patch("uvicorn.run") @patch("atexit.register") @@ -4383,59 +4383,37 @@ class TestRustForwardServer: assert (result.exit_code, len(lines)) == (1, 1) assert lines[0].startswith("Error: LITELLM_RUST is invalid: ") - @pytest.mark.parametrize( - ("argv", "environ"), - [(["--keepalive_timeout", "2"], {}), ([], {"KEEPALIVE_TIMEOUT": "7"}), ([], {})], - ) - def test_rust_retires_connections_by_the_keepalive_uvicorn_runs_with( - self, argv: list[str], environ: dict[str, str], monkeypatch: pytest.MonkeyPatch - ) -> None: - monkeypatch.delenv("KEEPALIVE_TIMEOUT", raising=False) - for name, value in environ.items(): - monkeypatch.setenv(name, value) - params, server = self._server(argv) + @staticmethod + def _served_by_python(params: dict[str, object]) -> tuple[object, ...]: + if params["run_gunicorn"]: + gunicorn: Final = GunicornConfig() + return gunicorn.keepalive, GracefulShutdownManager.get_timeout(), ",".join(gunicorn.forwarded_allow_ips) served: Final = uvicorn.Config( **ProxyInitializationHelpers.get_default_unvicorn_init_args( "127.0.0.1", 45555, keepalive_timeout=params["keepalive_timeout"] ) ) - assert server.keepalive_seconds == served.timeout_keep_alive + return served.timeout_keep_alive, served.timeout_graceful_shutdown, served.forwarded_allow_ips - def test_gunicorn_workers_keep_gunicorns_keepalive_and_rust_follows(self, monkeypatch: pytest.MonkeyPatch) -> None: - from gunicorn.config import Config as GunicornConfig - - monkeypatch.delenv("KEEPALIVE_TIMEOUT", raising=False) - _, server = self._server(["--run_gunicorn", "--keepalive_timeout", "9"]) - assert server.keepalive_seconds == GunicornConfig().keepalive - - @pytest.mark.parametrize("gunicorn", [True, False]) - def test_python_behind_rust_trusts_the_proxies_its_server_would_have( - self, gunicorn: bool, monkeypatch: pytest.MonkeyPatch + @pytest.mark.parametrize( + ("argv", "environ"), + [ + ([], {}), + (["--keepalive_timeout", "2"], {"GRACEFUL_SHUTDOWN_TIMEOUT": "12.5"}), + ([], {"KEEPALIVE_TIMEOUT": "7", "FORWARDED_ALLOW_IPS": "10.0.0.5"}), + (["--run_gunicorn", "--keepalive_timeout", "9"], {}), + ], + ) + def test_rust_follows_the_keepalive_drain_and_trust_python_serves_with( + self, argv: list[str], environ: dict[str, str], monkeypatch: pytest.MonkeyPatch ) -> None: - from gunicorn.config import Config as GunicornConfig - - from litellm.proxy.rust_forward.mode import InheritedFd, Supervised, rust_forward_mode - from litellm.proxy.rust_forward.platform import PrivateListener - from litellm.proxy.rust_forward.platform.posix import PosixPlatform - from litellm.proxy.rust_forward.supervisor.children import python_spec - - monkeypatch.delenv("FORWARDED_ALLOW_IPS", raising=False) - _, server = self._server(["--run_gunicorn"] if gunicorn else []) - mode: Final = rust_forward_mode( - python_spec( - PosixPlatform(), ("litellm",), PrivateListener(InheritedFd(17), 45555, (17,)), server - ).env_overrides - ) - trusted: Final = ( - GunicornConfig().forwarded_allow_ips - if gunicorn - else [uvicorn.Config(app="", log_config=None).forwarded_allow_ips] - ) - assert isinstance(mode, Supervised) - assert mode.trusted_proxies.split(",") == trusted - - def test_rust_drains_for_the_graceful_timeout_uvicorn_runs_with(self, monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "12.5") - _, server = self._server([]) - served: Final = ProxyInitializationHelpers.get_default_unvicorn_init_args("127.0.0.1", 45555) - assert server.graceful_timeout_seconds == served["timeout_graceful_shutdown"] + for name in ("KEEPALIVE_TIMEOUT", "GRACEFUL_SHUTDOWN_TIMEOUT", "FORWARDED_ALLOW_IPS"): + monkeypatch.delenv(name, raising=False) + for name, value in environ.items(): + monkeypatch.setenv(name, value) + params, server = self._server(argv) + assert ( + server.keepalive_seconds, + server.graceful_timeout_seconds, + server.trusted_proxies, + ) == self._served_by_python(params) From 9fd0d7c0c2a0120ede26fce4f5cbb4a97d7d3a10 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Fri, 6 Nov 2026 15:39:38 +0900 Subject: [PATCH 10/10] test(proxy): launch Rust forwarding integration proxies through one helper Both entry points now start through tests/integration/_support/rust_forward.py, which owns the port, config file, environment, readiness wait and the exit 0 check. The client IP test reads as a table of callers and no longer needs the shared rig just to borrow a master key. The lifecycle runner is generic, so the kill test returns its survivors instead of a placeholder Run Co-Authored-By: Claude Opus 5.5 --- tests/integration/_support/rust_forward.py | 117 +++++++++ .../configuration/test_rust_forward.py | 229 ++++++------------ .../test_rust_forward_lifecycle.py | 16 +- .../test_rust_forward_traffic.py | 108 ++------- 4 files changed, 214 insertions(+), 256 deletions(-) create mode 100644 tests/integration/_support/rust_forward.py diff --git a/tests/integration/_support/rust_forward.py b/tests/integration/_support/rust_forward.py new file mode 100644 index 00000000000..e85c5b67091 --- /dev/null +++ b/tests/integration/_support/rust_forward.py @@ -0,0 +1,117 @@ +import json +import os +import subprocess +import sys +from collections.abc import Iterator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass +from enum import Enum +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import httpx +from pydantic import JsonValue + +from tests.integration._support.client import eventually +from tests.integration._support.ports import allocate_port +from tests.integration._support.process import stop_root_process + +MASTER_KEY: Final = "sk-rust-forward-integration-key-not-a-secret" +WORKERS: Final = "2" +HEALTHCHECK_SECONDS: Final = "120" + + +class EntryPoint(Enum): + LITELLM = "litellm.proxy.proxy_cli" + GATEWAY = "gateway.launch" + + +@dataclass(frozen=True, slots=True) +class ForwardedProxy: + port: int + log: Path + + @property + def url(self) -> str: + return f"http://127.0.0.1:{self.port}" + + +def _serving(url: str) -> bool: + try: + return httpx.get(f"{url}/health/liveliness", timeout=5, trust_env=False).status_code == 200 + except httpx.TransportError: + return False + + +def _launch(entry_point: EntryPoint, config: Path) -> tuple[tuple[str, ...], Mapping[str, str]]: + match entry_point: + case EntryPoint.LITELLM: + return ( + ( + "--config", + str(config), + "--num_workers", + WORKERS, + "--timeout_worker_healthcheck", + HEALTHCHECK_SECONDS, + ), + MappingProxyType({}), + ) + case EntryPoint.GATEWAY: + return ( + ("--workers", WORKERS, "--timeout-worker-healthcheck", HEALTHCHECK_SECONDS), + MappingProxyType({"CONFIG_FILE_PATH": str(config), "LITELLM_PGBOUNCER_ENABLED": "false"}), + ) + + +@contextmanager +def forwarded_proxy( + entry_point: EntryPoint, + config: Mapping[str, JsonValue], + directory: Path, + arguments: tuple[str, ...] = (), + environment: Mapping[str, str] = MappingProxyType({}), +) -> Iterator[ForwardedProxy]: + """``entry_point`` on 0.0.0.0 with two workers, ``LITELLM_RUST=1`` and no database. It must exit 0 on SIGTERM.""" + proxy: Final = ForwardedProxy(allocate_port(), directory / f"{entry_point.name.lower()}.log") + config_file: Final = directory / "config.yaml" + config_file.write_text(json.dumps(config)) + launch_arguments, launch_environment = _launch(entry_point, config_file) + inherited: Final = { + name: value for name, value in os.environ.items() if not name.startswith("DATABASE_") and name != "DIRECT_URL" + } + with ( + proxy.log.open("wb") as output, + subprocess.Popen( + ( + sys.executable, + "-I", + "-m", + entry_point.value, + "--host", + "0.0.0.0", + "--port", + str(proxy.port), + *launch_arguments, + *arguments, + ), + env={ + **inherited, + "LITELLM_MASTER_KEY": MASTER_KEY, + "LITELLM_RUST": "true", + "STORE_MODEL_IN_DB": "False", + "PYTHON_DOTENV_DISABLED": "1", + **launch_environment, + **environment, + }, + stdout=output, + stderr=output, + ) as process, + ): + try: + assert eventually(lambda: _serving(proxy.url), bool, seconds=180), proxy.log.read_text() + yield proxy + finally: + assert stop_root_process(process), proxy.log.read_text() + assert process.returncode == 0, proxy.log.read_text() diff --git a/tests/integration/configuration/test_rust_forward.py b/tests/integration/configuration/test_rust_forward.py index b75da286bf7..009b3e423ed 100644 --- a/tests/integration/configuration/test_rust_forward.py +++ b/tests/integration/configuration/test_rust_forward.py @@ -1,180 +1,91 @@ -import json -import os -import subprocess -import sys +from dataclasses import dataclass from pathlib import Path from typing import Final import httpx import pytest -from tests.integration._support.client import Gateway, eventually +from tests.integration._support.client import eventually from tests.integration._support.ports import allocate_port -from tests.integration._support.process import stop_root_process +from tests.integration._support.rust_forward import MASTER_KEY, EntryPoint, forwarded_proxy TRUSTED_BALANCER: Final = "127.0.0.3" +CLIENT: Final = "127.0.0.2" -def ready(client: httpx.Client) -> bool: - try: - return client.get("/health/liveliness").status_code == 200 - except httpx.TransportError: - return False +@dataclass(frozen=True, slots=True) +class Caller: + connects_from: str + forwarded_for: str + resolved_ip: str -def client_ip(base_url: str, local_address: str, key: str, forwarded_for: str) -> httpx.Response: - with httpx.Client( - base_url=base_url, - transport=httpx.HTTPTransport(local_address=local_address), - trust_env=False, - timeout=5, - ) as client: - return client.get( +CALLERS: Final = ( + Caller(CLIENT, "1.2.3.4", resolved_ip=CLIENT), + Caller(TRUSTED_BALANCER, "203.0.113.9", resolved_ip="203.0.113.9"), + Caller(CLIENT, f"203.0.113.9, {TRUSTED_BALANCER}", resolved_ip=CLIENT), +) + + +def _client(url: str, local_address: str) -> httpx.Client: + return httpx.Client( + base_url=url, transport=httpx.HTTPTransport(local_address=local_address), trust_env=False, timeout=5 + ) + + +def _resolved_ip(url: str, caller: Caller) -> str: + with _client(url, caller.connects_from) as client: + response: Final = client.get( "/v1/mcp/network/client-ip", - headers={"Authorization": f"Bearer {key}", "X-Forwarded-For": forwarded_for}, + headers={"Authorization": f"Bearer {MASTER_KEY}", "X-Forwarded-For": caller.forwarded_for}, ) + assert response.status_code == 200, response.text + return response.json()["ip"] @pytest.mark.parametrize("server_arguments", [(), ("--run_gunicorn",)]) -def test_gateway_preserves_metrics_bind_and_mcp_client_ip( - gateway: Gateway, tmp_path: Path, server_arguments: tuple[str, ...] -) -> None: - config: Final = tmp_path / "config.yaml" - config.write_text( - json.dumps( - { - "model_list": [], - "general_settings": {"master_key": "os.environ/LITELLM_MASTER_KEY", "use_x_forwarded_for": False}, - } - ) - ) - metrics_port: Final = allocate_port() - public_port: Final = allocate_port() - environment: Final = { - **{key: value for key, value in os.environ.items() if key not in ("DATABASE_URL", "DATABASE_URL_READ_REPLICA")}, - "LITELLM_MASTER_KEY": gateway.key, - "STORE_MODEL_IN_DB": "False", - "PYTHON_DOTENV_DISABLED": "1", - "LITELLM_RUST": "true", - "FORWARDED_ALLOW_IPS": TRUSTED_BALANCER, - } - with (tmp_path / "proxy.log").open("wb") as output: - with subprocess.Popen( - ( - sys.executable, - "-I", - "-m", - "litellm.proxy.proxy_cli", - "--config", - str(config), - "--host", - "0.0.0.0", - "--port", - str(public_port), - "--num_workers", - "2", - "--timeout_worker_healthcheck", - "120", - "--prometheus_metrics_port", - str(metrics_port), - *server_arguments, - ), - env=environment, - stdout=output, - stderr=output, - ) as supervisor: - try: - with httpx.Client( - base_url=f"http://127.0.0.1:{public_port}", - transport=httpx.HTTPTransport(local_address="127.0.0.2"), - trust_env=False, - follow_redirects=True, - timeout=5, - ) as external: - assert eventually(lambda: ready(external), bool, seconds=180), (tmp_path / "proxy.log").read_text() - response: Final = external.get(f"http://127.0.0.2:{metrics_port}/metrics") - assert response.status_code == 200, response.text - public: Final = f"http://127.0.0.1:{public_port}" - spoofed: Final = client_ip(public, "127.0.0.2", gateway.key, "1.2.3.4") - assert spoofed.status_code == 200, spoofed.text - assert spoofed.json() == {"ip": "127.0.0.2"} - behind_balancer: Final = client_ip(public, TRUSTED_BALANCER, gateway.key, "203.0.113.9") - assert behind_balancer.status_code == 200, behind_balancer.text - assert behind_balancer.json() == {"ip": "203.0.113.9"} - borrowed: Final = client_ip(public, "127.0.0.2", gateway.key, f"203.0.113.9, {TRUSTED_BALANCER}") - assert borrowed.status_code == 200, borrowed.text - assert borrowed.json() == {"ip": "127.0.0.2"} - finally: - assert stop_root_process(supervisor), (tmp_path / "proxy.log").read_text() - assert supervisor.returncode == 0, (tmp_path / "proxy.log").read_text() - - -@pytest.mark.parametrize( - "server_arguments", - [("--no-proxy-headers",), ("--forwarded-allow-ips", "127.0.0.1")], -) -def test_gateway_launcher_private_hop_preserves_client_address( +def test_only_the_trusted_balancer_can_name_the_client_ip_behind_rust( tmp_path: Path, server_arguments: tuple[str, ...] ) -> None: - config: Final = tmp_path / "config.yaml" - config.write_text( - json.dumps({"model_list": [], "general_settings": {"master_key": "os.environ/LITELLM_MASTER_KEY"}}) - ) - public_port: Final = allocate_port() - environment: Final = { - **{key: value for key, value in os.environ.items() if not key.startswith("DATABASE_") and key != "DIRECT_URL"}, - "LITELLM_MASTER_KEY": "sk-rust-forwarding-integration-key-not-a-secret", - "LITELLM_RUST": "true", - "LITELLM_PGBOUNCER_ENABLED": "false", - "CONFIG_FILE_PATH": str(config), - "STORE_MODEL_IN_DB": "False", - "PYTHON_DOTENV_DISABLED": "1", - } - log: Final = tmp_path / "gateway.log" - with log.open("wb") as output: - with subprocess.Popen( - ( - sys.executable, - "-I", - "-m", - "gateway.launch", - "--host", - "0.0.0.0", - "--port", - str(public_port), - "--workers", - "2", - "--timeout-worker-healthcheck", - "120", - *server_arguments, + metrics_port: Final = allocate_port() + with forwarded_proxy( + EntryPoint.LITELLM, + { + "model_list": [], + "general_settings": {"master_key": "os.environ/LITELLM_MASTER_KEY", "use_x_forwarded_for": False}, + }, + tmp_path, + ("--prometheus_metrics_port", str(metrics_port), *server_arguments), + {"FORWARDED_ALLOW_IPS": TRUSTED_BALANCER}, + ) as proxy: + with _client(f"http://{CLIENT}:{metrics_port}", CLIENT) as external: + metrics: Final = external.get("/metrics", follow_redirects=True) + assert metrics.status_code == 200, metrics.text + assert tuple(_resolved_ip(proxy.url, caller) for caller in CALLERS) == tuple( + caller.resolved_ip for caller in CALLERS + ) + + +@pytest.mark.parametrize("server_arguments", [("--no-proxy-headers",), ("--forwarded-allow-ips", "127.0.0.1")]) +def test_gateway_launcher_logs_the_real_client_address_behind_rust( + tmp_path: Path, server_arguments: tuple[str, ...] +) -> None: + with forwarded_proxy( + EntryPoint.GATEWAY, + {"model_list": [], "general_settings": {"master_key": "os.environ/LITELLM_MASTER_KEY"}}, + tmp_path, + server_arguments, + ) as proxy: + with _client(proxy.url, CLIENT) as external: + response: Final = external.get( + "/health/liveliness?client-ip-proof", headers={"X-Forwarded-For": "198.51.100.10"} + ) + assert response.status_code == 200, response.text + logged: Final = eventually( + lambda: tuple( + line for line in proxy.log.read_text().splitlines() if "GET /health/liveliness?client-ip-proof " in line ), - env=environment, - stdout=output, - stderr=output, - ) as supervisor: - try: - with httpx.Client( - base_url=f"http://127.0.0.1:{public_port}", - transport=httpx.HTTPTransport(local_address="127.0.0.2"), - trust_env=False, - timeout=5, - ) as external: - assert eventually(lambda: ready(external), bool, seconds=180), log.read_text() - response: Final = external.get( - "/health/liveliness?client-ip-proof", headers={"X-Forwarded-For": "198.51.100.10"} - ) - assert response.status_code == 200, response.text - events: Final = eventually( - lambda: tuple( - line - for line in log.read_text().splitlines() - if "GET /health/liveliness?client-ip-proof " in line - ), - bool, - seconds=10, - ) - assert len(events) == 1, log.read_text() - assert "127.0.0.2:" in events[0], events - finally: - assert stop_root_process(supervisor), log.read_text() - assert supervisor.returncode == 0, log.read_text() + bool, + seconds=10, + ) + assert len(logged) == 1 and f"{CLIENT}:" in logged[0], logged diff --git a/tests/integration/configuration/test_rust_forward_lifecycle.py b/tests/integration/configuration/test_rust_forward_lifecycle.py index 7abafc6c6cd..90e2e7567e5 100644 --- a/tests/integration/configuration/test_rust_forward_lifecycle.py +++ b/tests/integration/configuration/test_rust_forward_lifecycle.py @@ -8,7 +8,7 @@ from collections.abc import Callable from dataclasses import dataclass from enum import Enum from pathlib import Path -from typing import Final +from typing import Final, TypeVar import httpx import psutil @@ -128,6 +128,7 @@ result = supervise(server, (sys.executable, str(scripts / "serve.py"), workers, raise SystemExit(result if isinstance(result, int) else 3) """ +T = TypeVar("T") EXIT_SECONDS: Final = 15 TIME_TOLERANCE_SECONDS: Final = 5 @@ -206,7 +207,7 @@ def _signal_after_ready(process: subprocess.Popen[bytes], port: int, scenario: S return Run(Outcome(_exit_status(process, EXIT_SECONDS), True), time.monotonic() - started) -def _run(argv: tuple[str, ...], environment: dict[str, str], act: Callable[[subprocess.Popen[bytes]], Run]) -> Run: +def _run(argv: tuple[str, ...], environment: dict[str, str], act: Callable[[subprocess.Popen[bytes]], T]) -> T: with subprocess.Popen(argv, env=environment, start_new_session=True) as process: try: return act(process) @@ -252,18 +253,17 @@ def test_children_exit_when_the_supervisor_is_killed(tmp_path: Path, workers: in scripts: Final = _scripts(tmp_path) port: Final = allocate_port() - def kill_supervisor(process: subprocess.Popen[bytes]) -> Run: + def survivors_of_killing_the_supervisor(process: subprocess.Popen[bytes]) -> tuple[int, ...]: assert eventually(lambda: _serving(port), bool, seconds=120) python: Final = int((tmp_path / "python.pid").read_text()) rust: Final = int((tmp_path / "rust.pid").read_text()) descendants: Final = (python, rust, *(child.pid for child in psutil.Process(python).children(recursive=True))) - started: Final = time.monotonic() kill_stray_process(process.pid) - assert eventually(lambda: tuple(pid for pid in descendants if _alive(pid)), lambda alive: not alive) == () - return Run(Outcome(None, False), time.monotonic() - started) + return eventually(lambda: tuple(pid for pid in descendants if _alive(pid)), lambda alive: not alive) - _run( + survivors: Final = _run( (sys.executable, str(scripts / "supervise.py"), str(workers), str(port), str(scripts)), _environment(tmp_path), - kill_supervisor, + survivors_of_killing_the_supervisor, ) + assert survivors == () diff --git a/tests/integration/configuration/test_rust_forward_traffic.py b/tests/integration/configuration/test_rust_forward_traffic.py index eeeeea366d8..4cd2f09e6d7 100644 --- a/tests/integration/configuration/test_rust_forward_traffic.py +++ b/tests/integration/configuration/test_rust_forward_traffic.py @@ -2,14 +2,11 @@ import asyncio import json import os import socket -import subprocess -import sys import threading import time from collections.abc import AsyncIterator, Callable, Iterator, Mapping from contextlib import AbstractContextManager, ExitStack from dataclasses import dataclass -from pathlib import Path from queue import SimpleQueue from types import MappingProxyType from typing import Final @@ -27,13 +24,10 @@ from starlette.websockets import WebSocket, WebSocketDisconnect from websockets.sync.client import connect from tests.integration._support.asgi import asgi_server -from tests.integration._support.client import eventually from tests.integration._support.mcp import mcp_peer -from tests.integration._support.ports import allocate_port -from tests.integration._support.process import stop_root_process +from tests.integration._support.rust_forward import MASTER_KEY, EntryPoint, forwarded_proxy from tests.integration._support.wire import Reply, Request, wire_server -MASTER_KEY: Final = "sk-rust-forward-traffic-key-not-a-secret" KEEPALIVE_SECONDS: Final = 1 FILES_MODEL: Final = "rust-forward-files" FILE_ID: Final = "file-rust-forward" @@ -59,7 +53,6 @@ class ScriptedProvider: class RustProxy: url: str port: int - log: Path provider: ScriptedProvider @@ -151,13 +144,6 @@ def _routable_address() -> str: return address -def _ready(client: httpx.Client) -> bool: - try: - return client.get("/health/liveliness").status_code == 200 - except httpx.TransportError: - return False - - def _files_peer(content: bytes) -> Callable[[Request], Reply]: file_object: Final = json.dumps( { @@ -205,9 +191,6 @@ def files_upstream() -> Iterator[str]: @pytest.fixture(scope="module") def rust_proxy(tmp_path_factory: pytest.TempPathFactory, files_upstream: str) -> Iterator[RustProxy]: - directory: Final = tmp_path_factory.mktemp("rust-forward-traffic") - log: Final = directory / "proxy.log" - port: Final = allocate_port() released: Final = threading.Event() hung_up: Final[SimpleQueue[str]] = SimpleQueue() with ExitStack() as stack: @@ -215,90 +198,37 @@ def rust_proxy(tmp_path_factory: pytest.TempPathFactory, files_upstream: str) -> provider: Final = ScriptedProvider( stack.enter_context(asgi_server(_provider_app(released, hung_up))), released, hung_up ) - config: Final = directory / "config.yaml" - config.write_text( - json.dumps( + proxy: Final = stack.enter_context( + forwarded_proxy( + EntryPoint.LITELLM, { "model_list": [ + _deployment(CHAT_MODEL, "openai/gpt-4o-mini", f"{provider.url}/v1"), { - "model_name": CHAT_MODEL, - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_base": f"{provider.url}/v1", - "api_key": "integration-provider-key", - }, - }, - { - "model_name": REALTIME_MODEL, - "litellm_params": { - "model": "openai/gpt-realtime", - "api_base": provider.url, - "api_key": "integration-provider-key", - }, + **_deployment(REALTIME_MODEL, "openai/gpt-realtime", provider.url), "model_info": {"mode": "realtime"}, }, - { - "model_name": FILES_MODEL, - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_base": f"{files_upstream}/v1", - "api_key": "integration-provider-key", - }, - }, + _deployment(FILES_MODEL, "openai/gpt-4o-mini", f"{files_upstream}/v1"), ], "mcp_servers": {"math": {**mcp.registration(), "available_on_public_internet": False}}, "general_settings": { "master_key": "os.environ/LITELLM_MASTER_KEY", "mcp_internal_ip_ranges": ["127.0.0.0/8"], }, - } + }, + tmp_path_factory.mktemp("rust-forward-traffic"), + ("--keepalive_timeout", str(KEEPALIVE_SECONDS)), ) ) - environment: Final = { - **{ - key: value - for key, value in os.environ.items() - if not key.startswith("DATABASE_") and key != "DIRECT_URL" - }, - "LITELLM_MASTER_KEY": MASTER_KEY, - "LITELLM_RUST": "true", - "STORE_MODEL_IN_DB": "False", - "PYTHON_DOTENV_DISABLED": "1", - } - output: Final = stack.enter_context(log.open("wb")) - supervisor: Final = stack.enter_context( - subprocess.Popen( - ( - sys.executable, - "-I", - "-m", - "litellm.proxy.proxy_cli", - "--config", - str(config), - "--host", - "0.0.0.0", - "--port", - str(port), - "--num_workers", - "2", - "--keepalive_timeout", - str(KEEPALIVE_SECONDS), - "--timeout_worker_healthcheck", - "120", - ), - env=environment, - stdout=output, - stderr=output, - ) - ) - try: - with httpx.Client(base_url=f"http://127.0.0.1:{port}", trust_env=False, timeout=5) as probe: - assert eventually(lambda: _ready(probe), bool, seconds=180), log.read_text() - yield RustProxy(f"http://127.0.0.1:{port}", port, log, provider) - finally: - released.set() - assert stop_root_process(supervisor), log.read_text() - assert supervisor.returncode == 0, log.read_text() + stack.callback(released.set) + yield RustProxy(proxy.url, proxy.port, provider) + + +def _deployment(name: str, model: str, api_base: str) -> dict[str, JsonValue]: + return { + "model_name": name, + "litellm_params": {"model": model, "api_base": api_base, "api_key": "integration-provider-key"}, + } def _still_open(stream: httpx.Response) -> bool: