diff --git a/tests/unit/gateway/test_launch.py b/tests/unit/gateway/test_launch.py index becafe07d5d..a08478d3c2d 100644 --- a/tests/unit/gateway/test_launch.py +++ b/tests/unit/gateway/test_launch.py @@ -12,7 +12,6 @@ from typing import Final, cast from unittest.mock import MagicMock, patch import pytest -import uvicorn from uvicorn.importer import import_from_string from uvicorn.main import main as uvicorn_main @@ -20,7 +19,14 @@ from gateway.launch import GATEWAY_APP, main, pool_database_url, python_server, from litellm.proxy.db.db_url_settings import DatabaseURLSettings from litellm.proxy.db.pgbouncer import PGBOUNCER_POOLED_ENV_VAR, PgBouncerError, PgBouncerSettings from litellm.proxy.proxy_server import app as proxy_app -from litellm.proxy.rust_forward.mode import AssignedPort, InheritedFd, Supervised, uvicorn_flags +from litellm.proxy.rust_forward.mode import ( + AssignedPort, + InheritedFd, + Mode, + RustForwardMode, + Supervised, + uvicorn_flags, +) from litellm.proxy.rust_forward.server import Listener from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager @@ -172,10 +178,6 @@ class TestUvicornArgv: served: Final = _uvicorn_params(uvicorn_argv(argv, os.environ))["timeout_keep_alive"] assert python_server(argv).keepalive_seconds == served - def test_without_a_keepalive_rust_follows_uvicorns_own_default(self, monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.delenv("KEEPALIVE_TIMEOUT", raising=False) - assert python_server(()).keepalive_seconds == uvicorn.Config(app="", log_config=None).timeout_keep_alive - @pytest.mark.parametrize("argv", [(), ("--timeout-graceful-shutdown", "12")]) def test_rust_drains_for_pythons_graceful_timeout_plus_teardown( self, argv: tuple[str, ...], monkeypatch: pytest.MonkeyPatch @@ -185,28 +187,27 @@ class TestUvicornArgv: graceful: Final = GracefulShutdownManager.get_timeout() if flag is None else float(str(flag)) assert python_server(argv).graceful_timeout_seconds == graceful - def test_supervised_listener_uses_the_inherited_loopback_socket(self): - supervised: Final = uvicorn_flags(Supervised(InheritedFd(17), True, "127.0.0.1")) - params: Final = _uvicorn_params((*uvicorn_argv(("--workers", "2"), {}), *supervised)) - assert params["fd"] == 17 - assert params["workers"] == 2 - - def test_assigned_port_listener_binds_only_loopback(self): - supervised: Final = uvicorn_flags(Supervised(AssignedPort(45556), True, "127.0.0.1")) - params: Final = _uvicorn_params((*uvicorn_argv(("--host", "0.0.0.0", "--port", "4000"), {}), *supervised)) - assert (params["host"], params["port"], params["fd"]) == ("127.0.0.1", 45556, None) - - @pytest.mark.parametrize("argv", [(), ("--proxy-headers",), ("--forwarded-allow-ips=*",)]) - def test_private_hop_leaves_proxy_headers_to_the_app(self, argv: tuple[str, ...]): - params: Final = _uvicorn_params( - (*uvicorn_argv(argv, {}), *uvicorn_flags(Supervised(InheritedFd(17), True, "*"))) - ) - assert params["proxy_headers"] is False - - def test_direct_listener_keeps_explicit_proxy_header_options(self): - params: Final = _uvicorn_params(uvicorn_argv(("--no-proxy-headers", "--forwarded-allow-ips=*"), {})) - assert params["proxy_headers"] is False - assert params["forwarded_allow_ips"] == "*" + @pytest.mark.parametrize( + ("mode", "served"), + [ + (Mode.DISABLED, {"host": "0.0.0.0", "port": 4000, "fd": None, "proxy_headers": True}), + ( + Supervised(InheritedFd(17), True, "*"), + {"host": "0.0.0.0", "port": 4000, "fd": 17, "proxy_headers": False}, + ), + ( + Supervised(AssignedPort(45556), True, "*"), + {"host": "127.0.0.1", "port": 45556, "fd": None, "proxy_headers": False}, + ), + ], + ) + def test_supervised_uvicorn_serves_the_private_hop_and_never_trusts_it( + self, mode: RustForwardMode, served: dict[str, object] + ) -> None: + argv: Final = ("--host", "0.0.0.0", "--port", "4000", "--workers", "2", "--proxy-headers") + params: Final = _uvicorn_params((*uvicorn_argv((*argv, "--forwarded-allow-ips=*"), {}), *uvicorn_flags(mode))) + assert {name: params[name] for name in served} == served + assert (params["workers"], params["forwarded_allow_ips"]) == (2, "*") def test_keepalive_env_reaches_uvicorn(self): params: Final = _uvicorn_params(uvicorn_argv(("--workers", "4"), {"KEEPALIVE_TIMEOUT": "75"})) diff --git a/tests/unit/proxy/rust_forward/supervisor/test_lifecycle.py b/tests/unit/proxy/rust_forward/supervisor/test_lifecycle.py index a71fb6b335e..4b479f3d22d 100644 --- a/tests/unit/proxy/rust_forward/supervisor/test_lifecycle.py +++ b/tests/unit/proxy/rust_forward/supervisor/test_lifecycle.py @@ -1,10 +1,11 @@ import functools +import itertools import signal from dataclasses import dataclass, replace from typing import Final import pytest -from hypothesis import example, given +from hypothesis import given from hypothesis import strategies as st from litellm.proxy.rust_forward.platform.posix import PosixPlatform @@ -16,7 +17,6 @@ from litellm.proxy.rust_forward.supervisor.lifecycle import ( ChildExited, Event, Lifecycle, - Phase, Probe, Send, Spawn, @@ -25,73 +25,127 @@ from litellm.proxy.rust_forward.supervisor.lifecycle import ( step, ) +PY: Final = Child.PYTHON +RUST: Final = Child.RUST +LISTENING: Final = Probe.PYTHON_LISTENING TERM: Final = signal.SIGTERM INT: Final = signal.SIGINT -POSIX: Final = PosixPlatform() +HUP: Final = signal.SIGHUP +STOPS: Final = PosixPlatform().stop_signals + + +def _replay(events: tuple[Event, ...]) -> tuple[tuple[tuple[Action, ...], ...], int | None]: + transitions: Final = tuple( + itertools.accumulate(events, lambda transition, event: step(transition[0], event, STOPS), initial=boot()) + ) + return tuple(actions for _, actions in transitions[1:]), exit_status(transitions[-1][0]) + + +@pytest.mark.parametrize( + ("trace", "status"), + [ + pytest.param( + ( + (LISTENING, (Spawn(RUST),)), + (Caught(TERM), (Send(RUST, TERM), Send(PY, TERM))), + (ChildExited(RUST, 0), ()), + (ChildExited(PY, -15), ()), + ), + -15, + id="TERM reaches Rust first, then Python, and the supervisor dies as Python did", + ), + pytest.param( + ( + (LISTENING, (Spawn(RUST),)), + (Caught(INT), (Send(RUST, INT), Send(PY, INT))), + (Caught(INT), (Send(RUST, INT), Send(PY, INT))), + (ChildExited(PY, 1), (Send(RUST, TERM),)), + (ChildExited(RUST, 0), ()), + ), + 1, + id="a second Ctrl-C reaches Python again to force the exit", + ), + pytest.param( + ( + (LISTENING, (Spawn(RUST),)), + (Caught(HUP), (Send(PY, HUP),)), + (Caught(signal.SIGTTIN), (Send(PY, signal.SIGTTIN),)), + (ChildExited(PY, 0), (Send(RUST, TERM),)), + (ChildExited(RUST, 0), ()), + ), + 0, + id="gunicorn's HUP and TTIN reach only Python and Rust keeps serving", + ), + pytest.param( + ( + (LISTENING, (Spawn(RUST),)), + (ChildExited(RUST, 0), (Send(PY, TERM),)), + (ChildExited(PY, -15), ()), + ), + FAILURE, + id="losing Rust while Python serves stops Python and fails", + ), + pytest.param( + ( + (Caught(TERM), (Send(PY, TERM),)), + (LISTENING, ()), + (ChildExited(PY, -15), ()), + ), + -15, + id="a stop before Python listens never starts Rust", + ), + pytest.param(((ChildExited(PY, 3), ()),), 3, id="Python failing to boot exits with its status"), + ], +) +def test_supervisor_relays_signals_and_exits_as_python_did( + trace: tuple[tuple[Event, tuple[Action, ...]], ...], status: int +) -> None: + events: Final = tuple(event for event, _ in trace) + assert _replay(events) == (tuple(actions for _, actions in trace), status) @dataclass(frozen=True, slots=True) -class World: - """What the operating system knows about the children, kept independently of ``Lifecycle``.""" +class Observed: + """What an outside observer of the two processes knows, without looking inside ``Lifecycle``.""" - alive: frozenset[Child] = frozenset({Child.PYTHON}) - spawned: frozenset[Child] = frozenset({Child.PYTHON}) - reported_listening: bool = False + state: Lifecycle + alive: frozenset[Child] = frozenset({PY}) + rust_spawns: int = 0 stop_requested: bool = False - python_status: int | None = None rust_lost: bool = False + python_status: int | None = None -def _expected_exit(world: World) -> int | None: - if world.alive or world.python_status is None: - return None - return FAILURE if world.rust_lost else world.python_status +def _observe(seen: Observed, event: Event) -> Observed: + state, actions = step(seen.state, event, STOPS) + death: Final = event if isinstance(event, ChildExited) and event.child in seen.alive else None + exited: Final = None if death is None else death.child + running: Final = seen.alive - {exited} + assert all(action.child in running for action in actions if isinstance(action, Send)), "signalled a dead child" + if isinstance(event, Caught): + to_rust: Final = (Send(RUST, event.signum),) if event.signum in STOPS and RUST in running else () + assert actions == (*to_rust, *((Send(PY, event.signum),) if PY in running else ())) + spawned: Final = Spawn(RUST) in actions + assert not spawned or (event is LISTENING and not seen.stop_requested), "started Rust out of turn" + stop_requested: Final = seen.stop_requested or (isinstance(event, Caught) and event.signum in STOPS) + after: Final = replace( + seen, + state=state, + alive=running | ({RUST} if spawned else frozenset()), + rust_spawns=seen.rust_spawns + int(spawned), + stop_requested=stop_requested, + rust_lost=seen.rust_lost or (exited is RUST and PY in running and not stop_requested), + python_status=death.status if death is not None and exited is PY else seen.python_status, + ) + settled: Final = None if after.alive or after.python_status is None else after.python_status + assert exit_status(state) == (FAILURE if after.rust_lost and settled is not None else settled) + return after -def _expected_relay(world: World, signum: signal.Signals) -> tuple[Action, ...]: - to_rust: Final = (Send(Child.RUST, signum),) if signum in STOPS and Child.RUST in world.alive else () - to_python: Final = (Send(Child.PYTHON, signum),) if Child.PYTHON in world.alive else () - return (*to_rust, *to_python) - - -def _expected_on_exit(world: World, child: Child) -> tuple[Action, ...]: - match child: - case Child.PYTHON: - return (Send(Child.RUST, TERM),) if Child.RUST in world.alive else () - case Child.RUST: - serving: Final = Child.PYTHON in world.alive and not world.stop_requested - return (Send(Child.PYTHON, TERM),) if serving else () - - -def _after(world: World, event: Event) -> World: - match event: - case Caught(signum): - return replace(world, stop_requested=world.stop_requested or signum in STOPS) - case Probe.PYTHON_LISTENING: - return replace(world, reported_listening=world.reported_listening or Child.PYTHON in world.alive) - case ChildExited(child, status) if child in world.alive: - lost: Final = child is Child.RUST and Child.PYTHON in world.alive and not world.stop_requested - return replace( - world, - alive=world.alive - {child}, - python_status=status if child is Child.PYTHON else world.python_status, - rust_lost=world.rust_lost or lost, - ) - case ChildExited(): - return world - - -def _spawned(world: World, actions: tuple[Action, ...]) -> World: - started: Final = frozenset(action.child for action in actions if isinstance(action, Spawn)) - return replace(world, alive=world.alive | started, spawned=world.spawned | started) - - -SIGNALS: Final = POSIX.relayed_signals -STOPS: Final = POSIX.stop_signals EVENTS: Final = st.lists( st.one_of( - st.sampled_from(SIGNALS).map(Caught), - st.just(Probe.PYTHON_LISTENING), + st.sampled_from(PosixPlatform().relayed_signals).map(Caught), + st.just(LISTENING), st.builds(ChildExited, st.sampled_from(Child), st.integers(-15, 3)), ), max_size=40, @@ -99,69 +153,6 @@ EVENTS: Final = st.lists( @given(EVENTS) -@example([Caught(TERM), Caught(signal.SIGHUP), Probe.PYTHON_LISTENING]) -@example([Probe.PYTHON_LISTENING, Caught(TERM), Caught(signal.SIGWINCH), ChildExited(Child.RUST, 0)]) -def test_lifecycle_relays_every_signal_and_ends_with_pythons_status(events: list[Event]) -> None: - state, actions = boot() - world = World() - assert actions == (Spawn(Child.PYTHON),) - for event in events: - before = _after(world, event) - state, actions = step(state, event, STOPS) - match event: - case Caught(signum): - assert actions == _expected_relay(before, signum) - case Probe.PYTHON_LISTENING: - starts_rust = Child.PYTHON in world.alive and not world.reported_listening and not world.stop_requested - assert actions == ((Spawn(Child.RUST),) if starts_rust else ()) - case ChildExited(child) if child in world.alive: - assert actions == _expected_on_exit(world, child) - case ChildExited(): - assert actions == () - assert all(action.child in before.alive for action in actions if isinstance(action, Send)) - assert all(action.child not in before.spawned for action in actions if isinstance(action, Spawn)) - world = _spawned(before, actions) - assert exit_status(state) == _expected_exit(world) - - -@given(EVENTS) -def test_rust_starts_at_most_once_and_only_while_python_listens_unstopped(events: list[Event]) -> None: - def advance(acc: tuple[Lifecycle, tuple[tuple[Event, Lifecycle, Action], ...]], event: Event): - state, seen = acc - following, actions = step(state, event, STOPS) - return following, (*seen, *((event, state, action) for action in actions)) - - _, history = functools.reduce(advance, events, (boot()[0], ())) - spawns: Final = tuple((event, state) for event, state, action in history if isinstance(action, Spawn)) - assert len(spawns) <= 1 - assert all( - event is Probe.PYTHON_LISTENING and state.python is Phase.SPAWNED and not state.stop_requested - for event, state in spawns - ) - - -@pytest.mark.parametrize( - ("events", "expected"), - [ - ((Probe.PYTHON_LISTENING, Caught(TERM), ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, -15)), -15), - ( - ( - Probe.PYTHON_LISTENING, - Caught(INT), - Caught(INT), - ChildExited(Child.PYTHON, 1), - ChildExited(Child.RUST, 0), - ), - 1, - ), - ( - (Probe.PYTHON_LISTENING, Caught(signal.SIGHUP), ChildExited(Child.PYTHON, -1), ChildExited(Child.RUST, 0)), - -1, - ), - ((Probe.PYTHON_LISTENING, ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, 0)), FAILURE), - ((ChildExited(Child.PYTHON, 3),), 3), - ], -) -def test_supervisor_exits_as_python_did_unless_rust_was_lost(events: tuple[Event, ...], expected: int) -> None: - final: Final = functools.reduce(lambda state, event: step(state, event, STOPS)[0], events, boot()[0]) - assert exit_status(final) == expected +def test_any_event_order_keeps_the_supervisor_transparent_to_python(events: list[Event]) -> None: + final: Final = functools.reduce(_observe, events, Observed(boot()[0])) + assert final.rust_spawns <= 1 diff --git a/tests/unit/proxy/rust_forward/supervisor/test_preflight.py b/tests/unit/proxy/rust_forward/supervisor/test_preflight.py index 71899262c30..423ee4df20a 100644 --- a/tests/unit/proxy/rust_forward/supervisor/test_preflight.py +++ b/tests/unit/proxy/rust_forward/supervisor/test_preflight.py @@ -4,6 +4,7 @@ from typing import Final import pytest +from litellm.proxy.rust_forward.platform import Platform from litellm.proxy.rust_forward.platform.posix import PosixPlatform from litellm.proxy.rust_forward.platform.windows import WindowsPlatform from litellm.proxy.rust_forward.server import Listener, PythonServer, Server @@ -17,58 +18,47 @@ POSIX_ONLY: Final = pytest.mark.skipif(sys.platform == "win32", reason="runs a P @pytest.mark.parametrize( - ("server", "reason"), + ("platform", "server", "refusal"), [ - (python_server(Server.HYPERCORN), "LITELLM_RUST does not support Hypercorn"), - (python_server(Server.GRANIAN), "LITELLM_RUST does not support Granian"), - (python_server(listener=Listener.TLS), "LITELLM_RUST does not support TLS on the public port"), - ( - python_server(listener=Listener.UNIX_SOCKET), - "LITELLM_RUST does not support a Unix socket on the public port", - ), + (POSIX, python_server(), None), + (POSIX, python_server(Server.GUNICORN), None), + (WINDOWS, python_server(), None), + (WINDOWS, python_server(Server.GUNICORN), "does not support gunicorn on Windows"), + (POSIX, python_server(Server.HYPERCORN), "does not support Hypercorn"), + (POSIX, python_server(Server.GRANIAN), "does not support Granian"), + (POSIX, python_server(listener=Listener.TLS), "does not support TLS on the public port"), + (POSIX, python_server(listener=Listener.UNIX_SOCKET), "does not support a Unix socket on the public port"), ( + POSIX, python_server(Server.GUNICORN, Listener.INHERITED_FD), - "LITELLM_RUST does not support an inherited file descriptor on the public port", + "does not support an inherited file descriptor on the public port", ), ], ) -def test_unsupported_servers_are_refused_before_anything_starts( - server: PythonServer, reason: str, capsys: pytest.CaptureFixture[str] +def test_only_plain_tcp_on_a_supported_server_can_be_forwarded( + platform: Platform, server: PythonServer, refusal: str | None ) -> None: - assert startup_check(POSIX, server, "gateway") == StartupError(reason) - prepared: Final[list[bool]] = [] - refused: Final = supervise(server, ("python",), before_start=lambda: prepared.append(True), platform=POSIX) - assert refused == StartupError(reason) - assert (capsys.readouterr().err, prepared) == ( - f"{reason}, so Python serves the public port without Rust forwarding\n", - [], - ) + expected: Final = "gateway" if refusal is None else StartupError(f"LITELLM_RUST {refusal}") + assert startup_check(platform, server, "gateway") == expected -@pytest.mark.parametrize("server", [Server.UVICORN, Server.GUNICORN]) -def test_plain_tcp_uvicorn_and_gunicorn_can_be_forwarded(server: Server) -> None: - assert startup_check(POSIX, python_server(server), "gateway") == "gateway" - - -def test_windows_forwards_uvicorn_and_refuses_gunicorn() -> None: - assert startup_check(WINDOWS, python_server(), "gateway") == "gateway" - assert startup_check(WINDOWS, python_server(Server.GUNICORN), "gateway") == StartupError( - "LITELLM_RUST does not support gunicorn on Windows" - ) - - -def test_missing_binary_is_refused_before_anything_starts(tmp_path: Path) -> None: +@pytest.mark.parametrize( + ("server", "refusal"), + [ + (python_server(Server.HYPERCORN), "does not support Hypercorn"), + (python_server(), "requires the litellm-gateway binary from the litellm-proxy package"), + ], +) +def test_a_refused_start_says_why_in_one_line_and_starts_nothing( + server: PythonServer, refusal: str, tmp_path: Path, capsys: pytest.CaptureFixture[str] +) -> None: prepared: Final[list[bool]] = [] refused: Final = supervise( - python_server(), - ("python",), - before_start=lambda: prepared.append(True), - scripts_dir=str(tmp_path), - platform=POSIX, + server, ("python",), before_start=lambda: prepared.append(True), scripts_dir=str(tmp_path), platform=POSIX ) - assert (refused, prepared) == ( - StartupError("LITELLM_RUST requires the litellm-gateway binary from the litellm-proxy package"), - [], + assert (refused, prepared) == (StartupError(f"LITELLM_RUST {refusal}"), []) + assert ( + capsys.readouterr().err == f"LITELLM_RUST {refusal}, so Python serves the public port without Rust forwarding\n" ) diff --git a/tests/unit/proxy/rust_forward/supervisor/test_runtime.py b/tests/unit/proxy/rust_forward/supervisor/test_runtime.py index 824ef9e3945..e3cd55d1130 100644 --- a/tests/unit/proxy/rust_forward/supervisor/test_runtime.py +++ b/tests/unit/proxy/rust_forward/supervisor/test_runtime.py @@ -13,7 +13,6 @@ from litellm.proxy.rust_forward.platform.windows import WindowsPlatform from litellm.proxy.rust_forward.supervisor import supervise from litellm.proxy.rust_forward.supervisor.children import ChildSpec from litellm.proxy.rust_forward.supervisor.lifecycle import ( - FAILURE, Action, Caught, Child, @@ -65,6 +64,7 @@ def test_rust_opens_the_public_port_only_once_python_listens(tmp_path: Path, mon TERM: Final = signal.SIGTERM POSIX_ONLY: Final = pytest.mark.skipif(sys.platform == "win32", reason="the child traps SIGTERM") STOPS: Final = frozenset({signal.SIGTERM, signal.SIGINT, signal.SIGQUIT}) +RUST_MISSING: Final = SpawnFailed(FileNotFoundError("litellm-gateway")) @dataclass @@ -72,12 +72,12 @@ class ScriptedExecutor: """Replays a fixed event script and records every action the driver carries out.""" events: list[Event] - failing: frozenset[Child] = frozenset() + failing: frozenset[Child] done: list[Action] = field(default_factory=list) async def spawn(self, child: Child) -> SpawnFailed | None: self.done.append(Spawn(child)) - return SpawnFailed(FileNotFoundError(child.value)) if child in self.failing else None + return RUST_MISSING if child in self.failing else None def send(self, child: Child, signum: signal.Signals) -> None: self.done.append(Send(child, signum)) @@ -86,34 +86,41 @@ class ScriptedExecutor: return self.events.pop(0) -async def test_driver_carries_out_each_action_in_order_and_ends_with_pythons_status() -> None: - executor: Final = ScriptedExecutor( - [Probe.PYTHON_LISTENING, Caught(TERM), ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, -15)] - ) - assert await drive(executor, STOPS) == -15 - assert executor.done == [Spawn(Child.PYTHON), Spawn(Child.RUST), Send(Child.RUST, TERM), Send(Child.PYTHON, TERM)] - assert executor.events == [] - - -async def test_driver_stops_python_and_fails_when_rust_is_lost() -> None: - executor: Final = ScriptedExecutor( - [Probe.PYTHON_LISTENING, ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, -15)] - ) - assert await drive(executor, STOPS) == FAILURE - assert executor.done == [Spawn(Child.PYTHON), Spawn(Child.RUST), Send(Child.PYTHON, TERM)] - - -async def test_driver_asks_for_no_event_once_the_status_is_settled() -> None: - executor: Final = ScriptedExecutor([ChildExited(Child.PYTHON, 3), Caught(TERM)]) - assert await drive(executor, STOPS) == 3 - assert executor.events == [Caught(TERM)] - - -async def test_driver_returns_a_failed_spawn_without_waiting_for_events() -> None: - executor: Final = ScriptedExecutor([Probe.PYTHON_LISTENING, Caught(TERM)], failing=frozenset({Child.RUST})) - outcome: Final = await drive(executor, STOPS) - assert isinstance(outcome, SpawnFailed) - assert executor.events == [Caught(TERM)] +@pytest.mark.parametrize( + ("events", "failing", "done", "outcome", "unread"), + [ + pytest.param( + [Probe.PYTHON_LISTENING, Caught(TERM), ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, -15)], + frozenset(), + [Spawn(Child.PYTHON), Spawn(Child.RUST), Send(Child.RUST, TERM), Send(Child.PYTHON, TERM)], + -15, + [], + id="carries out each action in order and ends with Python's status", + ), + pytest.param( + [ChildExited(Child.PYTHON, 3), Caught(TERM)], + frozenset(), + [Spawn(Child.PYTHON)], + 3, + [Caught(TERM)], + id="reads no event once the status is settled", + ), + pytest.param( + [Probe.PYTHON_LISTENING, Caught(TERM)], + frozenset({Child.RUST}), + [Spawn(Child.PYTHON), Spawn(Child.RUST)], + RUST_MISSING, + [Caught(TERM)], + id="returns a failed spawn without waiting for more events", + ), + ], +) +async def test_driver_executes_lifecycle_decisions( + events: list[Event], failing: frozenset[Child], done: list[Action], outcome: int | SpawnFailed, unread: list[Event] +) -> None: + executor: Final = ScriptedExecutor(events, failing) + assert await drive(executor, STOPS) == outcome + assert (executor.done, executor.events) == (done, unread) def _children(argv: tuple[str, ...]) -> MappingProxyType[Child, ChildSpec]: diff --git a/tests/unit/proxy/rust_forward/test_mode.py b/tests/unit/proxy/rust_forward/test_mode.py index 5acd76860fb..ffecd1948ae 100644 --- a/tests/unit/proxy/rust_forward/test_mode.py +++ b/tests/unit/proxy/rust_forward/test_mode.py @@ -24,23 +24,14 @@ from litellm.proxy.rust_forward.mode import ( @pytest.mark.parametrize( - ("environ", "expected"), + ("value", "expected"), [ - ({}, Mode.DISABLED), - ({"LITELLM_RUST": "false"}, Mode.DISABLED), - ({"LITELLM_RUST": "off"}, Mode.DISABLED), - ({"LITELLM_RUST": "0"}, Mode.DISABLED), - ({"LITELLM_RUST": "true"}, Mode.SUPERVISOR), - ({"LITELLM_RUST": "1"}, Mode.SUPERVISOR), - ({"LITELLM_RUST": "yes"}, Mode.SUPERVISOR), - ({"LITELLM_RUST": "on"}, Mode.SUPERVISOR), - ({"LITELLM_RUST": " 1 "}, Mode.SUPERVISOR), - ({"LITELLM_RUST": ""}, Mode.DISABLED), - ({"LITELLM_RUST": " "}, Mode.DISABLED), + *((value, Mode.DISABLED) for value in (None, "", " ", "0", "false", "off")), + *((value, Mode.SUPERVISOR) for value in ("1", "true", "yes", "on", " 1 ")), ], ) -def test_mode_is_resolved_from_environment(environ: dict[str, str], expected: RustForwardMode) -> None: - assert rust_forward_mode(environ) == expected +def test_mode_is_resolved_from_environment(value: str | None, expected: RustForwardMode) -> None: + assert rust_forward_mode({} if value is None else {"LITELLM_RUST": value}) == expected @pytest.mark.parametrize( @@ -84,32 +75,34 @@ def _served_uvicorn_config(mode: RustForwardMode) -> uvicorn.Config: return uvicorn.Config(**options) -def test_supervised_uvicorn_serves_the_inherited_socket_without_trusting_the_rust_hop() -> None: - config: Final = _served_uvicorn_config(Supervised(InheritedFd(17), True, "*")) - assert (config.fd, config.proxy_headers) == (17, False) - assert (config.app, config.factory) == (WORKER_APP, True) +PUBLIC: Final = ("0.0.0.0", 45555, None) -@pytest.mark.parametrize("mode", list(Mode)) -def test_unsupervised_servers_keep_their_own_socket_and_proxy_headers(mode: RustForwardMode) -> None: +@pytest.mark.parametrize( + ("mode", "app", "uvicorn_listener", "gunicorn_listener", "proxy_headers"), + [ + (Mode.DISABLED, "litellm.proxy.proxy_server:app", PUBLIC, "0.0.0.0:4000", True), + (Mode.SUPERVISOR, "litellm.proxy.proxy_server:app", PUBLIC, "0.0.0.0:4000", True), + (Supervised(InheritedFd(17), True, "*"), WORKER_APP, ("0.0.0.0", 45555, 17), "fd://17", False), + (Supervised(AssignedPort(45556), True, "*"), WORKER_APP, ("127.0.0.1", 45556, None), "127.0.0.1:45556", False), + ], +) +def test_supervised_server_serves_the_private_hop_and_never_trusts_it( + mode: RustForwardMode, + app: str, + uvicorn_listener: tuple[str, int, int | None], + gunicorn_listener: str, + proxy_headers: bool, +) -> None: config: Final = _served_uvicorn_config(mode) - assert (config.fd, config.proxy_headers, config.port) == (None, True, 45555) - assert gunicorn_bind(mode, "0.0.0.0:4000") == "0.0.0.0:4000" - assert server_proxy_headers(mode) is True - - -def test_supervised_gunicorn_binds_the_inherited_socket_without_trusting_the_rust_hop() -> None: - mode: Final = Supervised(InheritedFd(17), True, "*") - assert gunicorn_bind(mode, "0.0.0.0:4000") == "fd://17" - assert server_proxy_headers(mode) is False - - -def test_assigned_port_serves_only_the_loopback_hop_without_trusting_it() -> None: - mode: Final = Supervised(AssignedPort(45556), True, "*") - config: Final = _served_uvicorn_config(mode) - assert (config.host, config.port, config.fd, config.proxy_headers) == ("127.0.0.1", 45556, None, False) - assert gunicorn_bind(mode, "0.0.0.0:4000") == "127.0.0.1:45556" - assert server_proxy_headers(mode) is False + assert (config.app, (config.host, config.port, config.fd), config.proxy_headers) == ( + app, + uvicorn_listener, + proxy_headers, + ) + assert config.factory is (app == WORKER_APP) + assert gunicorn_bind(mode, "0.0.0.0:4000") == gunicorn_listener + assert server_proxy_headers(mode) is proxy_headers @pytest.mark.parametrize( diff --git a/tests/unit/proxy/rust_forward/test_peer.py b/tests/unit/proxy/rust_forward/test_peer.py index 01331e2c69e..fdc9695d0d7 100644 --- a/tests/unit/proxy/rust_forward/test_peer.py +++ b/tests/unit/proxy/rust_forward/test_peer.py @@ -95,16 +95,6 @@ async def test_python_behind_rust_sees_what_a_direct_connection_would( assert behind_rust == direct -async def test_untrusted_peer_cannot_set_the_scheme_behind_rust() -> None: - client, scheme, _ = await _seen_by_app( - True, - ("127.0.0.1", 60000), - ((b"x-forwarded-proto", b"https"), (b"x-litellm-rust-peer", b"203.0.113.7:51234")), - Supervised(InheritedFd(17), True, "127.0.0.1"), - ) - assert (client, scheme) == (("203.0.113.7", 51234), "http") - - async def test_peer_header_from_anyone_but_the_rust_hop_is_dropped_and_ignored() -> None: client, _, headers = await _seen_by_app( True, diff --git a/tests/unit/proxy/test_proxy_cli.py b/tests/unit/proxy/test_proxy_cli.py index df7bbd43324..4627fb625e6 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -26,6 +26,7 @@ import urllib.parse as urlparse import uvicorn import yaml +from gunicorn.config import Config as GunicornConfig from uvicorn.config import LOOP_FACTORIES from uvicorn.importer import import_from_string @@ -34,7 +35,8 @@ from litellm._logging import turn_on_json from litellm.constants import DEFAULT_MAX_REQUEST_HEADER_SIZE, GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS from litellm.proxy import proxy_cli from litellm.proxy.proxy_cli import ProxyInitializationHelpers, run_server -from litellm.proxy.shutdown.graceful_shutdown_manager import DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT +from litellm.proxy.rust_forward.server import Listener, Server +from litellm.proxy.shutdown.graceful_shutdown_manager import DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT, GracefulShutdownManager @pytest.fixture(autouse=True) @@ -4301,22 +4303,20 @@ class TestRustForwardServer: ) @pytest.mark.parametrize( - ("argv", "reason"), + ("argv", "server", "listener"), [ - ( - ["--ssl_keyfile_path", "key.pem", "--ssl_certfile_path", "cert.pem"], - "does not support TLS on the public port", - ), - (["--run_hypercorn"], "does not support Hypercorn"), - (["--run_granian"], "does not support Granian"), + ([], Server.UVICORN, Listener.TCP), + (["--run_gunicorn"], Server.GUNICORN, Listener.TCP), + (["--run_hypercorn"], Server.HYPERCORN, Listener.TCP), + (["--run_granian"], Server.GRANIAN, Listener.TCP), + (["--ssl_keyfile_path", "key.pem", "--ssl_certfile_path", "cert.pem"], Server.UVICORN, Listener.TLS), ], ) - def test_unsupported_servers_are_refused_before_starting_children(self, argv: list[str], reason: str) -> None: - from litellm.proxy.rust_forward.platform.posix import PosixPlatform - from litellm.proxy.rust_forward.supervisor.preflight import StartupError, startup_check - - _, server = self._server(argv) - assert startup_check(PosixPlatform(), server, "gateway") == StartupError(f"LITELLM_RUST {reason}") + def test_the_supervisor_sees_the_server_and_listener_python_would_use( + self, argv: list[str], server: Server, listener: Listener + ) -> None: + _, described = self._server(argv) + assert (described.server, described.listener) == (server, listener) @patch("uvicorn.run") @patch("atexit.register") @@ -4383,59 +4383,37 @@ class TestRustForwardServer: assert (result.exit_code, len(lines)) == (1, 1) assert lines[0].startswith("Error: LITELLM_RUST is invalid: ") - @pytest.mark.parametrize( - ("argv", "environ"), - [(["--keepalive_timeout", "2"], {}), ([], {"KEEPALIVE_TIMEOUT": "7"}), ([], {})], - ) - def test_rust_retires_connections_by_the_keepalive_uvicorn_runs_with( - self, argv: list[str], environ: dict[str, str], monkeypatch: pytest.MonkeyPatch - ) -> None: - monkeypatch.delenv("KEEPALIVE_TIMEOUT", raising=False) - for name, value in environ.items(): - monkeypatch.setenv(name, value) - params, server = self._server(argv) + @staticmethod + def _served_by_python(params: dict[str, object]) -> tuple[object, ...]: + if params["run_gunicorn"]: + gunicorn: Final = GunicornConfig() + return gunicorn.keepalive, GracefulShutdownManager.get_timeout(), ",".join(gunicorn.forwarded_allow_ips) served: Final = uvicorn.Config( **ProxyInitializationHelpers.get_default_unvicorn_init_args( "127.0.0.1", 45555, keepalive_timeout=params["keepalive_timeout"] ) ) - assert server.keepalive_seconds == served.timeout_keep_alive + return served.timeout_keep_alive, served.timeout_graceful_shutdown, served.forwarded_allow_ips - def test_gunicorn_workers_keep_gunicorns_keepalive_and_rust_follows(self, monkeypatch: pytest.MonkeyPatch) -> None: - from gunicorn.config import Config as GunicornConfig - - monkeypatch.delenv("KEEPALIVE_TIMEOUT", raising=False) - _, server = self._server(["--run_gunicorn", "--keepalive_timeout", "9"]) - assert server.keepalive_seconds == GunicornConfig().keepalive - - @pytest.mark.parametrize("gunicorn", [True, False]) - def test_python_behind_rust_trusts_the_proxies_its_server_would_have( - self, gunicorn: bool, monkeypatch: pytest.MonkeyPatch + @pytest.mark.parametrize( + ("argv", "environ"), + [ + ([], {}), + (["--keepalive_timeout", "2"], {"GRACEFUL_SHUTDOWN_TIMEOUT": "12.5"}), + ([], {"KEEPALIVE_TIMEOUT": "7", "FORWARDED_ALLOW_IPS": "10.0.0.5"}), + (["--run_gunicorn", "--keepalive_timeout", "9"], {}), + ], + ) + def test_rust_follows_the_keepalive_drain_and_trust_python_serves_with( + self, argv: list[str], environ: dict[str, str], monkeypatch: pytest.MonkeyPatch ) -> None: - from gunicorn.config import Config as GunicornConfig - - from litellm.proxy.rust_forward.mode import InheritedFd, Supervised, rust_forward_mode - from litellm.proxy.rust_forward.platform import PrivateListener - from litellm.proxy.rust_forward.platform.posix import PosixPlatform - from litellm.proxy.rust_forward.supervisor.children import python_spec - - monkeypatch.delenv("FORWARDED_ALLOW_IPS", raising=False) - _, server = self._server(["--run_gunicorn"] if gunicorn else []) - mode: Final = rust_forward_mode( - python_spec( - PosixPlatform(), ("litellm",), PrivateListener(InheritedFd(17), 45555, (17,)), server - ).env_overrides - ) - trusted: Final = ( - GunicornConfig().forwarded_allow_ips - if gunicorn - else [uvicorn.Config(app="", log_config=None).forwarded_allow_ips] - ) - assert isinstance(mode, Supervised) - assert mode.trusted_proxies.split(",") == trusted - - def test_rust_drains_for_the_graceful_timeout_uvicorn_runs_with(self, monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "12.5") - _, server = self._server([]) - served: Final = ProxyInitializationHelpers.get_default_unvicorn_init_args("127.0.0.1", 45555) - assert server.graceful_timeout_seconds == served["timeout_graceful_shutdown"] + for name in ("KEEPALIVE_TIMEOUT", "GRACEFUL_SHUTDOWN_TIMEOUT", "FORWARDED_ALLOW_IPS"): + monkeypatch.delenv(name, raising=False) + for name, value in environ.items(): + monkeypatch.setenv(name, value) + params, server = self._server(argv) + assert ( + server.keepalive_seconds, + server.graceful_timeout_seconds, + server.trusted_proxies, + ) == self._served_by_python(params)