mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
fix(proxy): let the OS assign Prisma engine ports (#50850)
* fix(proxy): let the OS assign Prisma engine ports Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): clear Prisma connection state after failed startup Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: mateo <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
87be092af7
commit
52cd9ef0e4
7 changed files with 340 additions and 8 deletions
|
|
@ -398,12 +398,26 @@ class PrismaWrapper:
|
|||
self._retirement_tasks = self._retirement_tasks.difference((retirement_task,))
|
||||
|
||||
async def connect(self, timeout: int | timedelta | None = None) -> None:
|
||||
if timeout is None:
|
||||
await self._original_prisma.connect()
|
||||
else:
|
||||
await self._original_prisma.connect(timeout)
|
||||
await self._connect_prisma(self._original_prisma, timeout)
|
||||
self._active_drain_tracker = self._instrument_prisma_client(self._original_prisma)
|
||||
|
||||
@staticmethod
|
||||
async def _connect_prisma(prisma_client: "Prisma", timeout: int | timedelta | None = None) -> None:
|
||||
new_connection: Final = prisma_client.is_connected() is False
|
||||
if new_connection:
|
||||
from litellm.proxy.db.prisma_engine import initialize_query_engine
|
||||
|
||||
initialize_query_engine(prisma_client)
|
||||
try:
|
||||
if timeout is None:
|
||||
await prisma_client.connect()
|
||||
else:
|
||||
await prisma_client.connect(timeout)
|
||||
except (Exception, asyncio.CancelledError):
|
||||
if new_connection:
|
||||
await prisma_client.disconnect()
|
||||
raise
|
||||
|
||||
@staticmethod
|
||||
async def _kill_engine_process(pid: int) -> None:
|
||||
"""Force-kill the engine subprocess to prevent DB connection pool leaks.
|
||||
|
|
@ -688,7 +702,7 @@ class PrismaWrapper:
|
|||
else:
|
||||
self._original_prisma = Prisma()
|
||||
|
||||
await self._original_prisma.connect()
|
||||
await self._connect_prisma(self._original_prisma)
|
||||
self._active_drain_tracker = self._instrument_prisma_client(self._original_prisma)
|
||||
self._engine_generation += 1
|
||||
|
||||
|
|
@ -719,7 +733,7 @@ class PrismaWrapper:
|
|||
if old_drain_tracker is None:
|
||||
old_drain_tracker = self._instrument_prisma_client(old_prisma)
|
||||
try:
|
||||
await replacement_prisma.connect()
|
||||
await self._connect_prisma(replacement_prisma)
|
||||
except (Exception, asyncio.CancelledError):
|
||||
self._schedule_engine_retirement(self._get_engine_pid(replacement_prisma), None)
|
||||
raise
|
||||
|
|
|
|||
136
litellm/proxy/db/prisma_engine.py
Normal file
136
litellm/proxy/db/prisma_engine.py
Normal file
|
|
@ -0,0 +1,136 @@
|
|||
import asyncio
|
||||
import os
|
||||
import signal
|
||||
import subprocess
|
||||
import sys
|
||||
from collections.abc import Sequence
|
||||
from datetime import timedelta
|
||||
from pathlib import Path
|
||||
from threading import Thread
|
||||
from typing import TYPE_CHECKING, BinaryIO, Final, Literal
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, ValidationError
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import Prisma
|
||||
from prisma.types import DatasourceOverride
|
||||
|
||||
|
||||
class _AddressFields(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
ip: Literal["127.0.0.1"]
|
||||
port: int = Field(gt=0, le=65535)
|
||||
message: str
|
||||
|
||||
|
||||
class _ListeningEvent(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
target: Literal["query_engine::server"]
|
||||
fields: _AddressFields
|
||||
|
||||
|
||||
def _listening_url(line: bytes) -> str | None:
|
||||
try:
|
||||
event: Final = _ListeningEvent.model_validate_json(line)
|
||||
except ValidationError:
|
||||
return None
|
||||
if not event.fields.message.startswith("Started query engine http server on "):
|
||||
return None
|
||||
return f"http://{event.fields.ip}:{event.fields.port}"
|
||||
|
||||
|
||||
def _publish_address(ready: asyncio.Future[str | None], address: str | None) -> None:
|
||||
if not ready.done():
|
||||
ready.set_result(address)
|
||||
|
||||
|
||||
def _forward_output(stream: BinaryIO, loop: asyncio.AbstractEventLoop, ready: asyncio.Future[str | None]) -> None:
|
||||
def publish(address: str | None) -> None:
|
||||
try:
|
||||
loop.call_soon_threadsafe(_publish_address, ready, address)
|
||||
except RuntimeError:
|
||||
pass
|
||||
|
||||
try:
|
||||
with stream:
|
||||
for line in stream:
|
||||
sys.stdout.write(line.decode("utf-8", errors="replace"))
|
||||
sys.stdout.flush()
|
||||
if (address := _listening_url(line)) is not None:
|
||||
publish(address)
|
||||
finally:
|
||||
publish(None)
|
||||
|
||||
|
||||
def _unblock_signals() -> None:
|
||||
signal.pthread_sigmask(signal.SIG_UNBLOCK, (signal.SIGINT, signal.SIGTERM))
|
||||
|
||||
|
||||
def initialize_query_engine(client: "Prisma") -> None:
|
||||
from prisma._constants import DEFAULT_CONNECT_TIMEOUT
|
||||
from prisma.builder import dumps
|
||||
from prisma.client import PACKAGED_SCHEMA_PATH
|
||||
from prisma.engine import errors
|
||||
from prisma.engine.query import QueryEngine
|
||||
from prisma.utils import (
|
||||
DEBUG,
|
||||
_env_bool, # pyright: ignore[reportPrivateUsage] # Preserve Prisma's playground flag semantics
|
||||
)
|
||||
|
||||
class PortOwningQueryEngine(QueryEngine):
|
||||
async def spawn(
|
||||
self,
|
||||
file: Path,
|
||||
timeout: timedelta = DEFAULT_CONNECT_TIMEOUT,
|
||||
datasources: Sequence["DatasourceOverride"] | None = None,
|
||||
) -> None:
|
||||
playground: Final = _env_bool("__PRISMA_PY_PLAYGROUND")
|
||||
environment: Final = {
|
||||
**os.environ,
|
||||
"PRISMA_DML_PATH": str(self.dml_path.absolute()),
|
||||
"RUST_LOG": "info" if DEBUG or playground else "error,query_engine::server=info",
|
||||
"RUST_LOG_FORMAT": "json",
|
||||
"PRISMA_CLIENT_ENGINE_TYPE": "binary",
|
||||
"PRISMA_ENGINE_PROTOCOL": "graphql",
|
||||
**({"OVERWRITE_DATASOURCES": dumps(datasources)} if datasources is not None else {}),
|
||||
**({"LOG_QUERIES": "y"} if self._log_queries else {}),
|
||||
}
|
||||
self.process = subprocess.Popen(
|
||||
(
|
||||
str(file.absolute()),
|
||||
"-p",
|
||||
"0",
|
||||
"--enable-metrics",
|
||||
"--enable-raw-queries",
|
||||
*(("--enable-playground",) if playground else ()),
|
||||
),
|
||||
env=environment,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=sys.stderr,
|
||||
preexec_fn=None if sys.platform == "win32" else _unblock_signals,
|
||||
)
|
||||
stream: Final = self.process.stdout
|
||||
if stream is None:
|
||||
raise errors.EngineConnectionError("Query engine output pipe is unavailable")
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
ready: Final[asyncio.Future[str | None]] = loop.create_future()
|
||||
Thread(target=_forward_output, args=(stream, loop, ready), daemon=True).start()
|
||||
try:
|
||||
async with asyncio.timeout(timeout.total_seconds()):
|
||||
address: Final = await ready
|
||||
if address is None:
|
||||
raise errors.EngineConnectionError(
|
||||
"Query engine exited before publishing its listening address"
|
||||
)
|
||||
self.url = address
|
||||
await self.request("GET", "/status")
|
||||
except TimeoutError as error:
|
||||
raise errors.EngineConnectionError(
|
||||
"Timed out waiting for the query engine's listening address"
|
||||
) from error
|
||||
|
||||
client._engine = PortOwningQueryEngine( # pyright: ignore[reportPrivateUsage] # Prisma exposes engine injection only through this setter # rebind-ok: install the lifecycle-owned transport
|
||||
dml_path=PACKAGED_SCHEMA_PATH,
|
||||
log_queries=client._log_queries, # pyright: ignore[reportPrivateUsage] # Preserve the generated client's query logging
|
||||
**client._http_config, # pyright: ignore[reportPrivateUsage] # Preserve the generated client's HTTP options
|
||||
)
|
||||
99
tests/integration/database/test_prisma_engine_port.py
Normal file
99
tests/integration/database/test_prisma_engine_port.py
Normal file
|
|
@ -0,0 +1,99 @@
|
|||
import asyncio
|
||||
import os
|
||||
import socket
|
||||
from datetime import timedelta
|
||||
from typing import Final
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
import pytest
|
||||
from prisma import Prisma
|
||||
from prisma.engine.errors import AlreadyConnectedError, EngineConnectionError
|
||||
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.proxy.db.prisma_client import PrismaWrapper
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from tests.integration._support.database import scratch_database
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("occupied", (False, True), ids=("normal", "occupied-port"))
|
||||
async def test_prisma_connect_owns_its_listening_port(occupied: bool, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as competitor:
|
||||
competitor.bind(("127.0.0.1", 0)) # test-quality-ok: keep the socket bound to force Prisma's port collision
|
||||
port: Final = competitor.getsockname()[1]
|
||||
if occupied:
|
||||
monkeypatch.setattr("prisma.engine.utils.get_open_port", lambda: port)
|
||||
monkeypatch.delenv("DATABASE_URL_READ_REPLICA", raising=False)
|
||||
client: Final = PrismaClient(os.environ["DATABASE_URL"], ProxyLogging(DualCache()))
|
||||
try:
|
||||
await client.db.connect(timeout=timedelta(seconds=5))
|
||||
with pytest.raises(AlreadyConnectedError):
|
||||
await client.db.connect(timeout=timedelta(seconds=5))
|
||||
assert await client.db.query_raw("SELECT 1 AS result") == [{"result": 1}]
|
||||
assert client.db.is_connected() is True
|
||||
finally:
|
||||
await client.db.disconnect()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_prisma_engines_query_their_own_databases() -> None:
|
||||
with scratch_database("engine_first") as first, scratch_database("engine_second") as second:
|
||||
clients: Final = tuple(PrismaWrapper(Prisma(datasource={"url": url})) for url in (first, second))
|
||||
try:
|
||||
await asyncio.gather(*(client.connect() for client in clients))
|
||||
results: Final = await asyncio.gather(
|
||||
*(client.query_raw("SELECT current_database() AS name") for client in clients)
|
||||
)
|
||||
assert results == [[{"name": urlsplit(url).path.lstrip("/")}] for url in (first, second)]
|
||||
finally:
|
||||
await asyncio.gather(*(client.disconnect() for client in clients))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_prisma_startup_database_failure_closes_the_engine(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as unavailable:
|
||||
unavailable.bind( # test-quality-ok: socket remains bound throughout the connection attempt
|
||||
("127.0.0.1", 0)
|
||||
)
|
||||
url: Final = (
|
||||
f"postgresql://invalid:invalid@127.0.0.1:{unavailable.getsockname()[1]}/unavailable?connect_timeout=1"
|
||||
)
|
||||
original: Final = Prisma(datasource={"url": url})
|
||||
client: Final = PrismaWrapper(original)
|
||||
connecting: Final = asyncio.create_task(client.connect(timeout=timedelta(seconds=5)))
|
||||
try:
|
||||
await asyncio.sleep(0)
|
||||
engine: Final = original._engine
|
||||
with pytest.raises(EngineConnectionError):
|
||||
await connecting
|
||||
assert original.is_connected() is False
|
||||
assert engine.process is None
|
||||
assert engine.session.should_close() is False
|
||||
monkeypatch.setattr(original, "_datasource", {"url": os.environ["DATABASE_URL"]})
|
||||
await client.connect(timeout=timedelta(seconds=5))
|
||||
assert await client.query_raw("SELECT 1 AS result") == [{"result": 1}]
|
||||
assert original.is_connected() is True
|
||||
finally:
|
||||
await client.disconnect()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancelled_prisma_startup_stops_its_engine() -> None:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as stalled_database:
|
||||
stalled_database.bind(("127.0.0.1", 0)) # test-quality-ok: socket remains bound during the stalled startup
|
||||
stalled_database.listen()
|
||||
url: Final = f"postgresql://invalid:invalid@127.0.0.1:{stalled_database.getsockname()[1]}/stalled"
|
||||
original: Final = Prisma(datasource={"url": url})
|
||||
client: Final = PrismaWrapper(original)
|
||||
connecting: Final = asyncio.create_task(client.connect())
|
||||
try:
|
||||
await asyncio.sleep(0)
|
||||
process: Final = original._engine.process
|
||||
assert process is not None
|
||||
connecting.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await connecting
|
||||
assert process.poll() is not None
|
||||
assert original.is_connected() is False
|
||||
finally:
|
||||
await client.disconnect()
|
||||
72
tests/unit/proxy/db/test_prisma_engine.py
Normal file
72
tests/unit/proxy/db/test_prisma_engine.py
Normal file
|
|
@ -0,0 +1,72 @@
|
|||
import asyncio
|
||||
from typing import Final, NoReturn
|
||||
|
||||
import pytest
|
||||
from prisma import Prisma
|
||||
from prisma.engine.errors import EngineConnectionError
|
||||
from prisma.engine.query import QueryEngine
|
||||
|
||||
from litellm.proxy.db.prisma_client import PrismaWrapper
|
||||
from litellm.proxy.db.prisma_engine import _listening_url
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("line", "expected"),
|
||||
(
|
||||
(
|
||||
b'{"target":"query_engine::server","fields":{"ip":"127.0.0.1","port":"31415",'
|
||||
b'"message":"Started query engine http server on http://127.0.0.1:31415"}}',
|
||||
"http://127.0.0.1:31415",
|
||||
),
|
||||
(b"", None),
|
||||
(b"query-engine version", None),
|
||||
(b'{"target":"query_engine","fields":{"message":"PANIC","reason":"Address already in use"}}', None),
|
||||
(
|
||||
b'{"target":"other","fields":{"ip":"127.0.0.1","port":"31415",'
|
||||
b'"message":"Started query engine http server on http://127.0.0.1:31415"}}',
|
||||
None,
|
||||
),
|
||||
(
|
||||
b'{"target":"query_engine::server","fields":{"ip":"0.0.0.0","port":"31415",'
|
||||
b'"message":"Started query engine http server on http://0.0.0.0:31415"}}',
|
||||
None,
|
||||
),
|
||||
(
|
||||
b'{"target":"query_engine::server","fields":{"ip":"127.0.0.1","port":"0",'
|
||||
b'"message":"Started query engine http server on http://127.0.0.1:0"}}',
|
||||
None,
|
||||
),
|
||||
(
|
||||
b'{"target":"query_engine::server","fields":{"ip":"127.0.0.1","port":"65536",'
|
||||
b'"message":"Started query engine http server on http://127.0.0.1:65536"}}',
|
||||
None,
|
||||
),
|
||||
(
|
||||
b'{"target":"query_engine::server","fields":{"ip":"127.0.0.1","port":"31415",'
|
||||
b'"message":"Other server event"}}',
|
||||
None,
|
||||
),
|
||||
),
|
||||
)
|
||||
def test_listening_url_only_accepts_an_engine_listening_event(line: bytes, expected: str | None) -> None:
|
||||
assert _listening_url(line) == expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("fault", (EngineConnectionError("engine startup failed"), asyncio.CancelledError()))
|
||||
async def test_failed_engine_startup_clears_connection_state(
|
||||
fault: EngineConnectionError | asyncio.CancelledError, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
def unavailable_engine(_: QueryEngine) -> NoReturn:
|
||||
raise fault
|
||||
|
||||
monkeypatch.setattr(QueryEngine, "_ensure_file", unavailable_engine)
|
||||
original: Final = Prisma(datasource={"url": "postgresql://unused@localhost/unused"})
|
||||
client: Final = PrismaWrapper(original)
|
||||
try:
|
||||
with pytest.raises(type(fault)) as caught:
|
||||
await client.connect()
|
||||
assert caught.value is fault
|
||||
assert original.is_connected() is False
|
||||
finally:
|
||||
await client.disconnect()
|
||||
|
|
@ -574,6 +574,9 @@ async def test_writer_recreate_passes_http_client_through(monkeypatch):
|
|||
def __init__(self, **kwargs):
|
||||
captured_kwargs.update(kwargs)
|
||||
|
||||
def is_connected(self):
|
||||
return True
|
||||
|
||||
async def connect(self):
|
||||
return None
|
||||
|
||||
|
|
@ -875,6 +878,9 @@ async def test_reader_recreate_uses_datasource_override(monkeypatch):
|
|||
def __init__(self, **kwargs):
|
||||
captured_kwargs.update(kwargs)
|
||||
|
||||
def is_connected(self):
|
||||
return True
|
||||
|
||||
async def connect(self):
|
||||
return None
|
||||
|
||||
|
|
@ -912,6 +918,9 @@ async def test_writer_recreate_does_not_use_datasource(monkeypatch):
|
|||
def __init__(self, **kwargs):
|
||||
captured_kwargs.update(kwargs)
|
||||
|
||||
def is_connected(self):
|
||||
return True
|
||||
|
||||
async def connect(self):
|
||||
return None
|
||||
|
||||
|
|
|
|||
|
|
@ -8298,6 +8298,8 @@ class _ReadOnlyReplicaPrisma:
|
|||
def __init__(self, journal: _PrismaJournal, role: str) -> None:
|
||||
self.journal = journal
|
||||
self.role = role
|
||||
self._log_queries: Final = False
|
||||
self._http_config: Final[dict[str, object]] = {}
|
||||
|
||||
def is_connected(self) -> bool:
|
||||
return False
|
||||
|
|
|
|||
|
|
@ -2713,7 +2713,7 @@ async def test_proxy_server_prisma_setup_invalid_db(monkeypatch):
|
|||
|
||||
Think 2-3 times before editing / deleting this test, it's important for PROD
|
||||
"""
|
||||
import httpx
|
||||
from prisma.engine.errors import EngineConnectionError
|
||||
|
||||
from litellm.proxy.proxy_server import ProxyStartupEvent
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
|
@ -2724,7 +2724,7 @@ async def test_proxy_server_prisma_setup_invalid_db(monkeypatch):
|
|||
|
||||
monkeypatch.setenv("DATABASE_URL", invalid_db_url)
|
||||
|
||||
with pytest.raises(httpx.ConnectError):
|
||||
with pytest.raises(EngineConnectionError):
|
||||
await ProxyStartupEvent._setup_prisma_client(
|
||||
database_url=invalid_db_url,
|
||||
proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache),
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue