mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-11 03:38:38 +00:00
9b53ee878f
* perf(ci): distribute unit cases and reuse security sweep connections * revert(ci): reject first speed trial after new failures * perf(ci): avoid repeated test discovery and parsing * perf(ci): isolate parallel security workers * revert(ci): restore measured security execution isolation * perf(tests): reuse source scans and deduplicate cache cleanup * perf(ci): restore parallel security with safe log batching * test(ci): isolate fixture state and schema comparisons * test(ci): make schema and scheduler fixtures deterministic * test(ci): drain queued logging before fixture teardown * test(ci): preserve logging delivery without idle worker restarts * test(ci): keep guardrail sinks available during parallel bursts * test(ci): isolate full proxy lifecycle from gateway imports * test(ci): preserve first failure responses and sweep evidence * test(ci): restore documentation URL settings between cases * test(ci): isolate open MCP tools in private gateways * test(ci): group private MCP cases and handle exited workers * test(ci): retain proxy exits and MCP failure details * test(ci): repair merge interactions with main Key the schema shard's full clone on artifact-name, since main removed the unit-flag input. Give each owned Redis its own log file again. Hand the realtime transcription tests main added a logging object whose dispatch handlers are awaitable, so the per-test logging drain no longer waits on a MagicMock task forever * perf(ci): shard owned integration cases and close lifecycle leaks * fix(ci): qualify shard execution and harden interrupted cleanup * fix(ci): bound worker supervision and own fault traffic * perf(ci): reuse unparametrized fixtures across bounded case groups * fix(ci): preserve boot fault allowances and complete exit records * test(ci): tolerate startup retries and retain readiness details * perf(ci): balance integration shards by recorded duration * fix(ci): load migration database helpers only for workers
82 lines
3 KiB
Python
82 lines
3 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections.abc import Awaitable, Callable
|
|
from contextlib import suppress
|
|
from threading import Thread
|
|
from typing import Final
|
|
|
|
Handler = Callable[[asyncio.StreamReader, asyncio.StreamWriter], Awaitable[None]]
|
|
|
|
|
|
async def close_writer(writer: asyncio.StreamWriter) -> None:
|
|
writer.close()
|
|
with suppress(ConnectionError, OSError):
|
|
await writer.wait_closed()
|
|
|
|
|
|
class LoopbackServer:
|
|
def __init__(self, port: int, handler: Handler) -> None:
|
|
self.loop: Final = asyncio.new_event_loop()
|
|
self._port: Final = port
|
|
self._handler: Final = handler
|
|
self._server: asyncio.Server | None = None
|
|
self._thread: Final = Thread(target=self._run, name=f"integration-relay-{port}", daemon=True)
|
|
self._connections: tuple[asyncio.Task[None], ...] = ()
|
|
|
|
def start(self) -> None:
|
|
try:
|
|
self._thread.start()
|
|
asyncio.run_coroutine_threadsafe(self._bind(), self.loop).result(timeout=10)
|
|
except BaseException:
|
|
self.stop()
|
|
raise
|
|
|
|
def stop(self) -> None:
|
|
if not self._thread.is_alive():
|
|
if not self.loop.is_closed():
|
|
self.loop.close()
|
|
return
|
|
try:
|
|
asyncio.run_coroutine_threadsafe(self._close(), self.loop).result(timeout=10)
|
|
finally:
|
|
self.loop.call_soon_threadsafe(self.loop.stop)
|
|
self._thread.join(10)
|
|
assert not self._thread.is_alive(), f"Database relay thread on port {self._port} did not stop"
|
|
|
|
def _run(self) -> None:
|
|
asyncio.set_event_loop(self.loop)
|
|
try:
|
|
self.loop.run_forever()
|
|
finally:
|
|
self.loop.run_until_complete(self.loop.shutdown_asyncgens())
|
|
self.loop.close()
|
|
|
|
async def _bind(self) -> None:
|
|
self._server = await asyncio.start_server(self._serve, "127.0.0.1", self._port)
|
|
|
|
async def _serve(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
|
|
task: Final = asyncio.current_task()
|
|
assert task is not None
|
|
self._connections = (*self._connections, task)
|
|
try:
|
|
await self._handler(reader, writer)
|
|
finally:
|
|
await close_writer(writer)
|
|
self._connections = tuple(connection for connection in self._connections if connection is not task)
|
|
|
|
async def _close(self) -> None:
|
|
server: Final = self._server
|
|
if server is not None:
|
|
server.close()
|
|
connections: Final = self._connections
|
|
for task in connections:
|
|
task.cancel()
|
|
await asyncio.gather(*connections, return_exceptions=True)
|
|
current: Final = asyncio.current_task()
|
|
remaining: Final = tuple(task for task in asyncio.all_tasks(self.loop) if task is not current)
|
|
for task in remaining:
|
|
task.cancel()
|
|
await asyncio.gather(*remaining, return_exceptions=True)
|
|
if server is not None:
|
|
await server.wait_closed()
|