diff --git a/litellm/proxy/rust_forward/platform/__init__.py b/litellm/proxy/rust_forward/platform/__init__.py index 4043adacdfb..b5a53d6efb0 100644 --- a/litellm/proxy/rust_forward/platform/__init__.py +++ b/litellm/proxy/rust_forward/platform/__init__.py @@ -24,6 +24,18 @@ class SpawnOptions: creationflags: int +class ChildProcess(Protocol): + @property + def pid(self) -> int: ... + + @property + def stdin(self) -> asyncio.StreamWriter | None: ... + + def send_signal(self, signum: int, /) -> None: ... + + async def wait(self) -> int: ... + + class Platform(Protocol): @property def name(self) -> str: ... @@ -46,9 +58,9 @@ class Platform(Protocol): def catch(self, loop: asyncio.AbstractEventLoop, on_signal: Callable[[signal.Signals], None]) -> None: ... - def signal_child(self, process: asyncio.subprocess.Process, signum: signal.Signals) -> None: ... + def signal_child(self, process: ChildProcess, signum: signal.Signals) -> None: ... - def stop_group(self, process: asyncio.subprocess.Process) -> None: ... + def stop_group(self, process: ChildProcess) -> None: ... def current_platform() -> Platform: diff --git a/litellm/proxy/rust_forward/platform/posix.py b/litellm/proxy/rust_forward/platform/posix.py index 0d0c6c96649..0cb71b01ae2 100644 --- a/litellm/proxy/rust_forward/platform/posix.py +++ b/litellm/proxy/rust_forward/platform/posix.py @@ -10,7 +10,7 @@ from typing import Final from pydantic import TypeAdapter from litellm.proxy.rust_forward.mode import RUST_HOP -from litellm.proxy.rust_forward.platform import PlatformListener, SpawnOptions +from litellm.proxy.rust_forward.platform import ChildProcess, PlatformListener, SpawnOptions from litellm.proxy.rust_forward.server import Server from litellm.proxy.server_listener import InheritedFd @@ -55,10 +55,10 @@ class PosixPlatform: for signum in self.relayed_signals: loop.add_signal_handler(signum, on_signal, signum) - def signal_child(self, process: asyncio.subprocess.Process, signum: signal.Signals) -> None: + def signal_child(self, process: ChildProcess, signum: signal.Signals) -> None: with contextlib.suppress(ProcessLookupError): process.send_signal(signum) - def stop_group(self, process: asyncio.subprocess.Process) -> None: + def stop_group(self, process: ChildProcess) -> None: with contextlib.suppress(ProcessLookupError): os.killpg(process.pid, signal.SIGTERM) diff --git a/litellm/proxy/rust_forward/platform/windows.py b/litellm/proxy/rust_forward/platform/windows.py index 163b96b9b20..841c9cbfd59 100644 --- a/litellm/proxy/rust_forward/platform/windows.py +++ b/litellm/proxy/rust_forward/platform/windows.py @@ -10,7 +10,7 @@ from typing import Final from pydantic import TypeAdapter from litellm.proxy.rust_forward.mode import RUST_HOP -from litellm.proxy.rust_forward.platform import PlatformListener, SpawnOptions +from litellm.proxy.rust_forward.platform import ChildProcess, PlatformListener, SpawnOptions from litellm.proxy.rust_forward.server import Server from litellm.proxy.server_listener import AssignedPort @@ -60,9 +60,9 @@ class WindowsPlatform: # A child in its own process group can only be sent CTRL_BREAK_EVENT, since Windows never delivers CTRL_C_EVENT # to a new group. Both children shut down on it as on Ctrl-C, but uvicorn force-exits only on a repeated SIGINT, # so a second Ctrl-C does not cut a Windows shutdown short the way it does on POSIX. - def signal_child(self, process: asyncio.subprocess.Process, signum: signal.Signals) -> None: + def signal_child(self, process: ChildProcess, signum: signal.Signals) -> None: self.stop_group(process) - def stop_group(self, process: asyncio.subprocess.Process) -> None: + def stop_group(self, process: ChildProcess) -> None: with contextlib.suppress(ProcessLookupError): process.send_signal(CTRL_BREAK_EVENT) diff --git a/litellm/proxy/rust_forward/supervisor/runtime.py b/litellm/proxy/rust_forward/supervisor/runtime.py index 7b727fc0d80..8a8bda6f7f2 100644 --- a/litellm/proxy/rust_forward/supervisor/runtime.py +++ b/litellm/proxy/rust_forward/supervisor/runtime.py @@ -14,7 +14,7 @@ from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStre 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.platform import ChildProcess, 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 ( @@ -46,6 +46,27 @@ class Executor(Protocol): async def next_event(self) -> Event: ... +class ChildControl(Protocol): + def signal_child(self, process: ChildProcess, signum: signal.Signals) -> None: ... + + def stop_group(self, process: ChildProcess) -> None: ... + + +class Spawner(Protocol): + async def __call__(self, spec: ChildSpec, /) -> ChildProcess: ... + + +async def spawn_subprocess(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, + ) + + async def run(platform: Platform, binary: str, command: Sequence[str], server: PythonServer) -> int | SpawnFailed: with platform.private_listener() as private: specs: Final = MappingProxyType( @@ -54,7 +75,7 @@ async def run(platform: Platform, binary: str, command: Sequence[str], server: P Child.RUST: rust_spec(platform, binary, private.port, server), } ) - async with open_children(platform, specs, private.port) as children: + async with open_children(platform, spawn_subprocess, specs, private.port) as children: platform.catch(asyncio.get_running_loop(), children.caught) return await drive(children, platform.stop_signals) @@ -100,14 +121,15 @@ def exit_like( @dataclass(frozen=True, slots=True) class Children: - platform: Platform + platform: ChildControl + spawner: Spawner specs: Mapping[Child, ChildSpec] private_port: int stack: contextlib.AsyncExitStack tasks: TaskGroup events_in: MemoryObjectSendStream[Event] events_out: MemoryObjectReceiveStream[Event] - processes: dict[Child, asyncio.subprocess.Process] # mutable-ok: the running children the stack owns + processes: dict[Child, ChildProcess] # mutable-ok: the running children the stack owns def caught(self, signum: signal.Signals) -> None: with contextlib.suppress(anyio.ClosedResourceError): @@ -115,7 +137,9 @@ class Children: async def spawn(self, child: Child) -> SpawnFailed | None: try: - process: Final = await self.stack.enter_async_context(_process(self.platform, self.specs[child])) + process: Final = await self.stack.enter_async_context( + _process(self.platform, self.spawner, self.specs[child]) + ) except OSError as error: return SpawnFailed(error) self.processes[child] = process @@ -137,28 +161,21 @@ class Children: # unbounded event buffer. @contextlib.asynccontextmanager async def open_children( - platform: Platform, specs: Mapping[Child, ChildSpec], private_port: int + platform: ChildControl, spawner: Spawner, specs: Mapping[Child, ChildSpec], private_port: int ) -> AsyncGenerator[Children]: events_in, events_out = anyio.create_memory_object_stream[Event](math.inf) with events_in, events_out: async with contextlib.AsyncExitStack() as stack: tasks: Final = await stack.enter_async_context(anyio.create_task_group()) - yield Children(platform, specs, private_port, stack, tasks, events_in, events_out, {}) + yield Children(platform, spawner, specs, private_port, stack, tasks, events_in, events_out, {}) # 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. Stopping is shielded so a cancelled supervisor still waits # for the child it is stopping. @contextlib.asynccontextmanager -async def _process(platform: Platform, spec: ChildSpec) -> 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, - ) +async def _process(platform: ChildControl, spawner: Spawner, spec: ChildSpec) -> AsyncGenerator[ChildProcess]: + process: Final = await spawner(spec) try: yield process finally: @@ -169,9 +186,7 @@ async def _process(platform: Platform, spec: ChildSpec) -> AsyncGenerator[asynci process.stdin.close() -async def _report_exit( - child: Child, process: asyncio.subprocess.Process, events: MemoryObjectSendStream[Event] -) -> None: +async def _report_exit(child: Child, process: ChildProcess, events: MemoryObjectSendStream[Event]) -> None: with events: events.send_nowait(ChildExited(child, await process.wait())) diff --git a/tests/unit/proxy/rust_forward/supervisor/test_runtime.py b/tests/unit/proxy/rust_forward/supervisor/test_runtime.py index b84b1302e3c..c5d1bbc7a36 100644 --- a/tests/unit/proxy/rust_forward/supervisor/test_runtime.py +++ b/tests/unit/proxy/rust_forward/supervisor/test_runtime.py @@ -9,7 +9,7 @@ from typing import Final import anyio import pytest -from litellm.proxy.rust_forward.platform import Platform, current_platform +from litellm.proxy.rust_forward.platform import ChildProcess, Platform, current_platform from litellm.proxy.rust_forward.platform.windows import WindowsPlatform from litellm.proxy.rust_forward.supervisor import supervise from litellm.proxy.rust_forward.supervisor.children import ChildSpec @@ -23,7 +23,7 @@ from litellm.proxy.rust_forward.supervisor.lifecycle import ( Send, Spawn, ) -from litellm.proxy.rust_forward.supervisor.runtime import SpawnFailed, drive, exit_like, open_children +from litellm.proxy.rust_forward.supervisor.runtime import SpawnFailed, drive, exit_like, open_children, spawn_subprocess from tests.unit.proxy.rust_forward.supervisor._support import python_server @@ -129,34 +129,80 @@ def _children(argv: tuple[str, ...]) -> MappingProxyType[Child, ChildSpec]: return MappingProxyType({Child.PYTHON: spec, Child.RUST: spec}) -SLOW_TO_STOP: Final = ( - sys.executable, - "-c", - "import signal, time; signal.signal(signal.SIGTERM, lambda *_: (time.sleep(0.5), exit(0))); time.sleep(30)", -) +@dataclass(slots=True) +class HeldChild: + exits: anyio.Event + waiting: int = 0 + stdin: None = None + + @property + def pid(self) -> int: + raise NotImplementedError + + def send_signal(self, signum: int, /) -> None: + raise NotImplementedError + + async def wait(self) -> int: + self.waiting += 1 + try: + await self.exits.wait() + finally: + self.waiting -= 1 + return 0 + + +@dataclass(frozen=True, slots=True) +class StopRecorder: + stopped: list[ChildProcess] = field(default_factory=list) + + def signal_child(self, process: ChildProcess, signum: signal.Signals) -> None: + raise NotImplementedError + + def stop_group(self, process: ChildProcess) -> None: + self.stopped.append(process) -@POSIX_ONLY async def test_a_cancelled_supervisor_still_waits_for_the_child_it_stops() -> None: - with anyio.CancelScope() as scope: - async with open_children(current_platform(), _children(SLOW_TO_STOP), 1) as children: - assert await children.spawn(Child.RUST) is None - process: Final = children.processes[Child.RUST] - await anyio.sleep(0.3) - scope.cancel() - await anyio.sleep(1) - assert process.returncode == 0 + child: Final = HeldChild(anyio.Event()) + control: Final = StopRecorder() + supervisor: Final = anyio.CancelScope() + returned: Final = anyio.Event() + + async def spawn(_spec: ChildSpec) -> HeldChild: + return child + + async def supervise_one_child() -> None: + with supervisor: + async with open_children(control, spawn, _children(("held",)), 1) as children: + assert await children.spawn(Child.RUST) is None + await anyio.sleep_forever() + returned.set() + + async with anyio.create_task_group() as tasks: + tasks.start_soon(supervise_one_child) + await anyio.wait_all_tasks_blocked() + supervisor.cancel() + await anyio.wait_all_tasks_blocked() + while_child_runs: Final = (tuple(control.stopped), child.waiting, returned.is_set()) + child.exits.set() + await anyio.wait_all_tasks_blocked() + once_child_exits: Final = returned.is_set() + assert (while_child_runs, once_child_exits) == (((child,), 1, False), True) async def test_exiting_children_are_reported_as_events() -> None: - async with open_children(current_platform(), _children((sys.executable, "-c", "exit(7)")), 1) as children: + async with open_children( + current_platform(), spawn_subprocess, _children((sys.executable, "-c", "exit(7)")), 1 + ) as children: assert await children.spawn(Child.RUST) is None with anyio.fail_after(5): assert await children.next_event() == ChildExited(Child.RUST, 7) async def test_a_child_that_cannot_start_is_returned_as_a_failure(tmp_path: Path) -> None: - async with open_children(current_platform(), _children((str(tmp_path / "missing"),)), 1) as children: + async with open_children( + current_platform(), spawn_subprocess, _children((str(tmp_path / "missing"),)), 1 + ) as children: outcome: Final = await children.spawn(Child.RUST) assert isinstance(outcome, SpawnFailed) assert isinstance(outcome.error, FileNotFoundError)