test(proxy): make the cancelled supervisor test deterministic (#51258)

The test spawned a real Python child that installs a SIGTERM handler, slept
0.3s, then cancelled the supervisor and expected the child to exit 0. CI runs
with --cov, and pytest-cov starts coverage in every spawned Python process,
so the child is often still starting when SIGTERM lands and dies with -15.

open_children now takes the process spawner the same way it takes the
platform, and only needs the platform to signal and stop children. The test
uses a fake child whose wait() blocks on an event and a fake platform that
records stop_group, so it checks the shielded stop-then-wait with no real
process and no sleep. It also checks the child is still being waited on,
because anyio's wait_all_tasks_blocked on asyncio can return one loop tick
before a task group finishes.

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:

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)

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)

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()))
 

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)