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