mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
refactor(proxy): lean on the stdlib, uvicorn and pydantic in Rust forwarding
The supervisor runs its children through one contextlib.AsyncExitStack loop instead of recursive async with blocks, and payload-free states and modes are Enum members The Python child's env contract is the Supervised type itself, written and read through one pydantic TypeAdapter, so the trusted proxies it carries are always set. Each entry point resolves them from its server library, uvicorn's Config or gunicorn's, instead of a default copied into this package gateway.launch validates uvicorn's parsed CLI options with one pydantic model. TLS now means what each server serves TLS with: a key or cert for uvicorn's CLI, both for proxy_cli. A cipher list alone, or one half of the pair under proxy_cli, no longer refuses forwarding, since Python alone serves plain HTTP there
This commit is contained in:
Yujong Lee
parent
e74261edc1
commit
cc4bce9cbe
13 changed files with 252 additions and 259 deletions
|
|
@ -80,40 +80,23 @@ def _serve(argv: Sequence[str]) -> None:
|
|||
|
||||
def python_server(argv: Sequence[str]) -> "PythonServer":
|
||||
"""What uvicorn would serve the public port with, read from the same flags and environment it reads."""
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.proxy.rust_forward.server import Listener, PythonServer, Server
|
||||
from litellm.proxy.rust_forward.server import PythonServer, Server, UvicornCliParams, uvicorn_trusted_proxies
|
||||
from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager
|
||||
|
||||
params: Final[Mapping[str, object]] = uvicorn_main.make_context(
|
||||
"uvicorn", list(uvicorn_argv(argv, os.environ))
|
||||
).params
|
||||
tls: Final = (
|
||||
params["ssl_keyfile"] is not None
|
||||
or params["ssl_certfile"] is not None
|
||||
or any(arg.split("=")[0] == "--ssl-ciphers" for arg in argv)
|
||||
params: Final = UvicornCliParams.model_validate(
|
||||
uvicorn_main.make_context("uvicorn", list(uvicorn_argv(argv, os.environ))).params
|
||||
)
|
||||
listener: Final = (
|
||||
Listener.TLS
|
||||
if tls
|
||||
else Listener.UNIX_SOCKET
|
||||
if params["uds"] is not None
|
||||
else Listener.INHERITED_FD
|
||||
if params["fd"] is not None
|
||||
else Listener.TCP
|
||||
)
|
||||
graceful_timeout: Final = TypeAdapter[int | None](int | None).validate_python(params["timeout_graceful_shutdown"])
|
||||
return PythonServer(
|
||||
server=Server.UVICORN,
|
||||
listener=listener,
|
||||
host=TypeAdapter(str).validate_python(params["host"]),
|
||||
port=TypeAdapter(int).validate_python(params["port"]),
|
||||
keepalive_seconds=TypeAdapter(int).validate_python(params["timeout_keep_alive"]),
|
||||
listener=params.listener,
|
||||
host=params.host,
|
||||
port=params.port,
|
||||
keepalive_seconds=params.timeout_keep_alive,
|
||||
graceful_timeout_seconds=GracefulShutdownManager.get_timeout()
|
||||
if graceful_timeout is None
|
||||
else graceful_timeout,
|
||||
proxy_headers=TypeAdapter(bool).validate_python(params["proxy_headers"]),
|
||||
trusted_proxies=TypeAdapter[str | None](str | None).validate_python(params["forwarded_allow_ips"]),
|
||||
if params.timeout_graceful_shutdown is None
|
||||
else params.timeout_graceful_shutdown,
|
||||
proxy_headers=params.proxy_headers,
|
||||
trusted_proxies=uvicorn_trusted_proxies(params.forwarded_allow_ips),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -132,9 +115,8 @@ def configure_database() -> None:
|
|||
|
||||
def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -> None:
|
||||
from litellm.proxy.rust_forward.mode import (
|
||||
Disabled,
|
||||
Mode,
|
||||
Supervised,
|
||||
Supervisor,
|
||||
rust_forward_mode,
|
||||
stop_with_supervisor,
|
||||
uvicorn_flags,
|
||||
|
|
@ -143,14 +125,14 @@ def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -
|
|||
|
||||
mode: Final = rust_forward_mode(os.environ)
|
||||
match mode:
|
||||
case Supervisor() if not isinstance(
|
||||
case Mode.SUPERVISOR if not isinstance(
|
||||
exit_code := supervise(
|
||||
python_server(argv), (sys.executable, "-m", "gateway.launch", *argv), before_start=configure_database
|
||||
),
|
||||
StartupError,
|
||||
):
|
||||
sys.exit(exit_code)
|
||||
case Disabled() | Supervisor():
|
||||
case Mode.DISABLED | Mode.SUPERVISOR:
|
||||
configure_database()
|
||||
serve(uvicorn_argv(argv, os.environ))
|
||||
case Supervised():
|
||||
|
|
|
|||
|
|
@ -854,7 +854,7 @@ class ProxyInitializationHelpers:
|
|||
import uvicorn
|
||||
from gunicorn.config import Config as GunicornConfig
|
||||
|
||||
from litellm.proxy.rust_forward.server import Listener, PythonServer, Server
|
||||
from litellm.proxy.rust_forward.server import Listener, PythonServer, Server, uvicorn_trusted_proxies
|
||||
from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager
|
||||
|
||||
server: Final = (
|
||||
|
|
@ -883,7 +883,7 @@ class ProxyInitializationHelpers:
|
|||
proxy_headers=True,
|
||||
trusted_proxies=",".join(TypeAdapter(list[str]).validate_python(GunicornConfig().forwarded_allow_ips))
|
||||
if server is Server.GUNICORN
|
||||
else None,
|
||||
else uvicorn_trusted_proxies(None),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -1359,7 +1359,7 @@ def run_server(
|
|||
from pydantic import ValidationError
|
||||
|
||||
from litellm.proxy.rust_forward.mode import (
|
||||
Supervisor,
|
||||
Mode,
|
||||
environment_error,
|
||||
gunicorn_bind,
|
||||
rust_forward_mode,
|
||||
|
|
@ -1373,7 +1373,7 @@ def run_server(
|
|||
except ValidationError as error:
|
||||
raise click.ClickException(environment_error(error)) from None
|
||||
stop_with_supervisor(forward_mode)
|
||||
if isinstance(forward_mode, Supervisor) and not any(
|
||||
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
|
||||
|
|
@ -1385,7 +1385,7 @@ def run_server(
|
|||
run_gunicorn=run_gunicorn,
|
||||
run_hypercorn=run_hypercorn,
|
||||
run_granian=run_granian,
|
||||
tls=any((ssl_keyfile_path, ssl_certfile_path, ciphers)),
|
||||
tls=None not in (ssl_keyfile_path, ssl_certfile_path),
|
||||
)
|
||||
supervised: Final = supervise(server, (sys.executable, "-m", "litellm.proxy.proxy_cli", *sys.argv[1:]))
|
||||
if not isinstance(supervised, StartupError):
|
||||
|
|
|
|||
|
|
@ -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. `mode.supervised_environment` and `mode.rust_forward_mode` are the two halves of the Python child's, and `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.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`
|
||||
|
||||
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 `async with` around each child owns its cleanup. 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. 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
|
||||
|
||||
- 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
|
||||
|
|
|
|||
|
|
@ -4,40 +4,28 @@ import signal
|
|||
import threading
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum, auto
|
||||
from types import MappingProxyType
|
||||
from typing import Final, TypeAlias
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, ValidationError
|
||||
from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm.rust_bridge.configuration import ENV_NAME, RustSwitch
|
||||
|
||||
GATEWAY_FD: Final = "_LITELLM_GATEWAY_FD"
|
||||
GATEWAY_PORT: Final = "_LITELLM_GATEWAY_PORT"
|
||||
PROXY_HEADERS: Final = "_LITELLM_GATEWAY_PROXY_HEADERS"
|
||||
TRUSTED_PROXIES: Final = "FORWARDED_ALLOW_IPS"
|
||||
DEFAULT_TRUSTED_PROXIES: Final = "127.0.0.1"
|
||||
SUPERVISED_ENV: Final = "_LITELLM_GATEWAY_SUPERVISED"
|
||||
RUST_HOP: Final = "127.0.0.1"
|
||||
|
||||
|
||||
class _ForwardEnvironment(BaseModel):
|
||||
class _Switch(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
rust: RustSwitch = Field(None, validation_alias=ENV_NAME)
|
||||
listener_fd: int | None = Field(None, validation_alias=GATEWAY_FD)
|
||||
listener_port: int | None = Field(None, validation_alias=GATEWAY_PORT)
|
||||
proxy_headers: bool = Field(True, validation_alias=PROXY_HEADERS)
|
||||
trusted_proxies: str = Field(DEFAULT_TRUSTED_PROXIES, validation_alias=TRUSTED_PROXIES)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Disabled:
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Supervisor:
|
||||
pass
|
||||
class Mode(Enum):
|
||||
DISABLED = auto()
|
||||
SUPERVISOR = auto()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -60,41 +48,20 @@ class Supervised:
|
|||
trusted_proxies: str
|
||||
|
||||
|
||||
RustForwardMode: TypeAlias = Disabled | Supervisor | Supervised
|
||||
RustForwardMode: TypeAlias = Mode | Supervised
|
||||
_SUPERVISED: Final = TypeAdapter(Supervised)
|
||||
|
||||
|
||||
def _listener_environment(listener: ChildListener) -> tuple[str, str]:
|
||||
match listener:
|
||||
case InheritedFd(fd):
|
||||
return GATEWAY_FD, str(fd)
|
||||
case AssignedPort(port):
|
||||
return GATEWAY_PORT, str(port)
|
||||
case _:
|
||||
assert_never(listener)
|
||||
|
||||
|
||||
def supervised_environment(
|
||||
listener: ChildListener, proxy_headers: bool, trusted_proxies: str | None
|
||||
) -> Mapping[str, str]:
|
||||
name, value = _listener_environment(listener)
|
||||
return MappingProxyType(
|
||||
{
|
||||
name: value,
|
||||
PROXY_HEADERS: str(proxy_headers).lower(),
|
||||
**({} if trusted_proxies is None else {TRUSTED_PROXIES: trusted_proxies}),
|
||||
}
|
||||
)
|
||||
def supervised_environment(supervised: Supervised) -> Mapping[str, str]:
|
||||
return MappingProxyType({SUPERVISED_ENV: _SUPERVISED.dump_json(supervised).decode()})
|
||||
|
||||
|
||||
def rust_forward_mode(environ: Mapping[str, str]) -> RustForwardMode:
|
||||
parsed: Final = _ForwardEnvironment.model_validate(environ)
|
||||
if parsed.listener_fd is not None:
|
||||
return Supervised(InheritedFd(parsed.listener_fd), parsed.proxy_headers, parsed.trusted_proxies)
|
||||
if parsed.listener_port is not None:
|
||||
return Supervised(AssignedPort(parsed.listener_port), parsed.proxy_headers, parsed.trusted_proxies)
|
||||
if parsed.rust:
|
||||
return Supervisor()
|
||||
return Disabled()
|
||||
switch: Final = _Switch.model_validate(environ)
|
||||
supervised: Final = environ.get(SUPERVISED_ENV)
|
||||
if supervised is not None:
|
||||
return _SUPERVISED.validate_json(supervised)
|
||||
return Mode.SUPERVISOR if switch.rust else Mode.DISABLED
|
||||
|
||||
|
||||
def _wait_for_eof(fd: int) -> None:
|
||||
|
|
@ -118,7 +85,7 @@ def stop_with_supervisor(mode: RustForwardMode, fd: int = 0, stop: Callable[[],
|
|||
match mode:
|
||||
case Supervised():
|
||||
threading.Thread(target=_stop_when_closed, args=(fd, stop), name="supervisor-stdin", daemon=True).start()
|
||||
case Disabled() | Supervisor():
|
||||
case Mode.DISABLED | Mode.SUPERVISOR:
|
||||
pass
|
||||
case _:
|
||||
assert_never(mode)
|
||||
|
|
@ -151,7 +118,7 @@ def uvicorn_options(mode: RustForwardMode, factory: str) -> Mapping[str, object]
|
|||
return MappingProxyType(
|
||||
{"app": factory, "factory": True, **_uvicorn_listener(listener), "proxy_headers": False}
|
||||
)
|
||||
case Disabled() | Supervisor():
|
||||
case Mode.DISABLED | Mode.SUPERVISOR:
|
||||
return MappingProxyType({})
|
||||
case _:
|
||||
assert_never(mode)
|
||||
|
|
@ -171,7 +138,7 @@ def uvicorn_flags(mode: RustForwardMode) -> tuple[str, ...]:
|
|||
match mode:
|
||||
case Supervised(listener):
|
||||
return (*_uvicorn_listener_flags(listener), "--no-proxy-headers")
|
||||
case Disabled() | Supervisor():
|
||||
case Mode.DISABLED | Mode.SUPERVISOR:
|
||||
return ()
|
||||
case _:
|
||||
assert_never(mode)
|
||||
|
|
@ -191,7 +158,7 @@ def gunicorn_bind(mode: RustForwardMode, public: str) -> str:
|
|||
match mode:
|
||||
case Supervised(listener):
|
||||
return _gunicorn_listener(listener)
|
||||
case Disabled() | Supervisor():
|
||||
case Mode.DISABLED | Mode.SUPERVISOR:
|
||||
return public
|
||||
case _:
|
||||
assert_never(mode)
|
||||
|
|
@ -201,7 +168,7 @@ def server_proxy_headers(mode: RustForwardMode) -> bool:
|
|||
match mode:
|
||||
case Supervised():
|
||||
return False
|
||||
case Disabled() | Supervisor():
|
||||
case Mode.DISABLED | Mode.SUPERVISOR:
|
||||
return True
|
||||
case _:
|
||||
assert_never(mode)
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ from typing import (
|
|||
|
||||
from typing_extensions import assert_never
|
||||
|
||||
from litellm.proxy.rust_forward.mode import RUST_HOP, Disabled, RustForwardMode, Supervised, Supervisor
|
||||
from litellm.proxy.rust_forward.mode import RUST_HOP, Mode, RustForwardMode, Supervised
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import starlette.types
|
||||
|
|
@ -63,7 +63,7 @@ def wrap(app: "starlette.types.ASGIApp", mode: RustForwardMode) -> ASGIApp:
|
|||
match mode:
|
||||
case Supervised():
|
||||
return RustPeerMiddleware(converted, mode)
|
||||
case Disabled() | Supervisor():
|
||||
case Mode.DISABLED | Mode.SUPERVISOR:
|
||||
return converted
|
||||
case _:
|
||||
assert_never(mode)
|
||||
|
|
|
|||
|
|
@ -1,5 +1,10 @@
|
|||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import Final
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter
|
||||
|
||||
_TRUSTED_PROXIES: Final = TypeAdapter(str)
|
||||
|
||||
|
||||
class Server(Enum):
|
||||
|
|
@ -27,4 +32,39 @@ class PythonServer:
|
|||
keepalive_seconds: int
|
||||
graceful_timeout_seconds: float
|
||||
proxy_headers: bool
|
||||
trusted_proxies: str | None
|
||||
trusted_proxies: str
|
||||
|
||||
|
||||
class UvicornCliParams(BaseModel):
|
||||
"""The options uvicorn's own CLI parsed that decide how it serves the public port."""
|
||||
|
||||
model_config = ConfigDict(frozen=True, extra="ignore")
|
||||
|
||||
host: str
|
||||
port: int
|
||||
uds: str | None
|
||||
fd: int | None
|
||||
ssl_keyfile: str | None
|
||||
ssl_certfile: str | None
|
||||
timeout_keep_alive: int
|
||||
timeout_graceful_shutdown: int | None
|
||||
proxy_headers: bool
|
||||
forwarded_allow_ips: str | None
|
||||
|
||||
@property
|
||||
def listener(self) -> Listener:
|
||||
if self.ssl_keyfile is not None or self.ssl_certfile is not None:
|
||||
return Listener.TLS
|
||||
if self.uds is not None:
|
||||
return Listener.UNIX_SOCKET
|
||||
if self.fd is not None:
|
||||
return Listener.INHERITED_FD
|
||||
return Listener.TCP
|
||||
|
||||
|
||||
def uvicorn_trusted_proxies(configured: str | None) -> str:
|
||||
import uvicorn
|
||||
|
||||
return _TRUSTED_PROXIES.validate_python(
|
||||
uvicorn.Config(app="", log_config=None, forwarded_allow_ips=configured).forwarded_allow_ips
|
||||
)
|
||||
|
|
|
|||
|
|
@ -7,15 +7,15 @@ import sys
|
|||
import sysconfig
|
||||
from collections.abc import AsyncGenerator, Callable, Coroutine, Mapping, Sequence
|
||||
from dataclasses import dataclass, replace
|
||||
from enum import Enum
|
||||
from enum import Enum, auto
|
||||
from functools import partial
|
||||
from types import MappingProxyType
|
||||
from typing import Final, TypeAlias
|
||||
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_environment
|
||||
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
|
||||
|
|
@ -52,19 +52,10 @@ class Child(Enum):
|
|||
FAILURE: Final = 1
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class NotStarted:
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Spawned:
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Listening:
|
||||
pass
|
||||
class Phase(Enum):
|
||||
NOT_STARTED = auto()
|
||||
SPAWNED = auto()
|
||||
LISTENING = auto()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -79,8 +70,8 @@ class Lost:
|
|||
status: int
|
||||
|
||||
|
||||
PythonState: TypeAlias = Spawned | Listening | Exited
|
||||
RustState: TypeAlias = NotStarted | Spawned | Exited | Lost
|
||||
PythonState: TypeAlias = Literal[Phase.SPAWNED, Phase.LISTENING] | Exited
|
||||
RustState: TypeAlias = Literal[Phase.NOT_STARTED, Phase.SPAWNED] | Exited | Lost
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -95,9 +86,8 @@ class Caught:
|
|||
signum: signal.Signals
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class PythonListening:
|
||||
pass
|
||||
class Probe(Enum):
|
||||
PYTHON_LISTENING = auto()
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -117,39 +107,39 @@ class Send:
|
|||
signum: signal.Signals
|
||||
|
||||
|
||||
Event: TypeAlias = Caught | PythonListening | ChildExited
|
||||
Event: TypeAlias = Caught | Literal[Probe.PYTHON_LISTENING] | ChildExited
|
||||
Action: TypeAlias = Spawn | Send
|
||||
Transition: TypeAlias = tuple[Lifecycle, tuple[Action, ...]]
|
||||
|
||||
|
||||
def boot() -> Transition:
|
||||
return Lifecycle(Spawned(), NotStarted(), stop_requested=False), (Spawn(Child.PYTHON),)
|
||||
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 isinstance(state.rust, Spawned) else ()
|
||||
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 not isinstance(state.python, Spawned):
|
||||
if state.python is not Phase.SPAWNED:
|
||||
return state, ()
|
||||
if state.stop_requested:
|
||||
return replace(state, python=Listening()), ()
|
||||
return replace(state, python=Listening(), rust=Spawned()), (Spawn(Child.RUST),)
|
||||
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 isinstance(state.rust, Spawned) else ())
|
||||
return exited, ((Send(Child.RUST, signal.SIGTERM),) if state.rust is Phase.SPAWNED else ())
|
||||
|
||||
|
||||
def _rust_exited(state: Lifecycle, status: int) -> Transition:
|
||||
if not isinstance(state.rust, Spawned):
|
||||
if state.rust is not Phase.SPAWNED:
|
||||
return state, ()
|
||||
if state.stop_requested or isinstance(state.python, Exited):
|
||||
return replace(state, rust=Exited(status)), ()
|
||||
|
|
@ -162,7 +152,7 @@ def step(state: Lifecycle, event: Event, stop_signals: frozenset[signal.Signals]
|
|||
match event:
|
||||
case Caught(signum):
|
||||
return _relay(state, signum, signum in stop_signals)
|
||||
case PythonListening():
|
||||
case Probe.PYTHON_LISTENING:
|
||||
return _python_listening(state)
|
||||
case ChildExited(child, status):
|
||||
match child:
|
||||
|
|
@ -178,7 +168,7 @@ def step(state: Lifecycle, event: Event, stop_signals: frozenset[signal.Signals]
|
|||
|
||||
def exit_status(state: Lifecycle) -> int | None:
|
||||
match state.python, state.rust:
|
||||
case Exited(status), NotStarted() | Exited():
|
||||
case Exited(status), Phase.NOT_STARTED | Exited():
|
||||
return status
|
||||
case Exited(), Lost():
|
||||
return FAILURE
|
||||
|
|
@ -198,7 +188,9 @@ def python_spec(
|
|||
) -> ChildSpec:
|
||||
return ChildSpec(
|
||||
argv=(*command, "--port", str(private.port)),
|
||||
env_overrides=supervised_environment(private.listener, server.proxy_headers, server.trusted_proxies),
|
||||
env_overrides=supervised_environment(
|
||||
Supervised(private.listener, server.proxy_headers, server.trusted_proxies)
|
||||
),
|
||||
spawn=platform.spawn_options(private.pass_fds),
|
||||
)
|
||||
|
||||
|
|
@ -251,7 +243,7 @@ async def _accepts(port: int) -> bool:
|
|||
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(PythonListening())
|
||||
events.put_nowait(Probe.PYTHON_LISTENING)
|
||||
|
||||
|
||||
@contextlib.asynccontextmanager
|
||||
|
|
@ -295,53 +287,22 @@ async def _launch_python(
|
|||
yield process
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Launching:
|
||||
state: Lifecycle
|
||||
child: Child
|
||||
rest: tuple[Action, ...]
|
||||
|
||||
|
||||
async def _advance(
|
||||
platform: Platform,
|
||||
transition: Transition,
|
||||
events: asyncio.Queue[Event],
|
||||
children: Mapping[Child, asyncio.subprocess.Process],
|
||||
) -> int | _Launching:
|
||||
state, actions = transition
|
||||
while True:
|
||||
for index, action in enumerate(actions):
|
||||
match action:
|
||||
case Send(child, signum):
|
||||
platform.signal_child(children[child], signum)
|
||||
case Spawn(child):
|
||||
return _Launching(state, child, actions[index + 1 :])
|
||||
case _:
|
||||
assert_never(action)
|
||||
status = exit_status(state)
|
||||
if status is not None:
|
||||
return status
|
||||
state, actions = step(state, await events.get(), platform.stop_signals)
|
||||
|
||||
|
||||
async def _run(
|
||||
platform: Platform,
|
||||
transition: Transition,
|
||||
events: asyncio.Queue[Event],
|
||||
children: Mapping[Child, asyncio.subprocess.Process],
|
||||
launches: Mapping[Child, Launch],
|
||||
) -> int:
|
||||
advanced: Final = await _advance(platform, transition, events, children)
|
||||
match advanced:
|
||||
case int():
|
||||
return advanced
|
||||
case _Launching(state, child, rest):
|
||||
async with launches[child](events) as process:
|
||||
return await _run(
|
||||
platform, (state, rest), events, MappingProxyType({**children, child: process}), launches
|
||||
)
|
||||
case _:
|
||||
assert_never(advanced)
|
||||
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:
|
||||
|
|
@ -356,7 +317,7 @@ async def _supervise(platform: Platform, binary: str, command: Sequence[str], se
|
|||
Child.RUST: partial(_launch, platform, Child.RUST, rust_spec(platform, binary, private.port, server)),
|
||||
}
|
||||
)
|
||||
return await _run(platform, boot(), events, MappingProxyType({}), launches)
|
||||
return await _run(platform, events, launches)
|
||||
|
||||
|
||||
def _exit_like(status: int) -> int:
|
||||
|
|
|
|||
|
|
@ -46,14 +46,14 @@ import uvicorn
|
|||
|
||||
|
||||
def listener():
|
||||
fd = os.environ.get("_LITELLM_GATEWAY_FD")
|
||||
if fd is None:
|
||||
return {"host": "127.0.0.1", "port": int(sys.argv[2])}
|
||||
from litellm.proxy.rust_forward.mode import rust_forward_mode, stop_with_supervisor
|
||||
from litellm.proxy.rust_forward.mode import Supervised, rust_forward_mode, stop_with_supervisor
|
||||
|
||||
stop_with_supervisor(rust_forward_mode(os.environ))
|
||||
mode = rust_forward_mode(os.environ)
|
||||
if not isinstance(mode, Supervised):
|
||||
return {"host": "127.0.0.1", "port": int(sys.argv[2])}
|
||||
stop_with_supervisor(mode)
|
||||
pathlib.Path(os.environ["MARKERS"], "python.pid").write_text(str(os.getpid()))
|
||||
return {"fd": int(fd)}
|
||||
return {"fd": mode.listener.fd}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
|
@ -123,7 +123,7 @@ from litellm.proxy.rust_forward.server import Listener, PythonServer, Server
|
|||
from litellm.proxy.rust_forward.supervisor import supervise
|
||||
|
||||
workers, port, scripts = sys.argv[1], int(sys.argv[2]), Path(sys.argv[3])
|
||||
server = PythonServer(Server.UVICORN, Listener.TCP, "127.0.0.1", port, 5, 30.0, True, None)
|
||||
server = PythonServer(Server.UVICORN, Listener.TCP, "127.0.0.1", port, 5, 30.0, True, "127.0.0.1")
|
||||
result = supervise(server, (sys.executable, str(scripts / "serve.py"), workers, str(port)), scripts_dir=str(scripts))
|
||||
raise SystemExit(result if isinstance(result, int) else 3)
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -126,7 +126,8 @@ class TestUvicornArgv:
|
|||
(("--fd", "0"), Listener.INHERITED_FD),
|
||||
(("--fd=17",), Listener.INHERITED_FD),
|
||||
(("--ssl-keyfile", "unused"), Listener.TLS),
|
||||
(("--ssl-ciphers", "TLSv1.2"), Listener.TLS),
|
||||
(("--ssl-certfile", "unused"), Listener.TLS),
|
||||
(("--ssl-ciphers", "TLSv1.2"), Listener.TCP),
|
||||
],
|
||||
)
|
||||
def test_forwarding_sees_the_public_listener_uvicorn_would_open(
|
||||
|
|
@ -135,18 +136,25 @@ class TestUvicornArgv:
|
|||
assert python_server(arguments).listener is listener
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("argv", "trust"),
|
||||
("argv", "environ", "trust"),
|
||||
[
|
||||
((), (True, None)),
|
||||
(("--forwarded-allow-ips", "10.0.0.5"), (True, "10.0.0.5")),
|
||||
(("--forwarded-allow-ips=*",), (True, "*")),
|
||||
(("--no-proxy-headers",), (False, None)),
|
||||
((), {}, (True, "127.0.0.1")),
|
||||
((), {"FORWARDED_ALLOW_IPS": "10.9.9.9"}, (True, "10.9.9.9")),
|
||||
(("--forwarded-allow-ips", "10.0.0.5"), {"FORWARDED_ALLOW_IPS": "10.9.9.9"}, (True, "10.0.0.5")),
|
||||
(("--forwarded-allow-ips=*",), {}, (True, "*")),
|
||||
(("--no-proxy-headers",), {}, (False, "127.0.0.1")),
|
||||
],
|
||||
)
|
||||
def test_forwarding_carries_operator_proxy_header_options_to_the_child(
|
||||
self, argv: tuple[str, ...], trust: tuple[bool, str | None], monkeypatch: pytest.MonkeyPatch
|
||||
def test_forwarding_carries_the_trust_uvicorn_resolves_to_the_child(
|
||||
self,
|
||||
argv: tuple[str, ...],
|
||||
environ: dict[str, str],
|
||||
trust: tuple[bool, str],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.delenv("FORWARDED_ALLOW_IPS", raising=False)
|
||||
for name, value in environ.items():
|
||||
monkeypatch.setenv(name, value)
|
||||
server: Final = python_server(("--host", "0.0.0.0", "--port", "4000", *argv))
|
||||
assert (server.host, server.port) == ("0.0.0.0", 4000)
|
||||
assert (server.proxy_headers, server.trusted_proxies) == trust
|
||||
|
|
@ -335,7 +343,7 @@ class TestMain:
|
|||
capsys: pytest.CaptureFixture[str],
|
||||
) -> None:
|
||||
monkeypatch.setenv("LITELLM_RUST", rust)
|
||||
monkeypatch.delenv("_LITELLM_GATEWAY_FD", raising=False)
|
||||
monkeypatch.delenv("_LITELLM_GATEWAY_SUPERVISED", raising=False)
|
||||
monkeypatch.delenv("LITELLM_PGBOUNCER_ENABLED", raising=False)
|
||||
served: Final[list[tuple[str, ...]]] = []
|
||||
main(("--uds", "/tmp/gateway.sock"), serve=lambda argv: served.append(tuple(argv)))
|
||||
|
|
|
|||
|
|
@ -8,12 +8,12 @@ from pydantic import ValidationError
|
|||
|
||||
from litellm.proxy.proxy_cli import WORKER_APP, ProxyInitializationHelpers
|
||||
from litellm.proxy.rust_forward.mode import (
|
||||
SUPERVISED_ENV,
|
||||
AssignedPort,
|
||||
Disabled,
|
||||
InheritedFd,
|
||||
Mode,
|
||||
RustForwardMode,
|
||||
Supervised,
|
||||
Supervisor,
|
||||
gunicorn_bind,
|
||||
rust_forward_mode,
|
||||
server_proxy_headers,
|
||||
|
|
@ -26,21 +26,17 @@ from litellm.proxy.rust_forward.mode import (
|
|||
@pytest.mark.parametrize(
|
||||
("environ", "expected"),
|
||||
[
|
||||
({}, Disabled()),
|
||||
({"LITELLM_RUST": "false"}, Disabled()),
|
||||
({"LITELLM_RUST": "off"}, Disabled()),
|
||||
({"LITELLM_RUST": "0"}, Disabled()),
|
||||
({"LITELLM_RUST": "true"}, Supervisor()),
|
||||
({"LITELLM_RUST": "1"}, Supervisor()),
|
||||
({"LITELLM_RUST": "yes"}, Supervisor()),
|
||||
({"LITELLM_RUST": "on"}, Supervisor()),
|
||||
({"LITELLM_RUST": " 1 "}, Supervisor()),
|
||||
({"LITELLM_RUST": ""}, Disabled()),
|
||||
({"LITELLM_RUST": " "}, Disabled()),
|
||||
(
|
||||
{"LITELLM_RUST": "true", "_LITELLM_GATEWAY_FD": "17", "FORWARDED_ALLOW_IPS": "127.0.0.1,10.0.0.5"},
|
||||
Supervised(InheritedFd(17), True, "127.0.0.1,10.0.0.5"),
|
||||
),
|
||||
({}, 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),
|
||||
],
|
||||
)
|
||||
def test_mode_is_resolved_from_environment(environ: dict[str, str], expected: RustForwardMode) -> None:
|
||||
|
|
@ -48,35 +44,36 @@ def test_mode_is_resolved_from_environment(environ: dict[str, str], expected: Ru
|
|||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("environ", "variable"),
|
||||
("environ", "invalid"),
|
||||
[
|
||||
({"LITELLM_RUST": "enabled"}, "LITELLM_RUST"),
|
||||
({"LITELLM_RUST": "2"}, "LITELLM_RUST"),
|
||||
({"_LITELLM_GATEWAY_FD": "seventeen"}, "_LITELLM_GATEWAY_FD"),
|
||||
({"_LITELLM_GATEWAY_FD": "17", "_LITELLM_GATEWAY_PROXY_HEADERS": "maybe"}, "_LITELLM_GATEWAY_PROXY_HEADERS"),
|
||||
({SUPERVISED_ENV: '{"listener": {"fd": "seventeen"}, "proxy_headers": true, "trusted_proxies": "*"}'}, "fd"),
|
||||
(
|
||||
{SUPERVISED_ENV: '{"listener": {"fd": 17}, "proxy_headers": "maybe", "trusted_proxies": "*"}'},
|
||||
"proxy_headers",
|
||||
),
|
||||
({SUPERVISED_ENV: '{"listener": {"fd": 17}, "proxy_headers": true}'}, "trusted_proxies"),
|
||||
],
|
||||
)
|
||||
def test_invalid_forwarding_value_stops_startup_naming_the_variable(environ: dict[str, str], variable: str) -> None:
|
||||
with pytest.raises(ValidationError, match=variable):
|
||||
def test_invalid_forwarding_value_stops_startup_naming_it(environ: dict[str, str], invalid: str) -> None:
|
||||
with pytest.raises(ValidationError, match=invalid):
|
||||
rust_forward_mode(environ)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("proxy_headers", "trusted_proxies", "inherited", "expected"),
|
||||
"written",
|
||||
[
|
||||
(True, None, {}, Supervised(InheritedFd(17), True, "127.0.0.1")),
|
||||
(True, None, {"FORWARDED_ALLOW_IPS": "10.0.0.5"}, Supervised(InheritedFd(17), True, "10.0.0.5")),
|
||||
(True, "*", {"FORWARDED_ALLOW_IPS": "10.0.0.5"}, Supervised(InheritedFd(17), True, "*")),
|
||||
(False, None, {}, Supervised(InheritedFd(17), False, "127.0.0.1")),
|
||||
(True, None, {}, Supervised(AssignedPort(45555), True, "127.0.0.1")),
|
||||
(False, "*", {}, Supervised(AssignedPort(45555), False, "*")),
|
||||
Supervised(InheritedFd(17), True, "127.0.0.1"),
|
||||
Supervised(InheritedFd(17), True, "127.0.0.1,10.0.0.5"),
|
||||
Supervised(InheritedFd(17), False, "127.0.0.1"),
|
||||
Supervised(AssignedPort(45555), True, "*"),
|
||||
Supervised(AssignedPort(45555), False, "*"),
|
||||
],
|
||||
)
|
||||
def test_child_reads_back_the_listener_and_trust_the_supervisor_wrote(
|
||||
proxy_headers: bool, trusted_proxies: str | None, inherited: dict[str, str], expected: Supervised
|
||||
) -> None:
|
||||
written: Final = supervised_environment(expected.listener, proxy_headers, trusted_proxies)
|
||||
assert rust_forward_mode({**inherited, **written}) == expected
|
||||
def test_child_reads_back_exactly_what_the_supervisor_wrote(written: Supervised) -> None:
|
||||
inherited: Final = {"LITELLM_RUST": "true", "FORWARDED_ALLOW_IPS": "10.9.9.9"}
|
||||
assert rust_forward_mode({**inherited, **supervised_environment(written)}) == written
|
||||
|
||||
|
||||
def _served_uvicorn_config(mode: RustForwardMode) -> uvicorn.Config:
|
||||
|
|
@ -93,7 +90,7 @@ def test_supervised_uvicorn_serves_the_inherited_socket_without_trusting_the_rus
|
|||
assert (config.app, config.factory) == (WORKER_APP, True)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", [Disabled(), Supervisor()])
|
||||
@pytest.mark.parametrize("mode", list(Mode))
|
||||
def test_unsupervised_servers_keep_their_own_socket_and_proxy_headers(mode: RustForwardMode) -> None:
|
||||
config: Final = _served_uvicorn_config(mode)
|
||||
assert (config.fd, config.proxy_headers, config.port) == (None, True, 45555)
|
||||
|
|
@ -117,7 +114,7 @@ def test_assigned_port_serves_only_the_loopback_hop_without_trusting_it() -> Non
|
|||
|
||||
@pytest.mark.parametrize(
|
||||
("mode", "stops"),
|
||||
[(Supervised(InheritedFd(17), False, "127.0.0.1"), True), (Disabled(), False), (Supervisor(), False)],
|
||||
[(Supervised(InheritedFd(17), False, "127.0.0.1"), True), (Mode.DISABLED, False), (Mode.SUPERVISOR, False)],
|
||||
)
|
||||
def test_only_a_supervised_server_stops_once_its_supervisor_is_gone(mode: RustForwardMode, stops: bool) -> None:
|
||||
read_end, write_end = os.pipe()
|
||||
|
|
|
|||
|
|
@ -20,7 +20,7 @@ if TYPE_CHECKING:
|
|||
)
|
||||
|
||||
from litellm.proxy.proxy_server import app as proxy_app
|
||||
from litellm.proxy.rust_forward.mode import InheritedFd, Supervised
|
||||
from litellm.proxy.rust_forward.mode import SUPERVISED_ENV, InheritedFd, Supervised, supervised_environment
|
||||
from litellm.proxy.rust_forward.peer import RustPeerMiddleware
|
||||
|
||||
Headers = tuple[tuple[bytes, bytes], ...]
|
||||
|
|
@ -124,7 +124,9 @@ async def _peer_probe(request: Request) -> dict[str, object]:
|
|||
async def test_every_entry_point_hands_routes_the_peer_rust_saw(
|
||||
entry_point: str, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setenv("_LITELLM_GATEWAY_FD", "17")
|
||||
monkeypatch.setenv(
|
||||
SUPERVISED_ENV, supervised_environment(Supervised(InheritedFd(17), True, "127.0.0.1"))[SUPERVISED_ENV]
|
||||
)
|
||||
monkeypatch.setattr(proxy_app.router, "lifespan_context", proxy_app.router.lifespan_context)
|
||||
monkeypatch.setattr(
|
||||
proxy_app.router, "routes", [APIRoute("/__rust_forward_peer_probe", _peer_probe), *proxy_app.router.routes]
|
||||
|
|
|
|||
|
|
@ -31,10 +31,10 @@ from litellm.proxy.rust_forward.supervisor import (
|
|||
ChildExited,
|
||||
Event,
|
||||
Lifecycle,
|
||||
PythonListening,
|
||||
Phase,
|
||||
Probe,
|
||||
Send,
|
||||
Spawn,
|
||||
Spawned,
|
||||
StartupError,
|
||||
boot,
|
||||
exit_status,
|
||||
|
|
@ -54,7 +54,7 @@ POSIX_ONLY: Final = pytest.mark.skipif(sys.platform == "win32", reason="runs a P
|
|||
|
||||
|
||||
def _server(
|
||||
server: Server = Server.UVICORN, listener: Listener = Listener.TCP, trusted_proxies: str | None = None
|
||||
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)
|
||||
|
||||
|
|
@ -168,12 +168,11 @@ def test_binary_is_found_only_in_the_environment_scripts_dir(tmp_path: Path, mon
|
|||
|
||||
|
||||
PYTHON_CHILD: Final = """
|
||||
import os, pathlib, socket, time
|
||||
import json, os, pathlib, socket, time
|
||||
time.sleep(0.5)
|
||||
pathlib.Path(os.environ["MARKERS"], "listening").touch()
|
||||
fd = os.environ.get("_LITELLM_GATEWAY_FD")
|
||||
port = os.environ.get("_LITELLM_GATEWAY_PORT")
|
||||
listener = socket.socket(fileno=int(fd)) if fd else socket.create_server(("127.0.0.1", int(port)))
|
||||
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)
|
||||
"""
|
||||
|
|
@ -240,7 +239,7 @@ 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 PythonListening():
|
||||
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
|
||||
|
|
@ -264,7 +263,7 @@ STOPS: Final = POSIX.stop_signals
|
|||
EVENTS: Final = st.lists(
|
||||
st.one_of(
|
||||
st.sampled_from(SIGNALS).map(Caught),
|
||||
st.just(PythonListening()),
|
||||
st.just(Probe.PYTHON_LISTENING),
|
||||
st.builds(ChildExited, st.sampled_from(Child), st.integers(-15, 3)),
|
||||
),
|
||||
max_size=40,
|
||||
|
|
@ -272,8 +271,8 @@ EVENTS: Final = st.lists(
|
|||
|
||||
|
||||
@given(EVENTS)
|
||||
@example([Caught(TERM), Caught(signal.SIGHUP), PythonListening()])
|
||||
@example([PythonListening(), Caught(TERM), Caught(signal.SIGWINCH), ChildExited(Child.RUST, 0)])
|
||||
@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()
|
||||
|
|
@ -284,7 +283,7 @@ def test_lifecycle_relays_every_signal_and_ends_with_pythons_status(events: list
|
|||
match event:
|
||||
case Caught(signum):
|
||||
assert actions == _expected_relay(before, signum)
|
||||
case PythonListening():
|
||||
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:
|
||||
|
|
@ -308,7 +307,7 @@ def test_rust_starts_at_most_once_and_only_while_python_listens_unstopped(events
|
|||
spawns: Final = tuple((event, state) for event, state, action in history if isinstance(action, Spawn))
|
||||
assert len(spawns) <= 1
|
||||
assert all(
|
||||
isinstance(event, PythonListening) and isinstance(state.python, Spawned) and not state.stop_requested
|
||||
event is Probe.PYTHON_LISTENING and state.python is Phase.SPAWNED and not state.stop_requested
|
||||
for event, state in spawns
|
||||
)
|
||||
|
||||
|
|
@ -316,10 +315,22 @@ def test_rust_starts_at_most_once_and_only_while_python_listens_unstopped(events
|
|||
@pytest.mark.parametrize(
|
||||
("events", "expected"),
|
||||
[
|
||||
((PythonListening(), Caught(TERM), ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, -15)), -15),
|
||||
((PythonListening(), Caught(INT), Caught(INT), ChildExited(Child.PYTHON, 1), ChildExited(Child.RUST, 0)), 1),
|
||||
((PythonListening(), Caught(signal.SIGHUP), ChildExited(Child.PYTHON, -1), ChildExited(Child.RUST, 0)), -1),
|
||||
((PythonListening(), ChildExited(Child.RUST, 0), ChildExited(Child.PYTHON, 0)), FAILURE),
|
||||
((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),
|
||||
],
|
||||
)
|
||||
|
|
|
|||
|
|
@ -4297,15 +4297,16 @@ class TestRustForwardServer:
|
|||
run_gunicorn=params["run_gunicorn"] is True,
|
||||
run_hypercorn=params["run_hypercorn"] is True,
|
||||
run_granian=params["run_granian"] is True,
|
||||
tls=any((params["ssl_keyfile_path"], params["ssl_certfile_path"], params["ciphers"])),
|
||||
tls=None not in (params["ssl_keyfile_path"], params["ssl_certfile_path"]),
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("argv", "reason"),
|
||||
[
|
||||
(["--ssl_keyfile_path", "unused"], "does not support TLS on the public port"),
|
||||
(["--ssl_certfile_path", "unused"], "does not support TLS on the public port"),
|
||||
(["--ciphers", "unused"], "does not support TLS on the public port"),
|
||||
(
|
||||
["--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"),
|
||||
],
|
||||
|
|
@ -4321,8 +4322,28 @@ class TestRustForwardServer:
|
|||
@patch("atexit.register")
|
||||
@patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database")
|
||||
@patch("litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False)
|
||||
@pytest.mark.parametrize(
|
||||
("argv", "reason", "served_keyfile"),
|
||||
[
|
||||
(
|
||||
["--ssl_keyfile_path", "key.pem", "--ssl_certfile_path", "cert.pem"],
|
||||
"does not support TLS on the public port",
|
||||
"key.pem",
|
||||
),
|
||||
(["--ssl_keyfile_path", "key.pem"], "requires the litellm-gateway binary", None),
|
||||
(["--ssl_certfile_path", "cert.pem"], "requires the litellm-gateway binary", None),
|
||||
(["--ciphers", "ECDHE"], "requires the litellm-gateway binary", None),
|
||||
],
|
||||
)
|
||||
def test_a_refused_forwarding_start_serves_from_python_alone(
|
||||
self, mock_should_update, mock_setup_db, mock_atexit_register, mock_uvicorn_run
|
||||
self,
|
||||
mock_should_update,
|
||||
mock_setup_db,
|
||||
mock_atexit_register,
|
||||
mock_uvicorn_run,
|
||||
argv: list[str],
|
||||
reason: str,
|
||||
served_keyfile: str | None,
|
||||
) -> None:
|
||||
from click.testing import CliRunner
|
||||
|
||||
|
|
@ -4333,20 +4354,24 @@ class TestRustForwardServer:
|
|||
save_worker_config=MagicMock(),
|
||||
)
|
||||
environ: Final = {
|
||||
**{k: v for k, v in os.environ.items() if k not in ("DATABASE_URL", "DIRECT_URL", "_LITELLM_GATEWAY_FD")},
|
||||
**{
|
||||
k: v
|
||||
for k, v in os.environ.items()
|
||||
if k not in ("DATABASE_URL", "DIRECT_URL", "_LITELLM_GATEWAY_SUPERVISED")
|
||||
},
|
||||
"LITELLM_RUST": "true",
|
||||
}
|
||||
with (
|
||||
patch.dict(os.environ, environ, clear=True),
|
||||
patch.dict("sys.modules", {"proxy_server": proxy_module, "litellm.proxy.proxy_server": proxy_module}),
|
||||
patch("litellm.proxy.rust_forward.supervisor.gateway_binary", return_value=None),
|
||||
):
|
||||
result: Final = CliRunner().invoke(
|
||||
run_server, ["--port", "45999", "--ssl_keyfile_path", "key.pem", "--ssl_certfile_path", "cert.pem"]
|
||||
)
|
||||
result: Final = CliRunner().invoke(run_server, ["--port", "45999", *argv])
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "LITELLM_RUST does not support TLS on the public port, so Python serves the public port" in result.output
|
||||
assert f"LITELLM_RUST {reason}" in result.output
|
||||
assert "so Python serves the public port without Rust forwarding" in result.output
|
||||
mock_uvicorn_run.assert_called_once()
|
||||
assert mock_uvicorn_run.call_args.kwargs["ssl_keyfile"] == "key.pem"
|
||||
assert mock_uvicorn_run.call_args.kwargs.get("ssl_keyfile") == served_keyfile
|
||||
assert mock_uvicorn_run.call_args.kwargs["port"] == 45999
|
||||
|
||||
def test_an_invalid_rust_switch_stops_startup_with_one_line_naming_it(self) -> None:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue