diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index f8b93b5cc11..d0726e0daa0 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -591,6 +591,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/AGENTS.md b/gateway/AGENTS.md new file mode 100644 index 00000000000..1bb4013179f --- /dev/null +++ b/gateway/AGENTS.md @@ -0,0 +1,16 @@ +# Gateway image entry point + +`python -m gateway.launch` is how the gateway image starts. `launch.py` hands its arguments to uvicorn's own CLI and serves `main.py`, which trims the proxy app to the LLM data plane + +- Importing `launch.py` stays cheap. Every uvicorn worker imports it again before it can answer its first health ping, so LiteLLM imports wait until `main` or `worker_app` runs +- PgBouncer starts exactly once per pod, before the workers start, and the workers inherit the loopback URL through `DATABASE_URL`. `configure_database` is the only place that does this + +## Rust forwarding + +`launch.py` is an entry point of Rust forwarding, and every forwarding decision lives in `litellm/proxy/rust_forward/`, whose `AGENTS.md` states the contract + +- `main` dispatches on `rust_forward_mode` and does nothing else about forwarding. With forwarding off or refused it serves exactly as before +- `python_server` reads the flags and environment through uvicorn's own CLI parser, so Rust gets the values uvicorn would run with. Never reparse a uvicorn flag by hand +- Under forwarding the supervisor process runs `configure_database`, as `before_start`, once forwarding is known to be possible. The supervised child never runs it and inherits the exported URL +- The supervised child appends `mode.uvicorn_flags` after the operator's flags, so the private listener and `--no-proxy-headers` win +- `worker_app` wraps the app with `peer.wrap`, which is the identity when forwarding is off diff --git a/gateway/launch.py b/gateway/launch.py index b9303fd9447..c8b5b53af92 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,16 +66,39 @@ 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 python_server(argv: Sequence[str]) -> "PythonServer": + """What uvicorn would serve the public port with, read from the same flags and environment it reads.""" + from litellm.proxy.rust_forward.server import PythonServer, Server, UvicornCliParams + from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager + + params: Final = UvicornCliParams.model_validate( + uvicorn_main.make_context("uvicorn", list(uvicorn_argv(argv, os.environ))).params + ) + return PythonServer( + server=Server.UVICORN, + listener=params.listener, + host=params.host, + port=params.port, + keepalive_seconds=params.timeout_keep_alive, + graceful_timeout_seconds=GracefulShutdownManager.get_timeout() + if params.timeout_graceful_shutdown is None + else params.timeout_graceful_shutdown, + trusted_proxies=params.trusted_proxies, + ) + + 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 @@ -88,10 +112,57 @@ def configure_database() -> None: export_pooled_database_url(pooled_url) -def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -> None: +def _serve_python_alone(argv: Sequence[str], serve: Callable[[Sequence[str]], None]) -> None: configure_database() serve(uvicorn_argv(argv, os.environ)) +def _python_server_or_exit(argv: Sequence[str]) -> "PythonServer": + """Report ``--help``, ``--version`` and a bad flag the way uvicorn's own CLI does, since it never gets to.""" + import click + + try: + return python_server(argv) + except click.exceptions.Exit as stopped: + sys.exit(stopped.exit_code) + except click.ClickException as error: + error.show() + sys.exit(error.exit_code) + + +def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -> None: + from pydantic import ValidationError + + from litellm.proxy.rust_forward.mode import ( + Mode, + Supervised, + environment_error, + rust_forward_mode, + stop_with_supervisor, + uvicorn_flags, + ) + from litellm.proxy.rust_forward.supervisor import supervise + from litellm.proxy.rust_forward.supervisor.preflight import StartupError + + try: + mode: Final = rust_forward_mode(os.environ) + except ValidationError as error: + sys.exit(environment_error(error)) + match mode: + case Mode.DISABLED: + _serve_python_alone(argv, serve) + case Mode.SUPERVISOR: + supervised: Final = supervise(_python_server_or_exit(argv), before_start=configure_database) + if isinstance(supervised, StartupError): + _serve_python_alone(argv, serve) + else: + sys.exit(supervised) + case Supervised(): + stop_with_supervisor(mode) + serve((*uvicorn_argv(argv, os.environ), *uvicorn_flags(mode))) + case _: + assert_never(mode) + + if __name__ == "__main__": main(sys.argv[1:]) diff --git a/litellm-rust/crates/gateway-forward/AGENTS.md b/litellm-rust/crates/gateway-forward/AGENTS.md index 2135980266f..3d07f9ee342 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 @@ -22,6 +23,8 @@ Every forwarding decision has one home. `rewrite.rs` owns what Python and the cl `Settings` reads only these variables, and only when `LITELLM_FORWARD_UPSTREAM` is set. The Python supervisor in `litellm/proxy/rust_forward/` sets every one of them from the values Python actually runs with, so all of them are required and a missing or invalid one fails startup naming it. The crate never reads Python's own settings (`HOST`, `PORT`, `KEEPALIVE_TIMEOUT`, `GRACEFUL_SHUTDOWN_TIMEOUT`) and keeps no copy of Python's defaults, since Python can get those values from CLI flags the environment never shows +Only the Python supervisor starts the gateway in forwarding mode, from `litellm` and from `python -m gateway.launch`, on a laptop and in every image. `supervisor/children.py` is the only writer of these variables, so adding, renaming or removing one changes `settings.rs` and `children.rust_spec` in the same change. Without `LITELLM_FORWARD_UPSTREAM` the binary runs as the standalone Rust gateway and this crate does nothing + | Variable | Meaning | |---|---| | `LITELLM_FORWARD_UPSTREAM` | Python's loopback origin, e.g. `http://127.0.0.1:53211` | 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/input.rs b/litellm-rust/crates/gateway-forward/src/input.rs new file mode 100644 index 00000000000..772d4610683 --- /dev/null +++ b/litellm-rust/crates/gateway-forward/src/input.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("stdin-eof".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/src/lib.rs b/litellm-rust/crates/gateway-forward/src/lib.rs index 4ac17b50544..35f25d19c3c 100644 --- a/litellm-rust/crates/gateway-forward/src/lib.rs +++ b/litellm-rust/crates/gateway-forward/src/lib.rs @@ -2,6 +2,7 @@ mod connections; mod connector; mod error; mod forward; +mod input; mod rewrite; mod server; mod settings; @@ -10,6 +11,7 @@ mod websocket; pub use error::Error; pub(crate) use error::Rejection; pub use forward::Forward; +pub use input::input_closed; pub use rewrite::PEER_HEADER; pub use settings::{Settings, Upstream}; 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 dc77cfe28eb..176bc68d247 100644 --- a/litellm-rust/crates/gateway/src/main.rs +++ b/litellm-rust/crates/gateway/src/main.rs @@ -48,7 +48,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/AGENTS.md b/litellm/proxy/AGENTS.md new file mode 100644 index 00000000000..62726c30a10 --- /dev/null +++ b/litellm/proxy/AGENTS.md @@ -0,0 +1,5 @@ +# Proxy + +## Rust forwarding + +`proxy_server.py`, routes, auth and middleware never know whether Rust is in front of the proxy, and none of them may branch on it. `proxy_cli.py` only describes what it would serve with as a `rust_forward.server.PythonServer`, calls `rust_forward.supervisor.supervise`, and merges the options `rust_forward.mode` returns, which are identities when forwarding is off. Every forwarding decision lives in `rust_forward/`, and `rust_forward/AGENTS.md` states the contract diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index b238cefd24f..f252e69d739 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -36,6 +36,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 @@ -680,7 +683,7 @@ 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, @@ -836,6 +839,65 @@ class ProxyInitializationHelpers: 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 litellm.proxy.rust_forward.server import Listener, PythonServer, Server, uvicorn_trusted_proxies + 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, trusted_proxies = ( + ProxyInitializationHelpers.gunicorn_keepalive_and_trusted_proxies() + if server is Server.GUNICORN + else ( + uvicorn.Config(app="", log_config=None).timeout_keep_alive + if keepalive_timeout is None + else keepalive_timeout, + uvicorn_trusted_proxies(None), + ) + ) + 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(), + trusted_proxies=trusted_proxies, + ) + + @staticmethod + def gunicorn_keepalive_and_trusted_proxies() -> tuple[int, str]: + """Loaded only when gunicorn is the server, because gunicorn imports ``fcntl``, which Windows lacks.""" + from gunicorn.config import Config as GunicornConfig + + config: Final = GunicornConfig() + keepalive: Final[object] = config.keepalive + forwarded_allow_ips: Final[object] = config.forwarded_allow_ips + return ( + TypeAdapter(int).validate_python(keepalive), + ",".join(TypeAdapter(list[str]).validate_python(forwarded_allow_ips)), + ) + @staticmethod def get_loop_type(): """Helper function to determine the event loop type based on platform""" @@ -1276,7 +1338,7 @@ def run_server( skip_server_startup, validate_config: bool, check_migrations_against_existing_data: bool, - keepalive_timeout, + keepalive_timeout: int | None, timeout_worker_healthcheck, max_request_header_size: int | None, max_requests_before_restart: int | None, @@ -1306,6 +1368,49 @@ def run_server( run_setup_wizard() return + from pydantic import ValidationError + + from litellm.proxy.rust_forward.mode import ( + Mode, + environment_error, + private_listener, + rust_forward_mode, + stop_with_supervisor, + uvicorn_options, + ) + + 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) + runs_one_command: Final = ( + version is True + or validate_config + or health is True + or test is True + or skip_server_startup is True + or check_migrations_against_existing_data + ) + if forward_mode is Mode.SUPERVISOR and not runs_one_command: + from litellm.proxy.rust_forward.supervisor import supervise + from litellm.proxy.rust_forward.supervisor.preflight import StartupError + + server: Final = ProxyInitializationHelpers.python_server( + host, + port, + keepalive_timeout, + run_gunicorn=run_gunicorn is True, + run_hypercorn=run_hypercorn is True, + run_granian=run_granian is True, + tls=None not in (ssl_keyfile_path, ssl_certfile_path), + ) + if prometheus_metrics_port == server.port: + raise click.UsageError("--prometheus_metrics_port must differ from --port") + supervised: Final = supervise(server) + if not isinstance(supervised, StartupError): + raise SystemExit(supervised) + args: Final = locals() try: from litellm.proxy.proxy_server import ( @@ -1773,6 +1878,8 @@ def run_server( timeout_worker_healthcheck=(timeout_worker_healthcheck if running_uvicorn else None), max_request_header_size=max_request_header_size, ) + for option, value in uvicorn_options(forward_mode, WORKER_APP).items(): + uvicorn_args[option] = value # Optional: recycle uvicorn workers after N requests if max_requests_before_restart is not None: uvicorn_args["limit_max_requests"] = max_requests_before_restart @@ -1837,16 +1944,19 @@ 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, + private_listener=private_listener(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..935ff3f896e --- /dev/null +++ b/litellm/proxy/rust_forward/AGENTS.md @@ -0,0 +1,48 @@ +# Rust forwarding, Python half + +With `LITELLM_RUST_FIRST` on, a supervisor runs the Rust gateway on the public port and Python on a private loopback listener. Rust forwards every request to Python unchanged, so Python still serves every route. 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 + +The same path runs everywhere. `litellm` (the main images and a local CLI run) and `python -m gateway.launch` (the gateway image) both call `supervisor.supervise`, and nothing else starts the Rust gateway in forwarding mode. In a container the supervisor is PID 1, so its signal and exit rules are deployment behavior and not CLI polish + +## Boundary + +A supervised Python server differs from Python alone in exactly two ways, and nothing else in the proxy may change because Rust is in front + +- It listens on the private loopback listener the supervisor gives it, never on the public address +- It restores the client address Rust saw, because its own TCP peer is always Rust + +`proxy_server.py`, routes, auth and middleware never know Rust is there. A change that needs one of them to branch on forwarding is in the wrong place + +## Modules + +- `server.py` describes what an entry point would serve with on the public port, as a `PythonServer`. The supervisor reads nothing else about the entry point +- `mode.py` is the Python child's side. `Supervised` is the whole contract the supervisor hands the child, and the `uvicorn_*` and `gunicorn_*` functions translate it into each server's own options +- `peer.py` restores the client address +- `supervisor/` starts, watches and stops both children. Its rules are in `supervisor/AGENTS.md` +- `platform/` holds every OS-specific operation. Its rules are in `platform/AGENTS.md` + +## Entry points + +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 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_FIRST` off +- In the supervised child they merge `mode.uvicorn_options` or `mode.uvicorn_flags` into their server settings, or hand `mode.private_listener` to `run_gunicorn_server`, 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 +- A server is started three ways here, through uvicorn keyword arguments, uvicorn's own CLI and gunicorn, and each one states the same two facts, the private listener and that the server itself trusts no peer's forwarding headers. The uvicorn translations are in `mode.py` and turn proxy headers off. `run_gunicorn_server` applies `server_listener.gunicorn_overrides` for the `PrivateListener` it is given, which sets gunicorn's `bind` and an empty `forwarded_allow_ips`. A new way to start a server gets its own translation, never a branch in the entry point + +## Client address + +Python must see each request as if the client had connected directly. `x-litellm-rust-peer` is the only Rust-specific input, and the proxy app itself never knows Rust is in front of it + +- Rust sends the client's TCP peer in `x-litellm-rust-peer` and leaves every forwarding header untouched. It never appends itself to `X-Forwarded-For`, because the proxy parses that header on its own (`use_x_forwarded_for`, `mcp_xff_num_trusted_hops`) and an extra hop would shift what those settings resolve +- Server-level proxy headers are off in a supervised server, through `proxy_headers: False` for uvicorn, `--no-proxy-headers` for its CLI and an empty `forwarded_allow_ips` for gunicorn. The server's TCP peer is always loopback, which uvicorn and gunicorn trust by default, so leaving them on would let any client set its own address and scheme +- `peer.RustPeerMiddleware` puts the real peer back into the scope and then runs uvicorn's own `ProxyHeadersMiddleware` with the operator's trust. It wraps the app from outside, so it runs before anything that reads the client address +- The operator's trust is one string, `trusted_proxies`. It is resolved by the server library the entry point uses (`server.uvicorn_trusted_proxies` or gunicorn's config), never defaulted in this package. `server.NO_TRUSTED_PROXIES` stands for proxy headers the operator turned off, since a trust list that matches nobody behaves the same +- `peer.py` honors the peer header only from the loopback hop and always removes it. Any process on the host can connect to the private listener and send that header, so loopback is the trust boundary of the private listener. That widens what Python alone allows only when the operator's trust list leaves loopback out + +## Environment contracts + +Each environment contract is written and read in one place + +- The Python child's contract 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 +- `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 there and a matching field in `settings.rs` in the same change +- `mode.rust_forward_mode` reads `LITELLM_RUST_FIRST` strictly, so an invalid value stops the CLI at startup with one line naming it. It is independent of `LITELLM_RUST`, which only selects the Rust paths inside the SDK in `litellm/rust_bridge/configuration.py`, so neither turns the other on 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..1c417db345f --- /dev/null +++ b/litellm/proxy/rust_forward/mode.py @@ -0,0 +1,142 @@ +import contextlib +import os +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, TypeAdapter, ValidationError +from typing_extensions import assert_never + +from litellm.proxy.server_listener import LOOPBACK_HOST, AssignedPort, InheritedFd, PrivateListener +from litellm.rust_bridge.configuration import RustSwitch + +SWITCH_ENV: Final = "LITELLM_RUST_FIRST" +SUPERVISED_ENV: Final = "_LITELLM_GATEWAY_SUPERVISED" +RUST_HOP: Final = LOOPBACK_HOST + + +class _Switch(BaseModel): + model_config = ConfigDict(frozen=True) + + rust: RustSwitch = Field(None, validation_alias=SWITCH_ENV) + + +class Mode(Enum): + DISABLED = auto() + SUPERVISOR = auto() + + +@dataclass(frozen=True, slots=True) +class Supervised: + listener: PrivateListener + trusted_proxies: str + + +RustForwardMode: TypeAlias = Mode | Supervised +_SUPERVISED: Final = TypeAdapter(Supervised) + + +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: + 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: + 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 _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: + case Supervised(): + threading.Thread(target=_stop_when_closed, args=(fd, stop), name="supervisor-stdin", daemon=True).start() + case Mode.DISABLED | Mode.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. + + +def _uvicorn_listener(listener: PrivateListener) -> 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 Mode.DISABLED | Mode.SUPERVISOR: + return MappingProxyType({}) + case _: + assert_never(mode) + + +def _uvicorn_listener_flags(listener: PrivateListener) -> 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 Mode.DISABLED | Mode.SUPERVISOR: + return () + case _: + assert_never(mode) + + +def private_listener(mode: RustForwardMode) -> PrivateListener | None: + match mode: + case Supervised(listener): + return listener + case Mode.DISABLED | Mode.SUPERVISOR: + return None + 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..1dd52e24564 --- /dev/null +++ b/litellm/proxy/rust_forward/peer.py @@ -0,0 +1,69 @@ +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, Mode, RustForwardMode, Supervised + +if TYPE_CHECKING: + import starlette.types + from uvicorn._types import ASGI3Application, ASGIReceiveCallable, ASGISendCallable, Scope, WWWScope + +PEER_HEADER: Final = b"x-litellm-rust-peer" + +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: + 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) + + 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: "starlette.types.ASGIApp", mode: RustForwardMode) -> ASGIApp: + converted: Final = served(app) + match mode: + case Supervised(): + return RustPeerMiddleware(converted, mode) + case Mode.DISABLED | Mode.SUPERVISOR: + return converted + case _: + assert_never(mode) diff --git a/litellm/proxy/rust_forward/platform/AGENTS.md b/litellm/proxy/rust_forward/platform/AGENTS.md new file mode 100644 index 00000000000..32d1e4ac5c9 --- /dev/null +++ b/litellm/proxy/rust_forward/platform/AGENTS.md @@ -0,0 +1,19 @@ +# Platforms + +Every OS-specific operation sits behind `Platform`, which `supervise` picks once with `current_platform()`. The supervisor, `mode.py` and the entry points never check the OS themselves, so a new OS difference becomes a `Platform` member and never an `os.name` check somewhere else + +A `Platform` answers these questions and nothing else. What the gateway binary is called, which servers it can supervise, how the private listener reaches Python, which spawn options the children get, which signals are relayed and which of them stop Rust, and how a child is signalled and stopped + +## POSIX + +`PosixPlatform` binds the private socket itself and hands Python the open descriptor (`server_listener.InheritedFd`). Nobody else can take the port between the supervisor picking it and Python serving it. The socket does not listen until Python does, so the readiness probe only succeeds against Python. Children run in their own session (`start_new_session`) only so that a terminal Ctrl-C reaches Python once, through the supervisor + +## Windows + +Windows is supported for running the proxy locally from the CLI. Nobody deploys the proxy on Windows, but `LITELLM_RUST_FIRST` is meant to become the default, and then this is the path every Windows laptop runs. Revisit the choices below before treating Windows as a deployment target + +- Python gets an assigned loopback port (`server_listener.AssignedPort`) instead of an open socket, which uvicorn on Windows cannot serve. Another process can 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 +- Each child runs in its own process group, which can only be sent `CTRL_BREAK_EVENT`. Both relayed signals arrive as that, so a second Ctrl-C does not force Python's exit the way it does on POSIX +- gunicorn is refused, since it does not run on Windows +- The Windows-only names are spelled out as constants in `windows.py` so the module imports and type checks on every OS. `tests/windows_tests/test_rust_forward_platform.py` pins them to the stdlib on the Windows runner +- No CI job runs the supervisor itself on Windows, because the `windows-sdk` job installs no proxy dependencies and no `litellm-gateway.exe`. Run a change to `windows.py` by hand on Windows until one does diff --git a/litellm/proxy/rust_forward/platform/__init__.py b/litellm/proxy/rust_forward/platform/__init__.py new file mode 100644 index 00000000000..4043adacdfb --- /dev/null +++ b/litellm/proxy/rust_forward/platform/__init__.py @@ -0,0 +1,58 @@ +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.server import Server +from litellm.proxy.server_listener import PrivateListener + + +@dataclass(frozen=True, slots=True) +class PlatformListener: + listener: PrivateListener + 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[PlatformListener]: ... + + def spawn_options(self, pass_fds: tuple[int, ...]) -> SpawnOptions: ... + + @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: ... + + 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..0d0c6c96649 --- /dev/null +++ b/litellm/proxy/rust_forward/platform/posix.py @@ -0,0 +1,64 @@ +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 +from litellm.proxy.rust_forward.platform import PlatformListener, SpawnOptions +from litellm.proxy.rust_forward.server import Server +from litellm.proxy.server_listener import InheritedFd + +_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[PlatformListener]: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as private: + private.bind((RUST_HOP, 0)) + _, port = _ADDRESS.validate_python(private.getsockname()) + yield PlatformListener(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) + + @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): + 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..163b96b9b20 --- /dev/null +++ b/litellm/proxy/rust_forward/platform/windows.py @@ -0,0 +1,68 @@ +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 +from litellm.proxy.rust_forward.platform import PlatformListener, SpawnOptions +from litellm.proxy.rust_forward.server import Server +from litellm.proxy.server_listener import AssignedPort + +_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[PlatformListener]: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: + probe.bind((RUST_HOP, 0)) + _, port = _ADDRESS.validate_python(probe.getsockname()) + yield PlatformListener(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) + + @property + def relayed_signals(self) -> tuple[signal.Signals, ...]: + return (signal.SIGINT, signal.Signals(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) + + 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..95a693291bf --- /dev/null +++ b/litellm/proxy/rust_forward/server.py @@ -0,0 +1,75 @@ +from dataclasses import dataclass +from enum import Enum +from typing import Final + +from pydantic import BaseModel, ConfigDict, TypeAdapter + +from litellm.proxy.server_listener import NO_TRUSTED_PROXIES + +_TRUSTED_PROXIES: Final = TypeAdapter(str) + + +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 + 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 + + @property + def trusted_proxies(self) -> str: + return uvicorn_trusted_proxies(self.forwarded_allow_ips) if self.proxy_headers else NO_TRUSTED_PROXIES + + +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/AGENTS.md b/litellm/proxy/rust_forward/supervisor/AGENTS.md new file mode 100644 index 00000000000..65af84481e5 --- /dev/null +++ b/litellm/proxy/rust_forward/supervisor/AGENTS.md @@ -0,0 +1,21 @@ +# Forwarding supervisor + +`supervise` serves the public port from Rust and runs the entry point's own command behind it, or says why that cannot start. That command is the supervisor's own interpreter command line (`OWN_COMMAND`, from `sys.orig_argv`), so the child has the same import path and interpreter flags as the process the operator started. An entry point never builds a second way to start itself. It is transparent to Python, so a supervised proxy behaves on signals and exits exactly as Python alone. `tests/integration/configuration/test_rust_forward_lifecycle.py` checks that equivalence against Python alone, and a change that breaks one of the rules below breaks it + +## Split by job + +- `preflight` decides whether forwarding can start. It is the only place that refuses, and a refusal is a returned `StartupError`, never an exit, so the entry point can serve from Python alone +- `children` says what each child runs with, and it is the only writer of both environment contracts. `python_spec` writes the `mode.Supervised` variable and hands the Python child the entry point's own arguments with the private port appended. `rust_spec` writes the `LITELLM_FORWARD_*` variables. Each child gets only those variables on top of the supervisor's own environment +- `lifecycle` decides. `step` is 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. It touches no process, socket or clock +- `runtime` acts. It only executes the actions `step` returns and feeds events back. The `contextlib.AsyncExitStack` it enters each child into owns their cleanup, last spawned first + +A new behavior is a new fact, event or action in `lifecycle` with a property test in `test_lifecycle.py`, never a condition in `runtime` + +## Rules + +- 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 +- 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 +- A container's PID 1 cannot die by a signal it sends itself, since the kernel drops it. uvicorn alone survives its own re-raise there and exits 0, so `runtime.exit_like` exits 0 when a stop signal killed Python and with 128 plus the signal number for any other signal, which is how a kill of Python alone is reported +- 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. Python alone takes no connections until it has booted, so the public port refuses them until then too +- Only a preflight refusal falls back to Python alone. `before_start` runs once forwarding is known to be possible, so a refused start leaves nothing behind, and a child that fails to spawn after that exits 1 diff --git a/litellm/proxy/rust_forward/supervisor/__init__.py b/litellm/proxy/rust_forward/supervisor/__init__.py new file mode 100644 index 00000000000..b6611e30ffa --- /dev/null +++ b/litellm/proxy/rust_forward/supervisor/__init__.py @@ -0,0 +1,55 @@ +import asyncio +import sys +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 SpawnFailed, exit_like, run + +OWN_COMMAND: Final = (sys.executable, *sys.orig_argv[1:]) + + +def _nothing_to_prepare() -> None: + pass + + +def supervise( + server: PythonServer, + command: Sequence[str] = OWN_COMMAND, + 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: + outcome: Final = asyncio.run(run(chosen, checked, command, server)) + except OSError as error: + return _startup_failed(error) + match outcome: + case SpawnFailed(error): + return _startup_failed(error) + case int(): + return exit_like(outcome, chosen.stop_signals) + 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/children.py b/litellm/proxy/rust_forward/supervisor/children.py new file mode 100644 index 00000000000..9d8f074f0f2 --- /dev/null +++ b/litellm/proxy/rust_forward/supervisor/children.py @@ -0,0 +1,43 @@ +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, PlatformListener, 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: PlatformListener, server: PythonServer +) -> ChildSpec: + return ChildSpec( + argv=(*command, "--port", str(private.port)), + env_overrides=supervised_environment(Supervised(private.listener, 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..a561e77ba8d --- /dev/null +++ b/litellm/proxy/rust_forward/supervisor/preflight.py @@ -0,0 +1,28 @@ +import os +import shutil +from dataclasses import dataclass + +from litellm.proxy.rust_forward.mode import SWITCH_ENV +from litellm.proxy.rust_forward.platform import Platform +from litellm.proxy.rust_forward.server import Listener, PythonServer, Server + + +@dataclass(frozen=True, slots=True) +class StartupError: + reason: str + + +def gateway_binary(platform: Platform, scripts_dir: str) -> str | None: + return shutil.which(os.path.join(scripts_dir, platform.gateway_name)) + + +def startup_check(platform: Platform, server: PythonServer, binary: str | None) -> str | StartupError: + if server.server in (Server.HYPERCORN, Server.GRANIAN): + return StartupError(f"{SWITCH_ENV} does not support {server.server.value}") + if server.server not in platform.servers: + return StartupError(f"{SWITCH_ENV} does not support {server.server.value} on {platform.name}") + if server.listener is not Listener.TCP: + return StartupError(f"{SWITCH_ENV} does not support {server.listener.value} on the public port") + if binary is None: + return StartupError(f"{SWITCH_ENV} 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..7b727fc0d80 --- /dev/null +++ b/litellm/proxy/rust_forward/supervisor/runtime.py @@ -0,0 +1,194 @@ +import asyncio +import contextlib +import math +import os +import signal +from collections.abc import AsyncGenerator, Callable, Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType +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 +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 + + +@dataclass(frozen=True, slots=True) +class SpawnFailed: + error: OSError + + +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: + specs: Final = MappingProxyType( + { + Child.PYTHON: python_spec(platform, command, private, server), + Child.RUST: rust_spec(platform, binary, private.port, server), + } + ) + 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 _die_by(signum: signal.Signals) -> None: + with contextlib.suppress(OSError): + signal.signal(signum, signal.SIG_DFL) + signal.raise_signal(signum) + + +def exit_like( + status: int, stop_signals: frozenset[signal.Signals], die_by: Callable[[signal.Signals], None] = _die_by +) -> int: + """A child killed by a signal kills the supervisor with the same signal, as uvicorn re-raises the one it caught. + + SIGKILL takes no handler and needs none to kill. The kernel drops the signal when the supervisor is a container's + PID 1. uvicorn alone survives its own re-raise there and exits 0, so a stop signal ends the supervisor with 0 too, + and any other signal with the status a shell reports for it.""" + if status >= 0: + return status + signum: Final = signal.Signals(-status) + die_by(signum) + return 0 if signum in stop_signals else 128 + signum + + +@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. Stopping is shielded so a cancelled supervisor still waits +# for the child it is stopping. +@contextlib.asynccontextmanager +async def _process(platform: Platform, spec: ChildSpec) -> 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, + ) + try: + yield process + finally: + with anyio.CancelScope(shield=True): + platform.stop_group(process) + await process.wait() + if process.stdin is not None: + process.stdin.close() + + +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_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: + 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/litellm/rust_bridge/AGENTS.md b/litellm/rust_bridge/AGENTS.md index 167b91c9b89..a13f57fae4b 100644 --- a/litellm/rust_bridge/AGENTS.md +++ b/litellm/rust_bridge/AGENTS.md @@ -46,6 +46,11 @@ - A failure after execution starts is terminal and never triggers fallback - Creating an unstarted native coroutine acquires no clients or credentials and captures no execution context +# The `LITELLM_RUST` switch + +- `decision` treats a blank or invalid `LITELLM_RUST` as unset and caches the parse by raw value. It runs inside SDK calls, so it never fails one over a switch and never validates per call +- It does not put the Rust gateway in front of Python. That is `LITELLM_RUST_FIRST`, which `litellm/proxy/rust_forward/mode.py` reads strictly at startup with this same `RustSwitch` type, stopping the CLI with one line naming the variable + # Native extension - `_native` ships in every wheel and the sdist builds it, so it is always importable 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/_support/process.py b/tests/integration/_support/process.py index 491fc87dce2..7a1bd7208d5 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/_support/rust_forward.py b/tests/integration/_support/rust_forward.py new file mode 100644 index 00000000000..14b557897c5 --- /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_FIRST=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_FIRST": "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 new file mode 100644 index 00000000000..cc4e0dc2770 --- /dev/null +++ b/tests/integration/configuration/test_rust_forward.py @@ -0,0 +1,110 @@ +from dataclasses import dataclass +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 +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" + + +@dataclass(frozen=True, slots=True) +class Caller: + connects_from: str + forwarded_for: str + resolved_ip: str + + +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 {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_only_the_trusted_balancer_can_name_the_client_ip_behind_rust( + tmp_path: Path, server_arguments: tuple[str, ...] +) -> None: + 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 + ), + bool, + seconds=10, + ) + assert len(logged) == 1 and f"{CLIENT}:" in logged[0], logged + + +@pytest.mark.parametrize("entry_point", list(EntryPoint), ids=[entry_point.name for entry_point in EntryPoint]) +def test_the_python_child_is_started_with_the_supervisors_own_command(tmp_path: Path, entry_point: EntryPoint) -> None: + with forwarded_proxy( + entry_point, + {"model_list": [], "general_settings": {"master_key": "os.environ/LITELLM_MASTER_KEY"}}, + tmp_path, + ): + supervisor: Final = next( + started for started in psutil.Process().children() if entry_point.value in started.cmdline() + ) + own_command: Final = supervisor.cmdline()[1:] + python_children: Final = [ + child.cmdline()[1:] for child in supervisor.children() if entry_point.value in child.cmdline() + ] + assert "-I" in own_command + assert [child[: len(own_command)] for child in python_children] == [own_command] 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..b3ffb8b7169 --- /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, TypeVar + +import httpx +import psutil +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") + +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(): + from litellm.proxy.rust_forward.mode import Supervised, rust_forward_mode, stop_with_supervisor + + 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": mode.listener.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, "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) +""" + +T = TypeVar("T") +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: + 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): + send_signal(process, signal.SIGINT) + time.sleep(0.5) + send_signal(process, signal.SIGINT) + status: Final = _exit_status(process, EXIT_SECONDS) + return Run(Outcome(status, False), time.monotonic() - started) + case Scenario.HUP: + 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) + send_signal(process, 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]], T]) -> T: + 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)): + kill_stray_process(leftover.pid) + + +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 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))) + kill_stray_process(process.pid) + return eventually(lambda: tuple(pid for pid in descendants if _alive(pid)), lambda alive: not alive) + + survivors: Final = _run( + (sys.executable, str(scripts / "supervise.py"), str(workers), str(port), str(scripts)), + _environment(tmp_path), + 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 new file mode 100644 index 00000000000..4cd2f09e6d7 --- /dev/null +++ b/tests/integration/configuration/test_rust_forward_traffic.py @@ -0,0 +1,387 @@ +import asyncio +import json +import os +import socket +import threading +import time +from collections.abc import AsyncIterator, Callable, Iterator, Mapping +from contextlib import AbstractContextManager, ExitStack +from dataclasses import dataclass +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.mcp import mcp_peer +from tests.integration._support.rust_forward import MASTER_KEY, EntryPoint, forwarded_proxy +from tests.integration._support.wire import Reply, Request, wire_server + +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 + 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 _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]: + 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 + ) + proxy: Final = stack.enter_context( + forwarded_proxy( + EntryPoint.LITELLM, + { + "model_list": [ + _deployment(CHAT_MODEL, "openai/gpt-4o-mini", f"{provider.url}/v1"), + { + **_deployment(REALTIME_MODEL, "openai/gpt-realtime", provider.url), + "model_info": {"mode": "realtime"}, + }, + _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)), + ) + ) + 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: + 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..95b3a94c3bc 100644 --- a/tests/unit/gateway/test_launch.py +++ b/tests/unit/gateway/test_launch.py @@ -15,10 +15,19 @@ import pytest 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 ( + Mode, + RustForwardMode, + Supervised, + uvicorn_flags, +) +from litellm.proxy.rust_forward.server import NO_TRUSTED_PROXIES, Listener +from litellm.proxy.server_listener import AssignedPort, InheritedFd +from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager DB_ENV: Final = { "DATABASE_HOST": "db.internal", @@ -114,6 +123,91 @@ 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-certfile", "unused"), Listener.TLS), + (("--ssl-ciphers", "TLSv1.2"), Listener.TCP), + ], + ) + def test_forwarding_sees_the_public_listener_uvicorn_would_open( + self, arguments: tuple[str, ...], listener: Listener + ) -> None: + assert python_server(arguments).listener is listener + + @pytest.mark.parametrize( + ("argv", "environ", "trust"), + [ + ((), {}, "127.0.0.1"), + ((), {"FORWARDED_ALLOW_IPS": "10.9.9.9"}, "10.9.9.9"), + (("--forwarded-allow-ips", "10.0.0.5"), {"FORWARDED_ALLOW_IPS": "10.9.9.9"}, "10.0.0.5"), + (("--forwarded-allow-ips=*",), {}, "*"), + (("--no-proxy-headers", "--forwarded-allow-ips=*"), {}, NO_TRUSTED_PROXIES), + ], + ) + def test_forwarding_carries_the_trust_uvicorn_resolves_to_the_child( + self, + argv: tuple[str, ...], + environ: dict[str, str], + trust: 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.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 + + @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 + + @pytest.mark.parametrize( + ("mode", "served"), + [ + (Mode.DISABLED, {"host": "0.0.0.0", "port": 4000, "fd": None, "proxy_headers": True}), + ( + Supervised(InheritedFd(17), "*"), + {"host": "0.0.0.0", "port": 4000, "fd": 17, "proxy_headers": False}, + ), + ( + Supervised(AssignedPort(45556), "*"), + {"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"})) assert params["app"] == GATEWAY_APP @@ -229,6 +323,70 @@ class TestMain: DatabaseURLSettings.from_env().apply_to_env() assert os.environ["DATABASE_URL"] == pooled + @pytest.mark.parametrize( + ("rust", "warning"), + [ + ("false", ""), + ( + "true", + "LITELLM_RUST_FIRST 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_FIRST", rust) + 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))) + assert _uvicorn_params(served[0])["uds"] == "/tmp/gateway.sock" + assert capsys.readouterr().err == warning + + def test_an_invalid_rust_switch_stops_the_gateway_naming_the_variable( + self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("LITELLM_RUST_FIRST", "maybe") + monkeypatch.delenv("_LITELLM_GATEWAY_SUPERVISED", raising=False) + served: Final[list[tuple[str, ...]]] = [] + with pytest.raises(SystemExit) as stopped: + main(("--workers", "4"), serve=lambda argv: served.append(tuple(argv))) + assert str(stopped.value).startswith("LITELLM_RUST_FIRST is invalid: ") + assert served == [] + + @pytest.mark.parametrize( + ("argv", "status", "said"), + [ + pytest.param(("--nope",), 2, "--nope", id="a bad flag is a usage error"), + pytest.param(("--help",), 0, "--timeout-keep-alive", id="help is printed once and exits cleanly"), + ], + ) + def test_the_supervisor_reports_uvicorn_arguments_the_way_uvicorn_does( + self, + argv: tuple[str, ...], + status: int, + said: str, + password_env: dict[str, str], + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], + ) -> None: + monkeypatch.setenv("LITELLM_RUST_FIRST", "true") + monkeypatch.delenv("_LITELLM_GATEWAY_SUPERVISED", raising=False) + served: Final[list[tuple[str, ...]]] = [] + with pytest.raises(SystemExit) as stopped: + main(argv, serve=lambda served_argv: served.append(tuple(served_argv))) + printed: Final = capsys.readouterr() + assert stopped.value.code == status + assert (printed.out + printed.err).count(said) == 1 + assert served == [] + 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/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..bac5a68b7fc --- /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, 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..f6ab9732b80 --- /dev/null +++ b/tests/unit/proxy/rust_forward/supervisor/test_children.py @@ -0,0 +1,54 @@ +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 Supervised, rust_forward_mode +from litellm.proxy.rust_forward.platform import PlatformListener +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 litellm.proxy.server_listener import AssignedPort, InheritedFd +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 = PlatformListener(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), "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), "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, PlatformListener(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..4b479f3d22d --- /dev/null +++ b/tests/unit/proxy/rust_forward/supervisor/test_lifecycle.py @@ -0,0 +1,158 @@ +import functools +import itertools +import signal +from dataclasses import dataclass, replace +from typing import Final + +import pytest +from hypothesis import 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, + Probe, + Send, + Spawn, + boot, + exit_status, + step, +) + +PY: Final = Child.PYTHON +RUST: Final = Child.RUST +LISTENING: Final = Probe.PYTHON_LISTENING +TERM: Final = signal.SIGTERM +INT: Final = signal.SIGINT +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 Observed: + """What an outside observer of the two processes knows, without looking inside ``Lifecycle``.""" + + state: Lifecycle + alive: frozenset[Child] = frozenset({PY}) + rust_spawns: int = 0 + stop_requested: bool = False + rust_lost: bool = False + python_status: int | None = None + + +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 + + +EVENTS: Final = st.lists( + st.one_of( + 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, +) + + +@given(EVENTS) +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 new file mode 100644 index 00000000000..82daf5b7ea0 --- /dev/null +++ b/tests/unit/proxy/rust_forward/supervisor/test_preflight.py @@ -0,0 +1,77 @@ +import sys +from pathlib import Path +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 +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( + ("platform", "server", "refusal"), + [ + (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), + "does not support an inherited file descriptor on the public port", + ), + ], +) +def test_only_plain_tcp_on_a_supported_server_can_be_forwarded( + platform: Platform, server: PythonServer, refusal: str | None +) -> None: + expected: Final = "gateway" if refusal is None else StartupError(f"LITELLM_RUST_FIRST {refusal}") + assert startup_check(platform, server, "gateway") == expected + + +@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( + server, ("python",), before_start=lambda: prepared.append(True), scripts_dir=str(tmp_path), platform=POSIX + ) + assert (refused, prepared) == (StartupError(f"LITELLM_RUST_FIRST {refusal}"), []) + assert ( + capsys.readouterr().err + == f"LITELLM_RUST_FIRST {refusal}, so Python serves the public port without Rust forwarding\n" + ) + + +@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..b84b1302e3c --- /dev/null +++ b/tests/unit/proxy/rust_forward/supervisor/test_runtime.py @@ -0,0 +1,189 @@ +import signal +import subprocess +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 ( + Action, + Caught, + Child, + ChildExited, + Event, + Probe, + Send, + Spawn, +) +from litellm.proxy.rust_forward.supervisor.runtime import SpawnFailed, drive, exit_like, open_children +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" + + +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 +class ScriptedExecutor: + """Replays a fixed event script and records every action the driver carries out.""" + + events: list[Event] + failing: frozenset[Child] + done: list[Action] = field(default_factory=list) + + async def spawn(self, child: Child) -> SpawnFailed | None: + self.done.append(Spawn(child)) + 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)) + + async def next_event(self) -> Event: + return self.events.pop(0) + + +@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]: + 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) + + +def _dropped(_signum: signal.Signals) -> None: + pass + + +@POSIX_ONLY +@pytest.mark.parametrize( + ("status", "exits"), + [ + pytest.param(3, 3, id="an exit status is passed on"), + pytest.param(-signal.SIGTERM, 0, id="a stop signal ends it the way uvicorn alone ends as PID 1"), + pytest.param(-signal.SIGKILL, 137, id="a kill is reported the way a shell reports it"), + ], +) +def test_a_supervisor_the_kernel_will_not_signal_still_exits_as_python_alone_would(status: int, exits: int) -> None: + assert exit_like(status, STOPS, die_by=_dropped) == exits + + +@POSIX_ONLY +@pytest.mark.parametrize("signum", [signal.SIGTERM, signal.SIGKILL]) +def test_the_supervisor_dies_by_the_signal_that_killed_python(signum: signal.Signals) -> None: + script: Final = ( + "from litellm.proxy.rust_forward.supervisor.runtime import exit_like\n" + f"raise SystemExit(exit_like({-signum}, frozenset()))" + ) + assert subprocess.run((sys.executable, "-I", "-c", script), check=False, timeout=60).returncode == -signum 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..59f9feb1d3b --- /dev/null +++ b/tests/unit/proxy/rust_forward/test_mode.py @@ -0,0 +1,127 @@ +import os +import threading +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 ( + SUPERVISED_ENV, + Mode, + RustForwardMode, + Supervised, + rust_forward_mode, + private_listener, + stop_with_supervisor, + supervised_environment, + uvicorn_options, +) +from litellm.proxy.rust_forward.server import NO_TRUSTED_PROXIES +from litellm.proxy.server_listener import AssignedPort, InheritedFd, PrivateListener + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + *((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(value: str | None, expected: RustForwardMode) -> None: + assert rust_forward_mode({} if value is None else {"LITELLM_RUST_FIRST": value}) == expected + + +@pytest.mark.parametrize( + ("environ", "expected"), + [ + ({"LITELLM_RUST": "true"}, Mode.DISABLED), + ({"LITELLM_RUST": "enabled"}, Mode.DISABLED), + ({"LITELLM_RUST": "false", "LITELLM_RUST_FIRST": "true"}, Mode.SUPERVISOR), + ], +) +def test_forwarding_is_independent_of_the_sdk_rust_switch(environ: dict[str, str], expected: RustForwardMode) -> None: + assert rust_forward_mode(environ) == expected + + +@pytest.mark.parametrize( + ("environ", "invalid"), + [ + ({"LITELLM_RUST_FIRST": "enabled"}, "LITELLM_RUST_FIRST"), + ({"LITELLM_RUST_FIRST": "2"}, "LITELLM_RUST_FIRST"), + ({SUPERVISED_ENV: '{"listener": {"fd": "seventeen"}, "trusted_proxies": "*"}'}, "fd"), + ({SUPERVISED_ENV: '{"listener": {"fd": 17}, "trusted_proxies": 7}'}, "trusted_proxies"), + ({SUPERVISED_ENV: '{"listener": {"fd": 17}}'}, "trusted_proxies"), + ], +) +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( + "written", + [ + Supervised(InheritedFd(17), "127.0.0.1"), + Supervised(InheritedFd(17), "127.0.0.1,10.0.0.5"), + Supervised(InheritedFd(17), NO_TRUSTED_PROXIES), + Supervised(AssignedPort(45555), "*"), + ], +) +def test_child_reads_back_exactly_what_the_supervisor_wrote(written: Supervised) -> None: + inherited: Final = {"LITELLM_RUST_FIRST": "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: + options: Final = { + **ProxyInitializationHelpers.get_default_unvicorn_init_args("0.0.0.0", 45555), + **uvicorn_options(mode, WORKER_APP), + } + return uvicorn.Config(**options) + + +PUBLIC: Final = ("0.0.0.0", 45555, None) + + +@pytest.mark.parametrize( + ("mode", "app", "uvicorn_listener", "gunicorn_listener", "proxy_headers"), + [ + (Mode.DISABLED, "litellm.proxy.proxy_server:app", PUBLIC, None, True), + (Mode.SUPERVISOR, "litellm.proxy.proxy_server:app", PUBLIC, None, True), + (Supervised(InheritedFd(17), "*"), WORKER_APP, ("0.0.0.0", 45555, 17), InheritedFd(17), False), + (Supervised(AssignedPort(45556), "*"), WORKER_APP, ("127.0.0.1", 45556, None), AssignedPort(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: PrivateListener | None, + proxy_headers: bool, +) -> None: + config: Final = _served_uvicorn_config(mode) + 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 private_listener(mode) == gunicorn_listener + + +@pytest.mark.parametrize( + ("mode", "stops"), + [(Supervised(InheritedFd(17), "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() + 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_peer.py b/tests/unit/proxy/rust_forward/test_peer.py new file mode 100644 index 00000000000..11a4737c6ce --- /dev/null +++ b/tests/unit/proxy/rust_forward/test_peer.py @@ -0,0 +1,135 @@ +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 SUPERVISED_ENV, Supervised, supervised_environment +from litellm.proxy.rust_forward.peer import RustPeerMiddleware +from litellm.proxy.rust_forward.server import NO_TRUSTED_PROXIES +from litellm.proxy.server_listener import InheritedFd + +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"]))) + + proxy_headers_off: Final = mode.trusted_proxies == NO_TRUSTED_PROXIES + direct: Final = probe if proxy_headers_off else ProxyHeadersMiddleware(probe, mode.trusted_proxies) + app: Final = RustPeerMiddleware(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), "127.0.0.1"), + Supervised(InheritedFd(17), "10.0.0.0/8"), + Supervised(InheritedFd(17), "*"), + Supervised(InheritedFd(17), NO_TRUSTED_PROXIES), + ], +) +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_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), "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(SUPERVISED_ENV, supervised_environment(Supervised(InheritedFd(17), "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] + ) + 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/test_proxy_cli.py b/tests/unit/proxy/test_proxy_cli.py index eec682f365f..d85d685c352 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,8 +35,9 @@ 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.rust_forward.server import Listener, Server from litellm.proxy.server_listener import AssignedPort, InheritedFd -from litellm.proxy.shutdown.graceful_shutdown_manager import DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT +from litellm.proxy.shutdown.graceful_shutdown_manager import DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT, GracefulShutdownManager @pytest.fixture(autouse=True) @@ -4382,3 +4384,154 @@ 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=None not in (params["ssl_keyfile_path"], params["ssl_certfile_path"]), + ) + + @pytest.mark.parametrize( + ("argv", "server", "listener"), + [ + ([], 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_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) + + def test_uvicorn_is_described_where_gunicorn_cannot_be_imported(self, monkeypatch: pytest.MonkeyPatch) -> None: + _, with_gunicorn = self._server([]) + monkeypatch.setitem(sys.modules, "gunicorn.config", None) + _, without_gunicorn = self._server([]) + assert without_gunicorn == with_gunicorn + + @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) + @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, + argv: list[str], + reason: str, + served_keyfile: str | None, + ) -> None: + from click.testing import CliRunner + + 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_SUPERVISED") + }, + "LITELLM_RUST_FIRST": "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", *argv]) + assert result.exit_code == 0, result.output + assert f"LITELLM_RUST_FIRST {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.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 click.testing import CliRunner + + with patch.dict(os.environ, {"LITELLM_RUST_FIRST": "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_FIRST is invalid: ") + + def test_a_metrics_port_on_the_public_port_is_refused_before_forwarding_starts(self) -> None: + from click.testing import CliRunner + + with ( + patch.dict(os.environ, {"LITELLM_RUST_FIRST": "true"}), + patch("litellm.proxy.rust_forward.supervisor.supervise") as supervise, + ): + result: Final = CliRunner().invoke(run_server, ["--port", "45999", "--prometheus_metrics_port", "45999"]) + assert result.exit_code == 2 + assert "--prometheus_metrics_port must differ from --port" in result.output + supervise.assert_not_called() + + @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"] + ) + ) + return served.timeout_keep_alive, served.timeout_graceful_shutdown, served.forwarded_allow_ips + + @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: + 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) diff --git a/tests/windows_tests/AGENTS.md b/tests/windows_tests/AGENTS.md new file mode 100644 index 00000000000..565dba91ce6 --- /dev/null +++ b/tests/windows_tests/AGENTS.md @@ -0,0 +1,7 @@ +# Windows tests + +The `windows-sdk` job in `.github/workflows/test-merge-smoke.yml` is the only place these run, on a Windows runner with no proxy dependencies installed. They prove the SDK works on Windows without them, so a test here must not import uvicorn, fastapi, psutil or anything else outside the `dev` group, and must not need the `litellm-gateway.exe` binary + +A test that cannot run off Windows carries a `skipif` on `sys.platform`, so collecting this folder on another OS never fails + +Windows-only names such as `CTRL_BREAK_EVENT` come from `litellm/proxy/rust_forward/platform/windows.py`, which spells them out so the type checker accepts them on every OS. `test_rust_forward_platform.py` pins those constants to the stdlib diff --git a/tests/windows_tests/test_rust_forward_platform.py b/tests/windows_tests/test_rust_forward_platform.py new file mode 100644 index 00000000000..1415a11f57c --- /dev/null +++ b/tests/windows_tests/test_rust_forward_platform.py @@ -0,0 +1,35 @@ +import signal +import subprocess +import sys +from pathlib import Path +from typing import Final + +import pytest + +from litellm.proxy.rust_forward.platform.windows import ( + CREATE_NEW_PROCESS_GROUP, + CTRL_BREAK_EVENT, + SIGBREAK, + WindowsPlatform, +) +from litellm.proxy.rust_forward.supervisor.preflight import gateway_binary + + +@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 + + +@pytest.mark.skipif(sys.platform != "win32", reason="only Windows searches the working directory for an executable") +def test_gateway_binary_in_the_working_directory_is_never_chosen( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + scripts: Final = tmp_path / "scripts" + scripts.mkdir() + (tmp_path / "litellm-gateway.exe").write_bytes(b"") + monkeypatch.chdir(tmp_path) + assert gateway_binary(WindowsPlatform(), str(scripts)) is None + installed: Final = scripts / "litellm-gateway.exe" + installed.write_bytes(b"") + assert gateway_binary(WindowsPlatform(), str(scripts)) == str(installed)