diff --git a/gateway/launch.py b/gateway/launch.py index 1adb2ef1034..e6b0c9dbeef 100644 --- a/gateway/launch.py +++ b/gateway/launch.py @@ -121,7 +121,8 @@ def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) - stop_with_supervisor, uvicorn_flags, ) - from litellm.proxy.rust_forward.supervisor import StartupError, supervise + from litellm.proxy.rust_forward.supervisor import supervise + from litellm.proxy.rust_forward.supervisor.preflight import StartupError mode: Final = rust_forward_mode(os.environ) match mode: diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 632a71ec5c8..afee156981b 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -1376,7 +1376,8 @@ def run_server( if forward_mode is Mode.SUPERVISOR and not any( (version, validate_config, health, test, skip_server_startup, check_migrations_against_existing_data) ): - from litellm.proxy.rust_forward.supervisor import StartupError, supervise + from litellm.proxy.rust_forward.supervisor import supervise + from litellm.proxy.rust_forward.supervisor.preflight import StartupError server: Final = ProxyInitializationHelpers.python_server( host, diff --git a/litellm/proxy/rust_forward/AGENTS.md b/litellm/proxy/rust_forward/AGENTS.md index 907b7dc1a3c..2084a19a66d 100644 --- a/litellm/proxy/rust_forward/AGENTS.md +++ b/litellm/proxy/rust_forward/AGENTS.md @@ -11,9 +11,9 @@ Every OS-specific operation lives behind `platform.Platform`, which `supervise` Windows support exists only for running the proxy locally from the CLI. Nobody deploys the proxy on Windows, and the Windows design leans on that. A port picked by the supervisor can be taken by another process before Python binds it, `CTRL_BREAK_EVENT` needs the console a CLI run always has, and gunicorn is refused. Revisit those choices before treating Windows as a deployment target -Each environment contract is written and read in one place. The Python child's is the `mode.Supervised` type itself: `mode.supervised_environment` serializes it into one variable and `mode.rust_forward_mode` validates it back, so a new field needs no parsing code. The trust it carries is resolved by the server library the entry point uses (`server.uvicorn_trusted_proxies` or gunicorn's config), never defaulted in this package. `supervisor.rust_spec` writes the `LITELLM_FORWARD_*` variables that `gateway-forward`'s `settings.rs` reads. Python owns every value Rust needs, so a new one gets a new variable in `rust_spec` +Each environment contract is written and read in one place. The Python child's is the `mode.Supervised` type itself: `mode.supervised_environment` serializes it into one variable and `mode.rust_forward_mode` validates it back, so a new field needs no parsing code. The trust it carries is resolved by the server library the entry point uses (`server.uvicorn_trusted_proxies` or gunicorn's config), never defaulted in this package. `supervisor.children.rust_spec` writes the `LITELLM_FORWARD_*` variables that `gateway-forward`'s `settings.rs` reads. Python owns every value Rust needs, so a new one gets a new variable in `children.rust_spec` -The supervisor is transparent to Python, so `python` with `LITELLM_RUST` behaves on signals and exits exactly as Python alone. Its whole lifecycle is `supervisor.step`, a pure function from a `Lifecycle` of facts and an event to the next facts and the actions to run, and `exit_status` reads the result from those facts. `_run` only executes actions, and the `contextlib.AsyncExitStack` it enters each child into owns their cleanup, last spawned first. Three rules decide everything there, and a change that breaks one of them breaks the equivalence `tests/integration/configuration/test_rust_forward_lifecycle.py` checks against Python alone +The supervisor is transparent to Python, so `python` with `LITELLM_RUST` behaves on signals and exits exactly as Python alone. The `supervisor` package splits by job: `preflight` decides whether forwarding can start, `children` says what each child runs with, `lifecycle` decides, and `runtime` acts. The whole lifecycle is `lifecycle.step`, a pure function from a `Lifecycle` of facts and an event to the next facts and the actions to run, and `exit_status` reads the result from those facts. `runtime` only executes actions, and the `contextlib.AsyncExitStack` it enters each child into owns their cleanup, last spawned first. Three rules decide everything there, and a change that breaks one of them breaks the equivalence `tests/integration/configuration/test_rust_forward_lifecycle.py` checks against Python alone - Signals pass through. Every signal in `Platform.relayed_signals` goes to Python unchanged, and the ones in `Platform.stop_signals` go to Rust first. The supervisor never reads meaning into a signal, so a repeated SIGINT reaches Python again and forces its exit as it would alone. Children run in their own session (`start_new_session`) only so that a terminal Ctrl-C reaches Python once, through the supervisor - Python's exit status is the result. When Python exits for any reason, Rust gets SIGTERM, and once Rust is gone the supervisor exits with Python's status, dying by the same signal when Python died by one. The one exception is Rust exiting while Python serves with no stop requested, which stops Python and exits 1 diff --git a/litellm/proxy/rust_forward/supervisor.py b/litellm/proxy/rust_forward/supervisor.py deleted file mode 100644 index d29db08dba0..00000000000 --- a/litellm/proxy/rust_forward/supervisor.py +++ /dev/null @@ -1,360 +0,0 @@ -import asyncio -import contextlib -import os -import shutil -import signal -import sys -import sysconfig -from collections.abc import AsyncGenerator, Callable, Coroutine, Mapping, Sequence -from dataclasses import dataclass, replace -from enum import Enum, auto -from functools import partial -from types import MappingProxyType -from typing import Final, Literal, TypeAlias - -from typing_extensions import assert_never - -from litellm.constants import GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS -from litellm.proxy.rust_forward.mode import RUST_HOP, Supervised, supervised_environment -from litellm.proxy.rust_forward.platform import Platform, PrivateListener, SpawnOptions, current_platform -from litellm.proxy.rust_forward.server import Listener, PythonServer, Server -from litellm.rust_bridge.configuration import ENV_NAME - -PROBE_INTERVAL_SECONDS: Final = 0.05 - - -def gateway_binary(platform: Platform, scripts_dir: str) -> str | None: - return shutil.which(platform.gateway_name, path=scripts_dir) - - -@dataclass(frozen=True, slots=True) -class StartupError: - reason: str - - -def startup_check(platform: Platform, server: PythonServer, binary: str | None) -> str | StartupError: - if server.server in (Server.HYPERCORN, Server.GRANIAN): - return StartupError(f"{ENV_NAME} does not support {server.server.value}") - if server.server not in platform.servers: - return StartupError(f"{ENV_NAME} does not support {server.server.value} on {platform.name}") - if server.listener is not Listener.TCP: - return StartupError(f"{ENV_NAME} does not support {server.listener.value} on the public port") - if binary is None: - return StartupError(f"{ENV_NAME} requires the litellm-gateway binary from the litellm-proxy package") - return binary - - -class Child(Enum): - RUST = "rust" - PYTHON = "python" - - -FAILURE: Final = 1 - - -class Phase(Enum): - NOT_STARTED = auto() - SPAWNED = auto() - LISTENING = auto() - - -@dataclass(frozen=True, slots=True) -class Exited: - status: int - - -@dataclass(frozen=True, slots=True) -class Lost: - """Rust exited while Python was serving and nobody had asked either to stop.""" - - status: int - - -PythonState: TypeAlias = Literal[Phase.SPAWNED, Phase.LISTENING] | Exited -RustState: TypeAlias = Literal[Phase.NOT_STARTED, Phase.SPAWNED] | Exited | Lost - - -@dataclass(frozen=True, slots=True) -class Lifecycle: - python: PythonState - rust: RustState - stop_requested: bool - - -@dataclass(frozen=True, slots=True) -class Caught: - signum: signal.Signals - - -class Probe(Enum): - PYTHON_LISTENING = auto() - - -@dataclass(frozen=True, slots=True) -class ChildExited: - child: Child - status: int - - -@dataclass(frozen=True, slots=True) -class Spawn: - child: Child - - -@dataclass(frozen=True, slots=True) -class Send: - child: Child - signum: signal.Signals - - -Event: TypeAlias = Caught | Literal[Probe.PYTHON_LISTENING] | ChildExited -Action: TypeAlias = Spawn | Send -Transition: TypeAlias = tuple[Lifecycle, tuple[Action, ...]] - - -def boot() -> Transition: - return Lifecycle(Phase.SPAWNED, Phase.NOT_STARTED, stop_requested=False), (Spawn(Child.PYTHON),) - - -def _relay(state: Lifecycle, signum: signal.Signals, stops: bool) -> Transition: - to_rust: Final = (Send(Child.RUST, signum),) if stops and state.rust is Phase.SPAWNED else () - to_python: Final = (Send(Child.PYTHON, signum),) if not isinstance(state.python, Exited) else () - return replace(state, stop_requested=state.stop_requested or stops), (*to_rust, *to_python) - - -# Python alone takes no connections until it has booted, so Rust only opens the public port once Python listens. -def _python_listening(state: Lifecycle) -> Transition: - if state.python is not Phase.SPAWNED: - return state, () - if state.stop_requested: - return replace(state, python=Phase.LISTENING), () - return replace(state, python=Phase.LISTENING, rust=Phase.SPAWNED), (Spawn(Child.RUST),) - - -def _python_exited(state: Lifecycle, status: int) -> Transition: - if isinstance(state.python, Exited): - return state, () - exited: Final = replace(state, python=Exited(status)) - return exited, ((Send(Child.RUST, signal.SIGTERM),) if state.rust is Phase.SPAWNED else ()) - - -def _rust_exited(state: Lifecycle, status: int) -> Transition: - if state.rust is not Phase.SPAWNED: - return state, () - if state.stop_requested or isinstance(state.python, Exited): - return replace(state, rust=Exited(status)), () - return replace(state, rust=Lost(status)), (Send(Child.PYTHON, signal.SIGTERM),) - - -def step(state: Lifecycle, event: Event, stop_signals: frozenset[signal.Signals]) -> Transition: - """The supervisor never gives a signal or an exit a meaning of its own. Every signal reaches Python as sent, - Python's exit status is the supervisor's, and Rust only follows Python, except that losing Rust stops Python.""" - match event: - case Caught(signum): - return _relay(state, signum, signum in stop_signals) - case Probe.PYTHON_LISTENING: - return _python_listening(state) - case ChildExited(child, status): - match child: - case Child.PYTHON: - return _python_exited(state, status) - case Child.RUST: - return _rust_exited(state, status) - case _: - assert_never(child) - case _: - assert_never(event) - - -def exit_status(state: Lifecycle) -> int | None: - match state.python, state.rust: - case Exited(status), Phase.NOT_STARTED | Exited(): - return status - case Exited(), Lost(): - return FAILURE - case _: - return None - - -@dataclass(frozen=True, slots=True) -class ChildSpec: - argv: tuple[str, ...] - env_overrides: Mapping[str, str] - spawn: SpawnOptions - - -def python_spec( - platform: Platform, command: Sequence[str], private: PrivateListener, server: PythonServer -) -> ChildSpec: - return ChildSpec( - argv=(*command, "--port", str(private.port)), - env_overrides=supervised_environment( - Supervised(private.listener, server.proxy_headers, server.trusted_proxies) - ), - spawn=platform.spawn_options(private.pass_fds), - ) - - -def rust_spec(platform: Platform, binary: str, upstream_port: int, server: PythonServer) -> ChildSpec: - return ChildSpec( - argv=(binary,), - env_overrides=MappingProxyType( - { - "LITELLM_FORWARD_UPSTREAM": f"http://{RUST_HOP}:{upstream_port}", - "LITELLM_FORWARD_HOST": server.host, - "LITELLM_FORWARD_PORT": str(server.port), - "LITELLM_FORWARD_UPSTREAM_KEEPALIVE": str(server.keepalive_seconds), - "LITELLM_FORWARD_SHUTDOWN_DRAIN": str( - server.graceful_timeout_seconds + GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS - ), - } - ), - spawn=platform.spawn_options(()), - ) - - -async def _spawn(spec: ChildSpec) -> asyncio.subprocess.Process: - return await asyncio.create_subprocess_exec( - *spec.argv, - stdin=asyncio.subprocess.PIPE, - env={**os.environ, **spec.env_overrides}, - pass_fds=spec.spawn.pass_fds, - start_new_session=spec.spawn.start_new_session, - creationflags=spec.spawn.creationflags, - ) - - -Launch: TypeAlias = Callable[[asyncio.Queue[Event]], contextlib.AbstractAsyncContextManager[asyncio.subprocess.Process]] - - -async def _report_exit(child: Child, process: asyncio.subprocess.Process, events: asyncio.Queue[Event]) -> None: - events.put_nowait(ChildExited(child, await process.wait())) - - -async def _accepts(port: int) -> bool: - try: - _, writer = await asyncio.wait_for(asyncio.open_connection(RUST_HOP, port), PROBE_INTERVAL_SECONDS) - except (OSError, asyncio.TimeoutError): - return False - writer.close() - return True - - -async def _report_once_listening(port: int, events: asyncio.Queue[Event]) -> None: - while not await _accepts(port): - await asyncio.sleep(PROBE_INTERVAL_SECONDS) - events.put_nowait(Probe.PYTHON_LISTENING) - - -@contextlib.asynccontextmanager -async def _in_background(work: Coroutine[None, None, None]) -> AsyncGenerator[None]: - task: Final = asyncio.create_task(work) - try: - yield - finally: - task.cancel() - await asyncio.gather(task, return_exceptions=True) - - -# A child's stdin is a pipe only the supervisor holds open, so the child sees EOF as soon as the supervisor is gone, -# however it went, and shuts down instead of outliving it. -@contextlib.asynccontextmanager -async def _launch( - platform: Platform, - child: Child, - spec: ChildSpec, - events: asyncio.Queue[Event], -) -> AsyncGenerator[asyncio.subprocess.Process]: - process: Final = await _spawn(spec) - waiter: Final = asyncio.create_task(_report_exit(child, process, events)) - try: - yield process - finally: - platform.stop_group(process) - await waiter - if process.stdin is not None: - process.stdin.close() - - -@contextlib.asynccontextmanager -async def _launch_python( - platform: Platform, spec: ChildSpec, port: int, events: asyncio.Queue[Event] -) -> AsyncGenerator[asyncio.subprocess.Process]: - async with ( - _launch(platform, Child.PYTHON, spec, events) as process, - _in_background(_report_once_listening(port, events)), - ): - yield process - - -async def _run(platform: Platform, events: asyncio.Queue[Event], launches: Mapping[Child, Launch]) -> int: - async with contextlib.AsyncExitStack() as stack: - children: Final[dict[Child, asyncio.subprocess.Process]] = {} # mutable-ok: indexes the stack's children - state, actions = boot() - while True: - for action in actions: - match action: - case Spawn(child): - children[child] = await stack.enter_async_context(launches[child](events)) - case Send(child, signum): - platform.signal_child(children[child], signum) - case _: - assert_never(action) - if (status := exit_status(state)) is not None: - return status - state, actions = step(state, await events.get(), platform.stop_signals) - - -async def _supervise(platform: Platform, binary: str, command: Sequence[str], server: PythonServer) -> int: - events: Final[asyncio.Queue[Event]] = asyncio.Queue() - platform.catch(asyncio.get_running_loop(), lambda signum: events.put_nowait(Caught(signum))) - with platform.private_listener() as private: - launches: Final[Mapping[Child, Launch]] = MappingProxyType( - { - Child.PYTHON: partial( - _launch_python, platform, python_spec(platform, command, private, server), private.port - ), - Child.RUST: partial(_launch, platform, Child.RUST, rust_spec(platform, binary, private.port, server)), - } - ) - return await _run(platform, events, launches) - - -def _exit_like(status: int) -> int: - """A child killed by a signal kills the supervisor with the same signal, as uvicorn re-raises the one it caught.""" - if status >= 0: - return status - signum: Final = signal.Signals(-status) - signal.signal(signum, signal.SIG_DFL) - signal.raise_signal(signum) - return 128 + signum - - -def _nothing_to_prepare() -> None: - pass - - -def supervise( - server: PythonServer, - command: Sequence[str], - before_start: Callable[[], None] = _nothing_to_prepare, - scripts_dir: str = sysconfig.get_path("scripts"), - platform: Platform | None = None, -) -> int | StartupError: - """Serve ``server``'s public port from Rust and run ``command`` behind it, or say why that cannot start. - - A refusal is returned for the caller to serve from Python alone, and ``before_start`` runs only once forwarding - is known to be possible, so a refused start leaves nothing behind. - """ - chosen: Final = current_platform() if platform is None else platform - checked: Final = startup_check(chosen, server, gateway_binary(chosen, scripts_dir)) - if isinstance(checked, StartupError): - sys.stderr.write(f"{checked.reason}, so Python serves the public port without Rust forwarding\n") - return checked - before_start() - try: - status: Final = asyncio.run(_supervise(chosen, checked, command, server)) - except OSError as error: - sys.stderr.write(f"Rust gateway startup failed: {error}\n") - return FAILURE - return _exit_like(status) diff --git a/litellm/proxy/rust_forward/supervisor/__init__.py b/litellm/proxy/rust_forward/supervisor/__init__.py new file mode 100644 index 00000000000..2819f837da3 --- /dev/null +++ b/litellm/proxy/rust_forward/supervisor/__init__.py @@ -0,0 +1,41 @@ +import asyncio +import sys +import sysconfig +from collections.abc import Callable, Sequence +from typing import Final + +from litellm.proxy.rust_forward.platform import Platform, current_platform +from litellm.proxy.rust_forward.server import PythonServer +from litellm.proxy.rust_forward.supervisor.lifecycle import FAILURE +from litellm.proxy.rust_forward.supervisor.preflight import StartupError, gateway_binary, startup_check +from litellm.proxy.rust_forward.supervisor.runtime import exit_like, run + + +def _nothing_to_prepare() -> None: + pass + + +def supervise( + server: PythonServer, + command: Sequence[str], + before_start: Callable[[], None] = _nothing_to_prepare, + scripts_dir: str = sysconfig.get_path("scripts"), + platform: Platform | None = None, +) -> int | StartupError: + """Serve ``server``'s public port from Rust and run ``command`` behind it, or say why that cannot start. + + A refusal is returned for the caller to serve from Python alone, and ``before_start`` runs only once forwarding + is known to be possible, so a refused start leaves nothing behind. + """ + chosen: Final = current_platform() if platform is None else platform + checked: Final = startup_check(chosen, server, gateway_binary(chosen, scripts_dir)) + if isinstance(checked, StartupError): + sys.stderr.write(f"{checked.reason}, so Python serves the public port without Rust forwarding\n") + return checked + before_start() + try: + status: Final = asyncio.run(run(chosen, checked, command, server)) + except OSError as error: + sys.stderr.write(f"Rust gateway startup failed: {error}\n") + return FAILURE + return exit_like(status) diff --git a/litellm/proxy/rust_forward/supervisor/children.py b/litellm/proxy/rust_forward/supervisor/children.py new file mode 100644 index 00000000000..c1da60519ed --- /dev/null +++ b/litellm/proxy/rust_forward/supervisor/children.py @@ -0,0 +1,45 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from types import MappingProxyType + +from litellm.constants import GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS +from litellm.proxy.rust_forward.mode import RUST_HOP, Supervised, supervised_environment +from litellm.proxy.rust_forward.platform import Platform, PrivateListener, SpawnOptions +from litellm.proxy.rust_forward.server import PythonServer + + +@dataclass(frozen=True, slots=True) +class ChildSpec: + argv: tuple[str, ...] + env_overrides: Mapping[str, str] + spawn: SpawnOptions + + +def python_spec( + platform: Platform, command: Sequence[str], private: PrivateListener, server: PythonServer +) -> ChildSpec: + return ChildSpec( + argv=(*command, "--port", str(private.port)), + env_overrides=supervised_environment( + Supervised(private.listener, server.proxy_headers, server.trusted_proxies) + ), + spawn=platform.spawn_options(private.pass_fds), + ) + + +def rust_spec(platform: Platform, binary: str, upstream_port: int, server: PythonServer) -> ChildSpec: + return ChildSpec( + argv=(binary,), + env_overrides=MappingProxyType( + { + "LITELLM_FORWARD_UPSTREAM": f"http://{RUST_HOP}:{upstream_port}", + "LITELLM_FORWARD_HOST": server.host, + "LITELLM_FORWARD_PORT": str(server.port), + "LITELLM_FORWARD_UPSTREAM_KEEPALIVE": str(server.keepalive_seconds), + "LITELLM_FORWARD_SHUTDOWN_DRAIN": str( + server.graceful_timeout_seconds + GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS + ), + } + ), + spawn=platform.spawn_options(()), + ) diff --git a/litellm/proxy/rust_forward/supervisor/lifecycle.py b/litellm/proxy/rust_forward/supervisor/lifecycle.py new file mode 100644 index 00000000000..e462c7df08a --- /dev/null +++ b/litellm/proxy/rust_forward/supervisor/lifecycle.py @@ -0,0 +1,138 @@ +import signal +from dataclasses import dataclass, replace +from enum import Enum, auto +from typing import Final, Literal, TypeAlias + +from typing_extensions import assert_never + + +class Child(Enum): + RUST = "rust" + PYTHON = "python" + + +FAILURE: Final = 1 + + +class Phase(Enum): + NOT_STARTED = auto() + SPAWNED = auto() + LISTENING = auto() + + +@dataclass(frozen=True, slots=True) +class Exited: + status: int + + +@dataclass(frozen=True, slots=True) +class Lost: + """Rust exited while Python was serving and nobody had asked either to stop.""" + + status: int + + +PythonState: TypeAlias = Literal[Phase.SPAWNED, Phase.LISTENING] | Exited +RustState: TypeAlias = Literal[Phase.NOT_STARTED, Phase.SPAWNED] | Exited | Lost + + +@dataclass(frozen=True, slots=True) +class Lifecycle: + python: PythonState + rust: RustState + stop_requested: bool + + +@dataclass(frozen=True, slots=True) +class Caught: + signum: signal.Signals + + +class Probe(Enum): + PYTHON_LISTENING = auto() + + +@dataclass(frozen=True, slots=True) +class ChildExited: + child: Child + status: int + + +@dataclass(frozen=True, slots=True) +class Spawn: + child: Child + + +@dataclass(frozen=True, slots=True) +class Send: + child: Child + signum: signal.Signals + + +Event: TypeAlias = Caught | Literal[Probe.PYTHON_LISTENING] | ChildExited +Action: TypeAlias = Spawn | Send +Transition: TypeAlias = tuple[Lifecycle, tuple[Action, ...]] + + +def boot() -> Transition: + return Lifecycle(Phase.SPAWNED, Phase.NOT_STARTED, stop_requested=False), (Spawn(Child.PYTHON),) + + +def _relay(state: Lifecycle, signum: signal.Signals, stops: bool) -> Transition: + to_rust: Final = (Send(Child.RUST, signum),) if stops and state.rust is Phase.SPAWNED else () + to_python: Final = (Send(Child.PYTHON, signum),) if not isinstance(state.python, Exited) else () + return replace(state, stop_requested=state.stop_requested or stops), (*to_rust, *to_python) + + +# Python alone takes no connections until it has booted, so Rust only opens the public port once Python listens. +def _python_listening(state: Lifecycle) -> Transition: + if state.python is not Phase.SPAWNED: + return state, () + if state.stop_requested: + return replace(state, python=Phase.LISTENING), () + return replace(state, python=Phase.LISTENING, rust=Phase.SPAWNED), (Spawn(Child.RUST),) + + +def _python_exited(state: Lifecycle, status: int) -> Transition: + if isinstance(state.python, Exited): + return state, () + exited: Final = replace(state, python=Exited(status)) + return exited, ((Send(Child.RUST, signal.SIGTERM),) if state.rust is Phase.SPAWNED else ()) + + +def _rust_exited(state: Lifecycle, status: int) -> Transition: + if state.rust is not Phase.SPAWNED: + return state, () + if state.stop_requested or isinstance(state.python, Exited): + return replace(state, rust=Exited(status)), () + return replace(state, rust=Lost(status)), (Send(Child.PYTHON, signal.SIGTERM),) + + +def step(state: Lifecycle, event: Event, stop_signals: frozenset[signal.Signals]) -> Transition: + """The supervisor never gives a signal or an exit a meaning of its own. Every signal reaches Python as sent, + Python's exit status is the supervisor's, and Rust only follows Python, except that losing Rust stops Python.""" + match event: + case Caught(signum): + return _relay(state, signum, signum in stop_signals) + case Probe.PYTHON_LISTENING: + return _python_listening(state) + case ChildExited(child, status): + match child: + case Child.PYTHON: + return _python_exited(state, status) + case Child.RUST: + return _rust_exited(state, status) + case _: + assert_never(child) + case _: + assert_never(event) + + +def exit_status(state: Lifecycle) -> int | None: + match state.python, state.rust: + case Exited(status), Phase.NOT_STARTED | Exited(): + return status + case Exited(), Lost(): + return FAILURE + case _: + return None diff --git a/litellm/proxy/rust_forward/supervisor/preflight.py b/litellm/proxy/rust_forward/supervisor/preflight.py new file mode 100644 index 00000000000..cfd5dd948e6 --- /dev/null +++ b/litellm/proxy/rust_forward/supervisor/preflight.py @@ -0,0 +1,27 @@ +import shutil +from dataclasses import dataclass + +from litellm.proxy.rust_forward.platform import Platform +from litellm.proxy.rust_forward.server import Listener, PythonServer, Server +from litellm.rust_bridge.configuration import ENV_NAME + + +@dataclass(frozen=True, slots=True) +class StartupError: + reason: str + + +def gateway_binary(platform: Platform, scripts_dir: str) -> str | None: + return shutil.which(platform.gateway_name, path=scripts_dir) + + +def startup_check(platform: Platform, server: PythonServer, binary: str | None) -> str | StartupError: + if server.server in (Server.HYPERCORN, Server.GRANIAN): + return StartupError(f"{ENV_NAME} does not support {server.server.value}") + if server.server not in platform.servers: + return StartupError(f"{ENV_NAME} does not support {server.server.value} on {platform.name}") + if server.listener is not Listener.TCP: + return StartupError(f"{ENV_NAME} does not support {server.listener.value} on the public port") + if binary is None: + return StartupError(f"{ENV_NAME} requires the litellm-gateway binary from the litellm-proxy package") + return binary diff --git a/litellm/proxy/rust_forward/supervisor/runtime.py b/litellm/proxy/rust_forward/supervisor/runtime.py new file mode 100644 index 00000000000..1e8aced9077 --- /dev/null +++ b/litellm/proxy/rust_forward/supervisor/runtime.py @@ -0,0 +1,142 @@ +import asyncio +import contextlib +import os +import signal +from collections.abc import AsyncGenerator, Callable, Coroutine, Mapping, Sequence +from functools import partial +from types import MappingProxyType +from typing import Final, TypeAlias + +from typing_extensions import assert_never + +from litellm.proxy.rust_forward.mode import RUST_HOP +from litellm.proxy.rust_forward.platform import Platform +from litellm.proxy.rust_forward.server import PythonServer +from litellm.proxy.rust_forward.supervisor.children import ChildSpec, python_spec, rust_spec +from litellm.proxy.rust_forward.supervisor.lifecycle import ( + Caught, + Child, + ChildExited, + Event, + Probe, + Send, + Spawn, + boot, + exit_status, + step, +) + +PROBE_INTERVAL_SECONDS: Final = 0.05 + +Launch: TypeAlias = Callable[[asyncio.Queue[Event]], contextlib.AbstractAsyncContextManager[asyncio.subprocess.Process]] + + +async def run(platform: Platform, binary: str, command: Sequence[str], server: PythonServer) -> int: + """Run Python and Rust under ``lifecycle.step`` until it decides the exit status.""" + events: Final[asyncio.Queue[Event]] = asyncio.Queue() + platform.catch(asyncio.get_running_loop(), lambda signum: events.put_nowait(Caught(signum))) + with platform.private_listener() as private: + launches: Final[Mapping[Child, Launch]] = MappingProxyType( + { + Child.PYTHON: partial( + _launch_python, platform, python_spec(platform, command, private, server), private.port + ), + Child.RUST: partial(_launch, platform, Child.RUST, rust_spec(platform, binary, private.port, server)), + } + ) + return await _drive(platform, events, launches) + + +def exit_like(status: int) -> int: + """A child killed by a signal kills the supervisor with the same signal, as uvicorn re-raises the one it caught.""" + if status >= 0: + return status + signum: Final = signal.Signals(-status) + signal.signal(signum, signal.SIG_DFL) + signal.raise_signal(signum) + return 128 + signum + + +async def _drive(platform: Platform, events: asyncio.Queue[Event], launches: Mapping[Child, Launch]) -> int: + async with contextlib.AsyncExitStack() as stack: + children: Final[dict[Child, asyncio.subprocess.Process]] = {} # mutable-ok: indexes the stack's children + state, actions = boot() + while True: + for action in actions: + match action: + case Spawn(child): + children[child] = await stack.enter_async_context(launches[child](events)) + case Send(child, signum): + platform.signal_child(children[child], signum) + case _: + assert_never(action) + if (status := exit_status(state)) is not None: + return status + state, actions = step(state, await events.get(), platform.stop_signals) + + +# A child's stdin is a pipe only the supervisor holds open, so the child sees EOF as soon as the supervisor is gone, +# however it went, and shuts down instead of outliving it. +@contextlib.asynccontextmanager +async def _launch( + platform: Platform, + child: Child, + spec: ChildSpec, + events: asyncio.Queue[Event], +) -> AsyncGenerator[asyncio.subprocess.Process]: + process: Final = await asyncio.create_subprocess_exec( + *spec.argv, + stdin=asyncio.subprocess.PIPE, + env={**os.environ, **spec.env_overrides}, + pass_fds=spec.spawn.pass_fds, + start_new_session=spec.spawn.start_new_session, + creationflags=spec.spawn.creationflags, + ) + waiter: Final = asyncio.create_task(_report_exit(child, process, events)) + try: + yield process + finally: + platform.stop_group(process) + await waiter + if process.stdin is not None: + process.stdin.close() + + +@contextlib.asynccontextmanager +async def _launch_python( + platform: Platform, spec: ChildSpec, port: int, events: asyncio.Queue[Event] +) -> AsyncGenerator[asyncio.subprocess.Process]: + async with ( + _launch(platform, Child.PYTHON, spec, events) as process, + _in_background(_report_once_listening(port, events)), + ): + yield process + + +async def _report_exit(child: Child, process: asyncio.subprocess.Process, events: asyncio.Queue[Event]) -> None: + events.put_nowait(ChildExited(child, await process.wait())) + + +async def _report_once_listening(port: int, events: asyncio.Queue[Event]) -> None: + while not await _accepts(port): + await asyncio.sleep(PROBE_INTERVAL_SECONDS) + events.put_nowait(Probe.PYTHON_LISTENING) + + +async def _accepts(port: int) -> bool: + try: + _, writer = await asyncio.wait_for(asyncio.open_connection(RUST_HOP, port), PROBE_INTERVAL_SECONDS) + except (OSError, asyncio.TimeoutError): + return False + writer.close() + return True + + +@contextlib.asynccontextmanager +async def _in_background(work: Coroutine[None, None, None]) -> AsyncGenerator[None]: + task: Final = asyncio.create_task(work) + try: + yield + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) diff --git a/tests/unit/proxy/rust_forward/platform/__init__.py b/tests/unit/proxy/rust_forward/platform/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/proxy/rust_forward/platform/test_windows.py b/tests/unit/proxy/rust_forward/platform/test_windows.py new file mode 100644 index 00000000000..986fb79ed53 --- /dev/null +++ b/tests/unit/proxy/rust_forward/platform/test_windows.py @@ -0,0 +1,13 @@ +import signal +import subprocess +import sys + +import pytest + +from litellm.proxy.rust_forward.platform.windows import CREATE_NEW_PROCESS_GROUP, CTRL_BREAK_EVENT, SIGBREAK + + +@pytest.mark.skipif(sys.platform != "win32", reason="these names exist only on Windows") +def test_windows_constants_match_the_stdlib() -> None: + assert (CTRL_BREAK_EVENT, SIGBREAK) == (signal.CTRL_BREAK_EVENT, signal.SIGBREAK) + assert CREATE_NEW_PROCESS_GROUP == subprocess.CREATE_NEW_PROCESS_GROUP diff --git a/tests/unit/proxy/rust_forward/supervisor/__init__.py b/tests/unit/proxy/rust_forward/supervisor/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/proxy/rust_forward/supervisor/_support.py b/tests/unit/proxy/rust_forward/supervisor/_support.py new file mode 100644 index 00000000000..cc4287f5ea7 --- /dev/null +++ b/tests/unit/proxy/rust_forward/supervisor/_support.py @@ -0,0 +1,7 @@ +from litellm.proxy.rust_forward.server import Listener, PythonServer, Server + + +def python_server( + server: Server = Server.UVICORN, listener: Listener = Listener.TCP, trusted_proxies: str = "127.0.0.1" +) -> PythonServer: + return PythonServer(server, listener, "0.0.0.0", 4000, 2, 2.5, False, trusted_proxies) diff --git a/tests/unit/proxy/rust_forward/supervisor/test_children.py b/tests/unit/proxy/rust_forward/supervisor/test_children.py new file mode 100644 index 00000000000..fc5d8cf703a --- /dev/null +++ b/tests/unit/proxy/rust_forward/supervisor/test_children.py @@ -0,0 +1,53 @@ +import socket +from typing import Final + +from litellm.constants import GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS +from litellm.proxy.proxy_cli import run_server +from litellm.proxy.rust_forward.mode import AssignedPort, InheritedFd, Supervised, rust_forward_mode +from litellm.proxy.rust_forward.platform import PrivateListener +from litellm.proxy.rust_forward.platform.posix import PosixPlatform +from litellm.proxy.rust_forward.platform.windows import CREATE_NEW_PROCESS_GROUP, WindowsPlatform +from litellm.proxy.rust_forward.supervisor.children import python_spec, rust_spec +from tests.unit.proxy.rust_forward.supervisor._support import python_server + +POSIX: Final = PosixPlatform() +WINDOWS: Final = WindowsPlatform() + + +def test_python_child_serves_the_private_socket_with_the_operators_trust() -> None: + private: Final = PrivateListener(InheritedFd(17), 45555, (17,)) + spec: Final = python_spec(POSIX, ("python",), private, python_server(trusted_proxies="10.0.0.5")) + assert rust_forward_mode(spec.env_overrides) == Supervised(InheritedFd(17), False, "10.0.0.5") + assert (spec.spawn.pass_fds, spec.spawn.start_new_session) == ((17,), True) + + +def test_windows_python_child_binds_the_assigned_port_in_its_own_process_group() -> None: + with WINDOWS.private_listener() as private: + pass + spec: Final = python_spec(WINDOWS, ("python",), private, python_server()) + assert rust_forward_mode(spec.env_overrides) == Supervised(AssignedPort(private.port), False, "127.0.0.1") + assert spec.argv[-2:] == ("--port", str(private.port)) + assert (spec.spawn.pass_fds, spec.spawn.creationflags) == ((), CREATE_NEW_PROCESS_GROUP) + with socket.create_server(("127.0.0.1", private.port)): + pass + + +def test_python_child_keeps_the_public_host_and_its_options() -> None: + arguments: Final = ("--host", "0.0.0.0", "--port", "4000", "--num_workers", "2", "--reload") + spec: Final = python_spec(POSIX, arguments, PrivateListener(InheritedFd(17), 45555, (17,)), python_server()) + params: Final = run_server.make_context("litellm", list(spec.argv)).params + assert (params["host"], params["port"]) == ("0.0.0.0", 45555) + assert (params["num_workers"], params["reload"]) == (2, True) + + +def test_rust_child_serves_public_port_and_forwards_to_private_listener() -> None: + spec: Final = rust_spec(POSIX, "gateway", 45555, python_server()) + assert spec.argv == ("gateway",) + assert spec.spawn.pass_fds == () + assert dict(spec.env_overrides) == { + "LITELLM_FORWARD_UPSTREAM": "http://127.0.0.1:45555", + "LITELLM_FORWARD_HOST": "0.0.0.0", + "LITELLM_FORWARD_PORT": "4000", + "LITELLM_FORWARD_UPSTREAM_KEEPALIVE": "2", + "LITELLM_FORWARD_SHUTDOWN_DRAIN": str(2.5 + GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS), + } diff --git a/tests/unit/proxy/rust_forward/supervisor/test_lifecycle.py b/tests/unit/proxy/rust_forward/supervisor/test_lifecycle.py new file mode 100644 index 00000000000..a71fb6b335e --- /dev/null +++ b/tests/unit/proxy/rust_forward/supervisor/test_lifecycle.py @@ -0,0 +1,167 @@ +import functools +import signal +from dataclasses import dataclass, replace +from typing import Final + +import pytest +from hypothesis import example, given +from hypothesis import strategies as st + +from litellm.proxy.rust_forward.platform.posix import PosixPlatform +from litellm.proxy.rust_forward.supervisor.lifecycle import ( + FAILURE, + Action, + Caught, + Child, + ChildExited, + Event, + Lifecycle, + Phase, + Probe, + Send, + Spawn, + boot, + exit_status, + step, +) + +TERM: Final = signal.SIGTERM +INT: Final = signal.SIGINT +POSIX: Final = PosixPlatform() + + +@dataclass(frozen=True, slots=True) +class World: + """What the operating system knows about the children, kept independently of ``Lifecycle``.""" + + alive: frozenset[Child] = frozenset({Child.PYTHON}) + spawned: frozenset[Child] = frozenset({Child.PYTHON}) + reported_listening: bool = False + stop_requested: bool = False + python_status: int | None = None + rust_lost: bool = False + + +def _expected_exit(world: World) -> int | None: + if world.alive or world.python_status is None: + return None + return FAILURE if world.rust_lost else world.python_status + + +def _expected_relay(world: World, signum: signal.Signals) -> tuple[Action, ...]: + to_rust: Final = (Send(Child.RUST, signum),) if signum in STOPS and Child.RUST in world.alive else () + to_python: Final = (Send(Child.PYTHON, signum),) if Child.PYTHON in world.alive else () + return (*to_rust, *to_python) + + +def _expected_on_exit(world: World, child: Child) -> tuple[Action, ...]: + match child: + case Child.PYTHON: + return (Send(Child.RUST, TERM),) if Child.RUST in world.alive else () + case Child.RUST: + serving: Final = Child.PYTHON in world.alive and not world.stop_requested + return (Send(Child.PYTHON, TERM),) if serving else () + + +def _after(world: World, event: Event) -> World: + match event: + case Caught(signum): + return replace(world, stop_requested=world.stop_requested or signum in STOPS) + case Probe.PYTHON_LISTENING: + return replace(world, reported_listening=world.reported_listening or Child.PYTHON in world.alive) + case ChildExited(child, status) if child in world.alive: + lost: Final = child is Child.RUST and Child.PYTHON in world.alive and not world.stop_requested + return replace( + world, + alive=world.alive - {child}, + python_status=status if child is Child.PYTHON else world.python_status, + rust_lost=world.rust_lost or lost, + ) + case ChildExited(): + return world + + +def _spawned(world: World, actions: tuple[Action, ...]) -> World: + started: Final = frozenset(action.child for action in actions if isinstance(action, Spawn)) + return replace(world, alive=world.alive | started, spawned=world.spawned | started) + + +SIGNALS: Final = POSIX.relayed_signals +STOPS: Final = POSIX.stop_signals +EVENTS: Final = st.lists( + st.one_of( + st.sampled_from(SIGNALS).map(Caught), + st.just(Probe.PYTHON_LISTENING), + st.builds(ChildExited, st.sampled_from(Child), st.integers(-15, 3)), + ), + max_size=40, +) + + +@given(EVENTS) +@example([Caught(TERM), Caught(signal.SIGHUP), Probe.PYTHON_LISTENING]) +@example([Probe.PYTHON_LISTENING, Caught(TERM), Caught(signal.SIGWINCH), ChildExited(Child.RUST, 0)]) +def test_lifecycle_relays_every_signal_and_ends_with_pythons_status(events: list[Event]) -> None: + state, actions = boot() + world = World() + assert actions == (Spawn(Child.PYTHON),) + for event in events: + before = _after(world, event) + state, actions = step(state, event, STOPS) + match event: + case Caught(signum): + assert actions == _expected_relay(before, signum) + case Probe.PYTHON_LISTENING: + starts_rust = Child.PYTHON in world.alive and not world.reported_listening and not world.stop_requested + assert actions == ((Spawn(Child.RUST),) if starts_rust else ()) + case ChildExited(child) if child in world.alive: + assert actions == _expected_on_exit(world, child) + case ChildExited(): + assert actions == () + assert all(action.child in before.alive for action in actions if isinstance(action, Send)) + assert all(action.child not in before.spawned for action in actions if isinstance(action, Spawn)) + world = _spawned(before, actions) + assert exit_status(state) == _expected_exit(world) + + +@given(EVENTS) +def test_rust_starts_at_most_once_and_only_while_python_listens_unstopped(events: list[Event]) -> None: + def advance(acc: tuple[Lifecycle, tuple[tuple[Event, Lifecycle, Action], ...]], event: Event): + state, seen = acc + following, actions = step(state, event, STOPS) + return following, (*seen, *((event, state, action) for action in actions)) + + _, history = functools.reduce(advance, events, (boot()[0], ())) + spawns: Final = tuple((event, state) for event, state, action in history if isinstance(action, Spawn)) + assert len(spawns) <= 1 + assert all( + event is Probe.PYTHON_LISTENING and state.python is Phase.SPAWNED and not state.stop_requested + for event, state in spawns + ) + + +@pytest.mark.parametrize( + ("events", "expected"), + [ + ((Probe.PYTHON_LISTENING, Caught(TERM), ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, -15)), -15), + ( + ( + Probe.PYTHON_LISTENING, + Caught(INT), + Caught(INT), + ChildExited(Child.PYTHON, 1), + ChildExited(Child.RUST, 0), + ), + 1, + ), + ( + (Probe.PYTHON_LISTENING, Caught(signal.SIGHUP), ChildExited(Child.PYTHON, -1), ChildExited(Child.RUST, 0)), + -1, + ), + ((Probe.PYTHON_LISTENING, ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, 0)), FAILURE), + ((ChildExited(Child.PYTHON, 3),), 3), + ], +) +def test_supervisor_exits_as_python_did_unless_rust_was_lost(events: tuple[Event, ...], expected: int) -> None: + final: Final = functools.reduce(lambda state, event: step(state, event, STOPS)[0], events, boot()[0]) + assert exit_status(final) == expected diff --git a/tests/unit/proxy/rust_forward/supervisor/test_preflight.py b/tests/unit/proxy/rust_forward/supervisor/test_preflight.py new file mode 100644 index 00000000000..71899262c30 --- /dev/null +++ b/tests/unit/proxy/rust_forward/supervisor/test_preflight.py @@ -0,0 +1,86 @@ +import sys +from pathlib import Path +from typing import Final + +import pytest + +from litellm.proxy.rust_forward.platform.posix import PosixPlatform +from litellm.proxy.rust_forward.platform.windows import WindowsPlatform +from litellm.proxy.rust_forward.server import Listener, PythonServer, Server +from litellm.proxy.rust_forward.supervisor import supervise +from litellm.proxy.rust_forward.supervisor.preflight import StartupError, gateway_binary, startup_check +from tests.unit.proxy.rust_forward.supervisor._support import python_server + +POSIX: Final = PosixPlatform() +WINDOWS: Final = WindowsPlatform() +POSIX_ONLY: Final = pytest.mark.skipif(sys.platform == "win32", reason="runs a POSIX shebang gateway") + + +@pytest.mark.parametrize( + ("server", "reason"), + [ + (python_server(Server.HYPERCORN), "LITELLM_RUST does not support Hypercorn"), + (python_server(Server.GRANIAN), "LITELLM_RUST does not support Granian"), + (python_server(listener=Listener.TLS), "LITELLM_RUST does not support TLS on the public port"), + ( + python_server(listener=Listener.UNIX_SOCKET), + "LITELLM_RUST does not support a Unix socket on the public port", + ), + ( + python_server(Server.GUNICORN, Listener.INHERITED_FD), + "LITELLM_RUST does not support an inherited file descriptor on the public port", + ), + ], +) +def test_unsupported_servers_are_refused_before_anything_starts( + server: PythonServer, reason: str, capsys: pytest.CaptureFixture[str] +) -> None: + assert startup_check(POSIX, server, "gateway") == StartupError(reason) + prepared: Final[list[bool]] = [] + refused: Final = supervise(server, ("python",), before_start=lambda: prepared.append(True), platform=POSIX) + assert refused == StartupError(reason) + assert (capsys.readouterr().err, prepared) == ( + f"{reason}, so Python serves the public port without Rust forwarding\n", + [], + ) + + +@pytest.mark.parametrize("server", [Server.UVICORN, Server.GUNICORN]) +def test_plain_tcp_uvicorn_and_gunicorn_can_be_forwarded(server: Server) -> None: + assert startup_check(POSIX, python_server(server), "gateway") == "gateway" + + +def test_windows_forwards_uvicorn_and_refuses_gunicorn() -> None: + assert startup_check(WINDOWS, python_server(), "gateway") == "gateway" + assert startup_check(WINDOWS, python_server(Server.GUNICORN), "gateway") == StartupError( + "LITELLM_RUST does not support gunicorn on Windows" + ) + + +def test_missing_binary_is_refused_before_anything_starts(tmp_path: Path) -> None: + prepared: Final[list[bool]] = [] + refused: Final = supervise( + python_server(), + ("python",), + before_start=lambda: prepared.append(True), + scripts_dir=str(tmp_path), + platform=POSIX, + ) + assert (refused, prepared) == ( + StartupError("LITELLM_RUST requires the litellm-gateway binary from the litellm-proxy package"), + [], + ) + + +@POSIX_ONLY +def test_binary_is_found_only_in_the_environment_scripts_dir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + scripts: Final = tmp_path / "scripts" + elsewhere: Final = tmp_path / "elsewhere" + scripts.mkdir() + elsewhere.mkdir() + binary: Final = elsewhere / "litellm-gateway" + binary.write_text("#!/bin/sh\nexit 0\n") + binary.chmod(0o755) + monkeypatch.setenv("PATH", str(elsewhere)) + assert gateway_binary(POSIX, str(scripts)) is None + assert gateway_binary(POSIX, str(elsewhere)) == str(binary) diff --git a/tests/unit/proxy/rust_forward/supervisor/test_runtime.py b/tests/unit/proxy/rust_forward/supervisor/test_runtime.py new file mode 100644 index 00000000000..ac47b3975ce --- /dev/null +++ b/tests/unit/proxy/rust_forward/supervisor/test_runtime.py @@ -0,0 +1,45 @@ +import sys +from pathlib import Path +from typing import Final + +import pytest + +from litellm.proxy.rust_forward.platform import Platform, current_platform +from litellm.proxy.rust_forward.platform.windows import WindowsPlatform +from litellm.proxy.rust_forward.supervisor import supervise +from tests.unit.proxy.rust_forward.supervisor._support import python_server + + +PYTHON_CHILD: Final = """ +import json, os, pathlib, socket, time +time.sleep(0.5) +pathlib.Path(os.environ["MARKERS"], "listening").touch() +private = json.loads(os.environ["_LITELLM_GATEWAY_SUPERVISED"])["listener"] +listener = socket.socket(fileno=private["fd"]) if "fd" in private else socket.create_server(("127.0.0.1", private["port"])) +listener.listen() +time.sleep(10) +""" +RUST_CHILD: Final = """ +import os, pathlib +markers = pathlib.Path(os.environ["MARKERS"]) +(markers / "rust_saw_python_listening").write_text(str((markers / "listening").exists())) +""" + + +def _fake_gateway(scripts: Path) -> Platform: + if sys.platform != "win32": + gateway: Final = scripts / "litellm-gateway" + gateway.write_text(f"#!{sys.executable} -I\n{RUST_CHILD}") + gateway.chmod(0o755) + return current_platform() + script: Final = scripts / "rust_child.py" + script.write_text(RUST_CHILD) + (scripts / "litellm-gateway.cmd").write_text(f'@"{sys.executable}" -I "{script}"\r\n') + return WindowsPlatform(gateway_name="litellm-gateway.cmd") + + +def test_rust_opens_the_public_port_only_once_python_listens(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("MARKERS", str(tmp_path)) + command: Final = (sys.executable, "-I", "-c", PYTHON_CHILD) + supervise(python_server(), command, scripts_dir=str(tmp_path), platform=_fake_gateway(tmp_path)) + assert (tmp_path / "rust_saw_python_listening").read_text() == "True" diff --git a/tests/unit/proxy/rust_forward/test_supervisor.py b/tests/unit/proxy/rust_forward/test_supervisor.py deleted file mode 100644 index bb14443c96e..00000000000 --- a/tests/unit/proxy/rust_forward/test_supervisor.py +++ /dev/null @@ -1,339 +0,0 @@ -import functools -import signal -import socket -import subprocess -import sys -from dataclasses import dataclass, replace -from pathlib import Path -from typing import Final - -import pytest -from hypothesis import example, given -from hypothesis import strategies as st - -from litellm.constants import GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS -from litellm.proxy.proxy_cli import run_server -from litellm.proxy.rust_forward.mode import AssignedPort, InheritedFd, Supervised, rust_forward_mode -from litellm.proxy.rust_forward.platform import Platform, PrivateListener, current_platform -from litellm.proxy.rust_forward.platform.posix import PosixPlatform -from litellm.proxy.rust_forward.platform.windows import ( - CREATE_NEW_PROCESS_GROUP, - CTRL_BREAK_EVENT, - SIGBREAK, - WindowsPlatform, -) -from litellm.proxy.rust_forward.server import Listener, PythonServer, Server -from litellm.proxy.rust_forward.supervisor import ( - FAILURE, - Action, - Caught, - Child, - ChildExited, - Event, - Lifecycle, - Phase, - Probe, - Send, - Spawn, - StartupError, - boot, - exit_status, - gateway_binary, - python_spec, - rust_spec, - startup_check, - step, - supervise, -) - -TERM: Final = signal.SIGTERM -INT: Final = signal.SIGINT -POSIX: Final = PosixPlatform() -WINDOWS: Final = WindowsPlatform() -POSIX_ONLY: Final = pytest.mark.skipif(sys.platform == "win32", reason="runs a POSIX shebang gateway") - - -def _server( - server: Server = Server.UVICORN, listener: Listener = Listener.TCP, trusted_proxies: str = "127.0.0.1" -) -> PythonServer: - return PythonServer(server, listener, "0.0.0.0", 4000, 2, 2.5, False, trusted_proxies) - - -def test_python_child_serves_the_private_socket_with_the_operators_trust() -> None: - private: Final = PrivateListener(InheritedFd(17), 45555, (17,)) - spec: Final = python_spec(POSIX, ("python",), private, _server(trusted_proxies="10.0.0.5")) - assert rust_forward_mode(spec.env_overrides) == Supervised(InheritedFd(17), False, "10.0.0.5") - assert (spec.spawn.pass_fds, spec.spawn.start_new_session) == ((17,), True) - - -def test_windows_python_child_binds_the_assigned_port_in_its_own_process_group() -> None: - with WINDOWS.private_listener() as private: - pass - spec: Final = python_spec(WINDOWS, ("python",), private, _server()) - assert rust_forward_mode(spec.env_overrides) == Supervised(AssignedPort(private.port), False, "127.0.0.1") - assert spec.argv[-2:] == ("--port", str(private.port)) - assert (spec.spawn.pass_fds, spec.spawn.creationflags) == ((), CREATE_NEW_PROCESS_GROUP) - with socket.create_server(("127.0.0.1", private.port)): - pass - - -@pytest.mark.skipif(sys.platform != "win32", reason="these names exist only on Windows") -def test_windows_constants_match_the_stdlib() -> None: - assert (CTRL_BREAK_EVENT, SIGBREAK) == (signal.CTRL_BREAK_EVENT, signal.SIGBREAK) - assert CREATE_NEW_PROCESS_GROUP == subprocess.CREATE_NEW_PROCESS_GROUP - - -def test_python_child_keeps_the_public_host_and_its_options() -> None: - arguments: Final = ("--host", "0.0.0.0", "--port", "4000", "--num_workers", "2", "--reload") - spec: Final = python_spec(POSIX, arguments, PrivateListener(InheritedFd(17), 45555, (17,)), _server()) - params: Final = run_server.make_context("litellm", list(spec.argv)).params - assert (params["host"], params["port"]) == ("0.0.0.0", 45555) - assert (params["num_workers"], params["reload"]) == (2, True) - - -def test_rust_child_serves_public_port_and_forwards_to_private_listener() -> None: - spec: Final = rust_spec(POSIX, "gateway", 45555, _server()) - assert spec.argv == ("gateway",) - assert spec.spawn.pass_fds == () - assert dict(spec.env_overrides) == { - "LITELLM_FORWARD_UPSTREAM": "http://127.0.0.1:45555", - "LITELLM_FORWARD_HOST": "0.0.0.0", - "LITELLM_FORWARD_PORT": "4000", - "LITELLM_FORWARD_UPSTREAM_KEEPALIVE": "2", - "LITELLM_FORWARD_SHUTDOWN_DRAIN": str(2.5 + GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS), - } - - -@pytest.mark.parametrize( - ("server", "reason"), - [ - (_server(Server.HYPERCORN), "LITELLM_RUST does not support Hypercorn"), - (_server(Server.GRANIAN), "LITELLM_RUST does not support Granian"), - (_server(listener=Listener.TLS), "LITELLM_RUST does not support TLS on the public port"), - (_server(listener=Listener.UNIX_SOCKET), "LITELLM_RUST does not support a Unix socket on the public port"), - ( - _server(Server.GUNICORN, Listener.INHERITED_FD), - "LITELLM_RUST does not support an inherited file descriptor on the public port", - ), - ], -) -def test_unsupported_servers_are_refused_before_anything_starts( - server: PythonServer, reason: str, capsys: pytest.CaptureFixture[str] -) -> None: - assert startup_check(POSIX, server, "gateway") == StartupError(reason) - prepared: Final[list[bool]] = [] - refused: Final = supervise(server, ("python",), before_start=lambda: prepared.append(True), platform=POSIX) - assert refused == StartupError(reason) - assert (capsys.readouterr().err, prepared) == ( - f"{reason}, so Python serves the public port without Rust forwarding\n", - [], - ) - - -@pytest.mark.parametrize("server", [Server.UVICORN, Server.GUNICORN]) -def test_plain_tcp_uvicorn_and_gunicorn_can_be_forwarded(server: Server) -> None: - assert startup_check(POSIX, _server(server), "gateway") == "gateway" - - -def test_windows_forwards_uvicorn_and_refuses_gunicorn() -> None: - assert startup_check(WINDOWS, _server(), "gateway") == "gateway" - assert startup_check(WINDOWS, _server(Server.GUNICORN), "gateway") == StartupError( - "LITELLM_RUST does not support gunicorn on Windows" - ) - - -def test_missing_binary_is_refused_before_anything_starts(tmp_path: Path) -> None: - prepared: Final[list[bool]] = [] - refused: Final = supervise( - _server(), ("python",), before_start=lambda: prepared.append(True), scripts_dir=str(tmp_path), platform=POSIX - ) - assert (refused, prepared) == ( - StartupError("LITELLM_RUST requires the litellm-gateway binary from the litellm-proxy package"), - [], - ) - - -@POSIX_ONLY -def test_binary_is_found_only_in_the_environment_scripts_dir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - scripts: Final = tmp_path / "scripts" - elsewhere: Final = tmp_path / "elsewhere" - scripts.mkdir() - elsewhere.mkdir() - binary: Final = elsewhere / "litellm-gateway" - binary.write_text("#!/bin/sh\nexit 0\n") - binary.chmod(0o755) - monkeypatch.setenv("PATH", str(elsewhere)) - assert gateway_binary(POSIX, str(scripts)) is None - assert gateway_binary(POSIX, str(elsewhere)) == str(binary) - - -PYTHON_CHILD: Final = """ -import json, os, pathlib, socket, time -time.sleep(0.5) -pathlib.Path(os.environ["MARKERS"], "listening").touch() -private = json.loads(os.environ["_LITELLM_GATEWAY_SUPERVISED"])["listener"] -listener = socket.socket(fileno=private["fd"]) if "fd" in private else socket.create_server(("127.0.0.1", private["port"])) -listener.listen() -time.sleep(10) -""" -RUST_CHILD: Final = """ -import os, pathlib -markers = pathlib.Path(os.environ["MARKERS"]) -(markers / "rust_saw_python_listening").write_text(str((markers / "listening").exists())) -""" - - -def _fake_gateway(scripts: Path) -> Platform: - if sys.platform != "win32": - gateway: Final = scripts / "litellm-gateway" - gateway.write_text(f"#!{sys.executable} -I\n{RUST_CHILD}") - gateway.chmod(0o755) - return current_platform() - script: Final = scripts / "rust_child.py" - script.write_text(RUST_CHILD) - (scripts / "litellm-gateway.cmd").write_text(f'@"{sys.executable}" -I "{script}"\r\n') - return WindowsPlatform(gateway_name="litellm-gateway.cmd") - - -def test_rust_opens_the_public_port_only_once_python_listens(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("MARKERS", str(tmp_path)) - command: Final = (sys.executable, "-I", "-c", PYTHON_CHILD) - supervise(_server(), command, scripts_dir=str(tmp_path), platform=_fake_gateway(tmp_path)) - assert (tmp_path / "rust_saw_python_listening").read_text() == "True" - - -@dataclass(frozen=True, slots=True) -class World: - """What the operating system knows about the children, kept independently of ``Lifecycle``.""" - - alive: frozenset[Child] = frozenset({Child.PYTHON}) - spawned: frozenset[Child] = frozenset({Child.PYTHON}) - reported_listening: bool = False - stop_requested: bool = False - python_status: int | None = None - rust_lost: bool = False - - -def _expected_exit(world: World) -> int | None: - if world.alive or world.python_status is None: - return None - return FAILURE if world.rust_lost else world.python_status - - -def _expected_relay(world: World, signum: signal.Signals) -> tuple[Action, ...]: - to_rust: Final = (Send(Child.RUST, signum),) if signum in STOPS and Child.RUST in world.alive else () - to_python: Final = (Send(Child.PYTHON, signum),) if Child.PYTHON in world.alive else () - return (*to_rust, *to_python) - - -def _expected_on_exit(world: World, child: Child) -> tuple[Action, ...]: - match child: - case Child.PYTHON: - return (Send(Child.RUST, TERM),) if Child.RUST in world.alive else () - case Child.RUST: - serving: Final = Child.PYTHON in world.alive and not world.stop_requested - return (Send(Child.PYTHON, TERM),) if serving else () - - -def _after(world: World, event: Event) -> World: - match event: - case Caught(signum): - return replace(world, stop_requested=world.stop_requested or signum in STOPS) - case Probe.PYTHON_LISTENING: - return replace(world, reported_listening=world.reported_listening or Child.PYTHON in world.alive) - case ChildExited(child, status) if child in world.alive: - lost: Final = child is Child.RUST and Child.PYTHON in world.alive and not world.stop_requested - return replace( - world, - alive=world.alive - {child}, - python_status=status if child is Child.PYTHON else world.python_status, - rust_lost=world.rust_lost or lost, - ) - case ChildExited(): - return world - - -def _spawned(world: World, actions: tuple[Action, ...]) -> World: - started: Final = frozenset(action.child for action in actions if isinstance(action, Spawn)) - return replace(world, alive=world.alive | started, spawned=world.spawned | started) - - -SIGNALS: Final = POSIX.relayed_signals -STOPS: Final = POSIX.stop_signals -EVENTS: Final = st.lists( - st.one_of( - st.sampled_from(SIGNALS).map(Caught), - st.just(Probe.PYTHON_LISTENING), - st.builds(ChildExited, st.sampled_from(Child), st.integers(-15, 3)), - ), - max_size=40, -) - - -@given(EVENTS) -@example([Caught(TERM), Caught(signal.SIGHUP), Probe.PYTHON_LISTENING]) -@example([Probe.PYTHON_LISTENING, Caught(TERM), Caught(signal.SIGWINCH), ChildExited(Child.RUST, 0)]) -def test_lifecycle_relays_every_signal_and_ends_with_pythons_status(events: list[Event]) -> None: - state, actions = boot() - world = World() - assert actions == (Spawn(Child.PYTHON),) - for event in events: - before = _after(world, event) - state, actions = step(state, event, STOPS) - match event: - case Caught(signum): - assert actions == _expected_relay(before, signum) - case Probe.PYTHON_LISTENING: - starts_rust = Child.PYTHON in world.alive and not world.reported_listening and not world.stop_requested - assert actions == ((Spawn(Child.RUST),) if starts_rust else ()) - case ChildExited(child) if child in world.alive: - assert actions == _expected_on_exit(world, child) - case ChildExited(): - assert actions == () - assert all(action.child in before.alive for action in actions if isinstance(action, Send)) - assert all(action.child not in before.spawned for action in actions if isinstance(action, Spawn)) - world = _spawned(before, actions) - assert exit_status(state) == _expected_exit(world) - - -@given(EVENTS) -def test_rust_starts_at_most_once_and_only_while_python_listens_unstopped(events: list[Event]) -> None: - def advance(acc: tuple[Lifecycle, tuple[tuple[Event, Lifecycle, Action], ...]], event: Event): - state, seen = acc - following, actions = step(state, event, STOPS) - return following, (*seen, *((event, state, action) for action in actions)) - - _, history = functools.reduce(advance, events, (boot()[0], ())) - spawns: Final = tuple((event, state) for event, state, action in history if isinstance(action, Spawn)) - assert len(spawns) <= 1 - assert all( - event is Probe.PYTHON_LISTENING and state.python is Phase.SPAWNED and not state.stop_requested - for event, state in spawns - ) - - -@pytest.mark.parametrize( - ("events", "expected"), - [ - ((Probe.PYTHON_LISTENING, Caught(TERM), ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, -15)), -15), - ( - ( - Probe.PYTHON_LISTENING, - Caught(INT), - Caught(INT), - ChildExited(Child.PYTHON, 1), - ChildExited(Child.RUST, 0), - ), - 1, - ), - ( - (Probe.PYTHON_LISTENING, Caught(signal.SIGHUP), ChildExited(Child.PYTHON, -1), ChildExited(Child.RUST, 0)), - -1, - ), - ((Probe.PYTHON_LISTENING, ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, 0)), FAILURE), - ((ChildExited(Child.PYTHON, 3),), 3), - ], -) -def test_supervisor_exits_as_python_did_unless_rust_was_lost(events: tuple[Event, ...], expected: int) -> None: - final: Final = functools.reduce(lambda state, event: step(state, event, STOPS)[0], events, boot()[0]) - assert exit_status(final) == expected diff --git a/tests/unit/proxy/test_proxy_cli.py b/tests/unit/proxy/test_proxy_cli.py index 5ad21572605..df7bbd43324 100644 --- a/tests/unit/proxy/test_proxy_cli.py +++ b/tests/unit/proxy/test_proxy_cli.py @@ -4313,7 +4313,7 @@ class TestRustForwardServer: ) def test_unsupported_servers_are_refused_before_starting_children(self, argv: list[str], reason: str) -> None: from litellm.proxy.rust_forward.platform.posix import PosixPlatform - from litellm.proxy.rust_forward.supervisor import StartupError, startup_check + from litellm.proxy.rust_forward.supervisor.preflight import StartupError, startup_check _, server = self._server(argv) assert startup_check(PosixPlatform(), server, "gateway") == StartupError(f"LITELLM_RUST {reason}") @@ -4417,7 +4417,7 @@ class TestRustForwardServer: from litellm.proxy.rust_forward.mode import InheritedFd, Supervised, rust_forward_mode from litellm.proxy.rust_forward.platform import PrivateListener from litellm.proxy.rust_forward.platform.posix import PosixPlatform - from litellm.proxy.rust_forward.supervisor import python_spec + from litellm.proxy.rust_forward.supervisor.children import python_spec monkeypatch.delenv("FORWARDED_ALLOW_IPS", raising=False) _, server = self._server(["--run_gunicorn"] if gunicorn else [])