mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-03 07:23:33 +00:00
180 lines
7.1 KiB
Python
180 lines
7.1 KiB
Python
import random
|
|
import re
|
|
import time
|
|
from collections.abc import Generator, Mapping
|
|
from contextlib import contextmanager
|
|
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING, Final
|
|
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
|
|
|
|
from litellm_proxy_extras._logging import logger
|
|
from litellm_proxy_extras._redaction import _redact_credentials
|
|
from litellm_proxy_extras.prisma_toolchain import (
|
|
MIGRATION_LOCK_TIMEOUT_ENV_VAR,
|
|
migration_lock_timeout,
|
|
)
|
|
|
|
MIGRATION_LOCK_KEY: Final = int.from_bytes(b"llm_mig2", "big")
|
|
PRISMA_ADVISORY_LOCK_KEY: Final = 72707369
|
|
_SERVER_ANSWERED_MARKERS: Final = ("FATAL:", "SSL", "certificate")
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class AdvisoryLock:
|
|
key: int
|
|
name: str
|
|
holder: str
|
|
|
|
|
|
COORDINATOR_LOCK: Final = AdvisoryLock(
|
|
MIGRATION_LOCK_KEY,
|
|
"v2 migration coordinator lock",
|
|
"another v2 migration resolver",
|
|
)
|
|
PRISMA_LOCK: Final = AdvisoryLock(
|
|
PRISMA_ADVISORY_LOCK_KEY,
|
|
"Prisma migration lock",
|
|
"another Prisma migration",
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
import psycopg
|
|
|
|
|
|
def _migration_url(database_url: str, direct_url: "str | None") -> str:
|
|
if not direct_url:
|
|
return database_url
|
|
schema: Final = next((value for key, value in parse_qsl(urlsplit(database_url).query) if key == "schema"), "public")
|
|
direct: Final = urlsplit(direct_url)
|
|
parameters: Final = tuple((key, value) for key, value in parse_qsl(direct.query) if key != "schema")
|
|
return urlunsplit(direct._replace(query=urlencode((*parameters, ("schema", schema)))))
|
|
|
|
|
|
def migration_environment(environment: Mapping[str, str]) -> Mapping[str, str]:
|
|
database_url: Final = environment.get("DATABASE_URL")
|
|
if not database_url:
|
|
return environment
|
|
migration_url: Final = _migration_url(database_url, environment.get("DIRECT_URL"))
|
|
return {**environment, "DATABASE_URL": migration_url, "PRISMA_SCHEMA_DISABLE_ADVISORY_LOCK": "1"}
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class _LockResult:
|
|
acquired: bool
|
|
|
|
|
|
def _try_lock(connection: "psycopg.Connection[tuple[object, ...]]", key: int = MIGRATION_LOCK_KEY) -> bool:
|
|
from psycopg.rows import class_row
|
|
|
|
with connection.cursor(row_factory=class_row(_LockResult)) as cursor:
|
|
row: Final = cursor.execute("SELECT pg_try_advisory_xact_lock(%s) AS acquired", (key,)).fetchone()
|
|
return row is not None and row.acquired
|
|
|
|
|
|
def _coordination_failure(database_url: str, exc: "psycopg.Error", *, connecting: bool) -> str:
|
|
import psycopg
|
|
|
|
prefix: Final = "Lost or could not establish v2 migration coordination with the database: "
|
|
raw: Final = str(exc)
|
|
unquoted: Final = re.sub(r"\"[^\"]*\"|'[^']*'", "", raw)
|
|
detail: Final = _redact_credentials(raw).strip()
|
|
if (
|
|
not connecting
|
|
or not isinstance(exc, psycopg.OperationalError)
|
|
or any(marker in unquoted for marker in _SERVER_ANSWERED_MARKERS)
|
|
):
|
|
return prefix + detail
|
|
parsed_url: Final = urlsplit(database_url)
|
|
try:
|
|
port: Final = parsed_url.port
|
|
except ValueError:
|
|
return prefix + f"P1001: Can't reach database server. {detail}"
|
|
host: Final = parsed_url.hostname or "localhost"
|
|
return prefix + f"P1001: Can't reach database server at `{host}:{port or 5432}`. {detail}"
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class MigrationCoordinator:
|
|
connection: "psycopg.Connection[tuple[object, ...]]"
|
|
|
|
def check_connection(self) -> None:
|
|
self.connection.execute("SELECT 1")
|
|
|
|
def release_statement_snapshot(self) -> None:
|
|
# a parameterized statement leaves its portal snapshot registered on the backend until the
|
|
# next statement, and CREATE INDEX CONCURRENTLY waits for every such snapshot to go away
|
|
self.connection.execute("SELECT 1")
|
|
|
|
def acquire_prisma_lock(self) -> None:
|
|
deadline: Final = time.monotonic() + migration_lock_timeout()
|
|
while time.monotonic() < deadline:
|
|
if _try_lock(self.connection, PRISMA_ADVISORY_LOCK_KEY):
|
|
return
|
|
time.sleep(min(random.uniform(0.5, 1.5), max(0.0, deadline - time.monotonic())))
|
|
raise RuntimeError(
|
|
"Timed out waiting for Prisma's lock to recover migration history. LiteLLM startup has stopped. "
|
|
"Another migration or a pooled database session may still hold the lock. Check the database lock holder. "
|
|
"When using a transaction pooler, configure DIRECT_URL to reach the same database without the pooler."
|
|
)
|
|
|
|
|
|
@contextmanager
|
|
def migration_lock(
|
|
database_url: str,
|
|
lock: AdvisoryLock = COORDINATOR_LOCK,
|
|
) -> Generator[MigrationCoordinator, None, None]:
|
|
import psycopg
|
|
|
|
wait_seconds: Final = migration_lock_timeout()
|
|
deadline: Final = time.monotonic() + wait_seconds
|
|
try:
|
|
connection: Final = psycopg.connect(database_url, connect_timeout=10, autocommit=True)
|
|
except psycopg.Error as exc:
|
|
raise RuntimeError(_coordination_failure(database_url, exc, connecting=True)) from None
|
|
try:
|
|
with connection:
|
|
coordinator: Final = MigrationCoordinator(connection)
|
|
logger.info("Waiting for the %s (up to %ss)", lock.name, wait_seconds)
|
|
while time.monotonic() < deadline:
|
|
with connection.transaction():
|
|
if _try_lock(connection, lock.key):
|
|
logger.info("Acquired the %s", lock.name)
|
|
|
|
yield coordinator
|
|
coordinator.check_connection()
|
|
return
|
|
time.sleep(min(random.uniform(0.5, 1.5), max(0.0, deadline - time.monotonic())))
|
|
except psycopg.Error as exc:
|
|
raise RuntimeError(_coordination_failure(database_url, exc, connecting=False)) from None
|
|
raise RuntimeError(
|
|
f"Timed out waiting for {lock.holder} after {wait_seconds}s. "
|
|
f"Check the running migration or increase {MIGRATION_LOCK_TIMEOUT_ENV_VAR}."
|
|
)
|
|
|
|
|
|
@contextmanager
|
|
def prisma_migration_lock(database_url: str) -> Generator[None, None, None]:
|
|
"""Hold only Prisma's key as an xact lock on a pinned connection with Prisma's session lock disabled.
|
|
|
|
This keeps deploy independent of the coordinator key held by request-log index builds."""
|
|
with migration_lock(database_url, PRISMA_LOCK) as coordinator:
|
|
coordinator.release_statement_snapshot()
|
|
yield
|
|
|
|
|
|
@contextmanager
|
|
def held_migration_lock(connection: "psycopg.Connection[tuple[object, ...]]") -> Generator[bool, None, None]:
|
|
"""A session-level, non-blocking hold of the migration coordinator lock on an autocommit
|
|
connection, for DDL that cannot run inside a transaction (`CREATE INDEX CONCURRENTLY`).
|
|
Yields whether the lock was acquired; a v2 resolver or another migration job's index build
|
|
holding it yields False. Released on exit."""
|
|
from psycopg.rows import class_row
|
|
|
|
with connection.cursor(row_factory=class_row(_LockResult)) as cursor:
|
|
row: Final = cursor.execute("SELECT pg_try_advisory_lock(%s) AS acquired", (MIGRATION_LOCK_KEY,)).fetchone()
|
|
acquired: Final = row is not None and row.acquired
|
|
try:
|
|
yield acquired
|
|
finally:
|
|
if acquired:
|
|
connection.execute("SELECT pg_advisory_unlock(%s)", (MIGRATION_LOCK_KEY,))
|