diff --git a/gateway/launch.py b/gateway/launch.py index f7f98bd1156..449276ca089 100644 --- a/gateway/launch.py +++ b/gateway/launch.py @@ -132,16 +132,20 @@ def configure_database() -> None: def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -> None: from litellm.proxy.rust_forward.mode import Disabled, Supervised, Supervisor, rust_forward_mode, uvicorn_flags - from litellm.proxy.rust_forward.supervisor import supervise + from litellm.proxy.rust_forward.supervisor import StartupError, supervise mode: Final = rust_forward_mode(os.environ) match mode: - case Disabled(): + case Supervisor() if not isinstance( + exit_code := supervise( + python_server(argv), (sys.executable, "-m", "gateway.launch", *argv), before_start=configure_database + ), + StartupError, + ): + sys.exit(exit_code) + case Disabled() | Supervisor(): configure_database() serve(uvicorn_argv(argv, os.environ)) - case Supervisor(): - command: Final = (sys.executable, "-m", "gateway.launch", *argv) - sys.exit(supervise(python_server(argv), command, before_start=configure_database)) case Supervised(): serve((*uvicorn_argv(argv, os.environ), *uvicorn_flags(mode))) case _: diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index d3b3e073fa5..e99b06d6c8e 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -1368,7 +1368,7 @@ def run_server( if isinstance(forward_mode, Supervisor) and not any( (version, validate_config, health, test, skip_server_startup, check_migrations_against_existing_data) ): - from litellm.proxy.rust_forward.supervisor import supervise + from litellm.proxy.rust_forward.supervisor import StartupError, supervise server: Final = ProxyInitializationHelpers.python_server( host, @@ -1379,9 +1379,9 @@ def run_server( run_granian=run_granian, tls=any((ssl_keyfile_path, ssl_certfile_path, ciphers)), ) - raise SystemExit( - supervise(server, (sys.executable, "-m", "litellm.proxy.proxy_cli", *sys.argv[1:])) - ) + supervised: Final = supervise(server, (sys.executable, "-m", "litellm.proxy.proxy_cli", *sys.argv[1:])) + if not isinstance(supervised, StartupError): + raise SystemExit(supervised) args: Final = locals() try: diff --git a/litellm/proxy/rust_forward/AGENTS.md b/litellm/proxy/rust_forward/AGENTS.md index 294d94dd21b..64d9fb5673a 100644 --- a/litellm/proxy/rust_forward/AGENTS.md +++ b/litellm/proxy/rust_forward/AGENTS.md @@ -4,7 +4,7 @@ With `LITELLM_RUST` on, the supervisor here runs the Rust gateway on the public Every forwarding decision lives here. Entry points (`litellm/proxy/proxy_cli.py`, `gateway/launch.py`, `litellm_cli/main.py`) only translate between their own flags and this package: -- In the supervisor they describe what they would serve with as a `server.PythonServer` and hand it to `supervisor.supervise`, which decides whether forwarding can start, what Rust and the Python child receive, and how to report a refusal +- In the supervisor they describe what they would serve with as a `server.PythonServer` and hand it to `supervisor.supervise`, which decides whether forwarding can start and what Rust and the Python child receive. A refusal (no gateway binary, or a server or listener Rust cannot front) is reported and returned, and the entry point then serves from Python alone exactly as with `LITELLM_RUST` off - In the supervised child they merge `mode.uvicorn_options`, `mode.uvicorn_flags` or `mode.gunicorn_bind` and `mode.server_proxy_headers` into their server settings, and wrap the app with `peer.wrap` in their worker factory. These are identities when forwarding is off, so entry points never branch on it Every OS-specific operation lives behind `platform.Platform`, which `supervise` picks once with `platform.current_platform()`: the gateway binary name, how the private listener reaches Python, the spawn options, and how children are signalled and stopped. `PosixPlatform` hands Python an inherited socket (`mode.InheritedFd`). `WindowsPlatform` hands it an assigned loopback port (`mode.AssignedPort`), starts each child in its own process group and stops it with `CTRL_BREAK_EVENT`. The supervisor and `mode.py` never check the OS themselves diff --git a/litellm/proxy/rust_forward/supervisor.py b/litellm/proxy/rust_forward/supervisor.py index 3331b3f0c0b..5a53d3c562e 100644 --- a/litellm/proxy/rust_forward/supervisor.py +++ b/litellm/proxy/rust_forward/supervisor.py @@ -40,9 +40,7 @@ def startup_check(platform: Platform, server: PythonServer, binary: str | None) if server.listener is not Listener.TCP: return StartupError(f"{ENV_NAME} does not support {server.listener.value} on the public port") if binary is None: - return StartupError( - f"{ENV_NAME} requires the litellm-gateway binary. Install litellm[proxy] in this environment" - ) + return StartupError(f"{ENV_NAME} requires the litellm-gateway binary from the litellm-proxy package") return binary @@ -314,16 +312,17 @@ def supervise( before_start: Callable[[], None] = _nothing_to_prepare, scripts_dir: str = sysconfig.get_path("scripts"), platform: Platform | None = None, -) -> int: +) -> int | StartupError: """Serve ``server``'s public port from Rust and run ``command`` behind it, or say why that cannot start. - ``before_start`` runs only once forwarding is known to be possible, so a refused start leaves nothing behind. + A refusal is returned for the caller to serve from Python alone, and ``before_start`` runs only once forwarding + is known to be possible, so a refused start leaves nothing behind. """ chosen: Final = current_platform() if platform is None else platform checked: Final = startup_check(chosen, server, gateway_binary(chosen, scripts_dir)) if isinstance(checked, StartupError): - sys.stderr.write(f"{checked.reason}\n") - return ExitCode.CHILD_FAILED + sys.stderr.write(f"{checked.reason}, so Python serves the public port without Rust forwarding\n") + return checked before_start() try: return asyncio.run(_supervise(chosen, checked, command, server)) diff --git a/tests/unit/gateway/test_launch.py b/tests/unit/gateway/test_launch.py index 50dab887688..8a916560c1b 100644 --- a/tests/unit/gateway/test_launch.py +++ b/tests/unit/gateway/test_launch.py @@ -134,18 +134,6 @@ class TestUvicornArgv: ) -> None: assert python_server(arguments).listener is listener - def test_an_unsupported_listener_stops_forwarding_before_the_database_is_prepared( - self, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] - ) -> None: - monkeypatch.setenv("LITELLM_RUST", "true") - monkeypatch.delenv("_LITELLM_GATEWAY_FD", raising=False) - monkeypatch.setenv("LITELLM_PGBOUNCER_ENABLED", "true") - monkeypatch.setenv("LITELLM_PGBOUNCER_BINARY", "/nonexistent/pgbouncer") - with pytest.raises(SystemExit) as stopped: - main(("--uds", "/tmp/gateway.sock"), serve=lambda argv: None) - assert stopped.value.code == 1 - assert capsys.readouterr().err == "LITELLM_RUST does not support a Unix socket on the public port\n" - @pytest.mark.parametrize( ("argv", "trust"), [ @@ -327,15 +315,32 @@ class TestMain: DatabaseURLSettings.from_env().apply_to_env() assert os.environ["DATABASE_URL"] == pooled - def test_forwarding_disabled_keeps_uvicorn_socket_options( - self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch + @pytest.mark.parametrize( + ("rust", "warning"), + [ + ("false", ""), + ( + "true", + "LITELLM_RUST does not support a Unix socket on the public port, " + "so Python serves the public port without Rust forwarding\n", + ), + ], + ) + def test_python_alone_keeps_uvicorn_socket_options( + self, + rust: str, + warning: str, + password_env: dict[str, str], + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], ) -> None: - monkeypatch.setenv("LITELLM_RUST", "false") + monkeypatch.setenv("LITELLM_RUST", rust) monkeypatch.delenv("_LITELLM_GATEWAY_FD", raising=False) monkeypatch.delenv("LITELLM_PGBOUNCER_ENABLED", raising=False) served: Final[list[tuple[str, ...]]] = [] main(("--uds", "/tmp/gateway.sock"), serve=lambda argv: served.append(tuple(argv))) assert _uvicorn_params(served[0])["uds"] == "/tmp/gateway.sock" + assert capsys.readouterr().err == warning def test_a_pooler_that_cannot_start_stops_the_gateway_before_uvicorn( self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch, tmp_path: Path diff --git a/tests/unit/proxy/rust_forward/test_supervisor.py b/tests/unit/proxy/rust_forward/test_supervisor.py index d0ba6872bbd..ea69ce86951 100644 --- a/tests/unit/proxy/rust_forward/test_supervisor.py +++ b/tests/unit/proxy/rust_forward/test_supervisor.py @@ -120,8 +120,11 @@ def test_unsupported_servers_are_refused_before_anything_starts( assert startup_check(POSIX, server, "gateway") == StartupError(reason) prepared: Final[list[bool]] = [] refused: Final = supervise(server, ("python",), before_start=lambda: prepared.append(True), platform=POSIX) - assert refused == ExitCode.CHILD_FAILED - assert (capsys.readouterr().err, prepared) == (f"{reason}\n", []) + assert refused == StartupError(reason) + assert (capsys.readouterr().err, prepared) == ( + f"{reason}, so Python serves the public port without Rust forwarding\n", + [], + ) @pytest.mark.parametrize("server", [Server.UVICORN, Server.GUNICORN]) @@ -136,9 +139,14 @@ def test_windows_forwards_uvicorn_and_refuses_gunicorn() -> None: ) -def test_missing_binary_gives_install_instruction() -> None: - assert startup_check(POSIX, _server(), None) == StartupError( - "LITELLM_RUST requires the litellm-gateway binary. Install litellm[proxy] in this environment" +def test_missing_binary_is_refused_before_anything_starts(tmp_path: Path) -> None: + prepared: Final[list[bool]] = [] + refused: Final = supervise( + _server(), ("python",), before_start=lambda: prepared.append(True), scripts_dir=str(tmp_path), platform=POSIX + ) + assert (refused, prepared) == ( + StartupError("LITELLM_RUST requires the litellm-gateway binary from the litellm-proxy package"), + [], ) diff --git a/tests/unit/proxy/test_proxy_cli.py b/tests/unit/proxy/test_proxy_cli.py index fe0e46caaf3..c7d83a00188 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -4303,23 +4303,51 @@ class TestRustForwardServer: @pytest.mark.parametrize( ("argv", "reason"), [ - (["--ssl_keyfile_path", "unused"], "does not support TLS"), - (["--ssl_certfile_path", "unused"], "does not support TLS"), - (["--ciphers", "unused"], "does not support TLS"), + (["--ssl_keyfile_path", "unused"], "does not support TLS on the public port"), + (["--ssl_certfile_path", "unused"], "does not support TLS on the public port"), + (["--ciphers", "unused"], "does not support TLS on the public port"), (["--run_hypercorn"], "does not support Hypercorn"), (["--run_granian"], "does not support Granian"), ], ) - def test_unsupported_servers_fail_before_starting_children( - self, argv: list[str], reason: str, monkeypatch: pytest.MonkeyPatch + def test_unsupported_servers_are_refused_before_starting_children(self, argv: list[str], reason: str) -> None: + from litellm.proxy.rust_forward.platform.posix import PosixPlatform + from litellm.proxy.rust_forward.supervisor import StartupError, startup_check + + _, server = self._server(argv) + assert startup_check(PosixPlatform(), server, "gateway") == StartupError(f"LITELLM_RUST {reason}") + + @patch("uvicorn.run") + @patch("atexit.register") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False) + def test_a_refused_forwarding_start_serves_from_python_alone( + self, mock_should_update, mock_setup_db, mock_atexit_register, mock_uvicorn_run ) -> None: from click.testing import CliRunner - monkeypatch.setenv("LITELLM_RUST", "true") - monkeypatch.delenv("_LITELLM_GATEWAY_FD", raising=False) - result: Final = CliRunner().invoke(run_server, argv) - assert result.exit_code == 1 - assert f"LITELLM_RUST {reason}" in result.output + proxy_module: Final = MagicMock( + app=MagicMock(), + ProxyConfig=MagicMock(), + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), + ) + environ: Final = { + **{k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL", "_LITELLM_GATEWAY_FD")}, + "LITELLM_RUST": "true", + } + with ( + patch.dict(os.environ, environ, clear=True), + patch.dict("sys.modules", {"proxy_server": proxy_module, "litellm.proxy.proxy_server": proxy_module}), + ): + result: Final = CliRunner().invoke( + run_server, ["--port", "45999", "--ssl_keyfile_path", "key.pem", "--ssl_certfile_path", "cert.pem"] + ) + assert result.exit_code == 0, result.output + assert "LITELLM_RUST does not support TLS on the public port, so Python serves the public port" in result.output + mock_uvicorn_run.assert_called_once() + assert mock_uvicorn_run.call_args.kwargs["ssl_keyfile"] == "key.pem" + assert mock_uvicorn_run.call_args.kwargs["port"] == 45999 @pytest.mark.parametrize( ("argv", "environ"),