mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-05 07:37:50 +00:00
fix(proxy): treat Postgres restart and shutdown errors as DB unavailable, not poison spend-log rows (#50132)
* fix(proxy): treat Postgres restart and shutdown errors as DB unavailable, not poison spend-log rows Postgres refuses every new connection with SQLSTATE 57P03 while it starts up, shuts down or replays WAL, and prisma surfaces that as a base DataError with no code. The spend-log writer took it for a poison row, bisected the batch down to single rows and dropped each one. Spend logs now retry and requeue on 57P03 and on server-terminated sessions (57P01, 57P02, 57P05), tool and model usage rollups requeue when the write provably never reached the server, the tool registry requeues its rows, and auth answers 503 instead of 401 while the auth database restarts. Co-authored-by: Dinu Arateanu <dinu.arateanu@rbinternational.com> * fix(proxy): requeue tool discovery rows only when the write never reached Postgres, keep the auth fallback to 57P03 A terminated session (57P01, 57P02, 57P05) may have committed the tool upsert, and the upsert adds call_count on conflict, so replaying it counts the calls twice. The non-Redis tool discovery flush now uses the same never-sent rule as the rollups, which also stops a permanently rejected row (54000) from coming back every tick. One terminated session does not prove the database is down, so it no longer opens the allow_requests_on_db_unavailable fallback, it still answers 503 * fix(proxy): a full server (53300, P2024) also never took the write, so counters and tool rows are sent again Postgres refuses the connection before any statement runs when it is out of slots, and a prisma pool timeout never reaches the server either, so those failures are as safe to replay as 57P03 and P1001 --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Dinu Arateanu <dinu.arateanu@rbinternational.com>
This commit is contained in:
Kerry Lu
GitHub
parent
6173ae16db
commit
b080c59e78
8 changed files with 435 additions and 22 deletions
|
|
@ -55,3 +55,10 @@ class ToolDiscoveryQueue:
|
|||
"""Return and clear all pending rows."""
|
||||
items, self._pending = list(self._pending.values()), {}
|
||||
return items
|
||||
|
||||
def requeue(self, items: Iterable[ToolDiscoveryQueueItem]) -> None:
|
||||
"""Put rows from a failed write back, ahead of and merged with anything queued since the flush."""
|
||||
since_flush: Final = self._pending
|
||||
self._pending = {}
|
||||
for item in (*items, *since_flush.values()):
|
||||
self.add_update(item)
|
||||
|
|
|
|||
|
|
@ -31,6 +31,24 @@ _CONNECTION_CAPACITY_PHRASES: Final = (
|
|||
"too many connections for database",
|
||||
"remaining connection slots are reserved",
|
||||
)
|
||||
_NOT_ACCEPTING_CONNECTIONS_SQLSTATE: Final = "57P03"
|
||||
_NOT_ACCEPTING_CONNECTIONS_PHRASES: Final = (
|
||||
"the database system is starting up",
|
||||
"the database system is shutting down",
|
||||
"the database system is not yet accepting connections",
|
||||
"the database system is not accepting connections",
|
||||
"the database system is in recovery mode",
|
||||
)
|
||||
_SERVER_TERMINATED_SQLSTATES: Final = frozenset({"57P01", "57P02", "57P05"})
|
||||
_SERVER_TERMINATED_PHRASES: Final = (
|
||||
"terminating connection due to administrator command",
|
||||
"terminating connection due to immediate shutdown command",
|
||||
"terminating connection because of unexpected sigquit signal",
|
||||
"terminating connection because of crash of another server process",
|
||||
"terminating connection due to idle-session timeout",
|
||||
)
|
||||
_UNREACHABLE_PRISMA_CODE: Final = "P1001"
|
||||
_UNREACHABLE_PHRASES: Final = ("can't reach database server", "cannot reach database server")
|
||||
_PRISMA_POOL_TIMEOUT_CODE: Final = "P2024"
|
||||
_LOCK_TIMEOUT_SQLSTATE: Final = "55P03"
|
||||
_LOCK_TIMEOUT_PHRASE: Final = "canceling statement due to lock timeout"
|
||||
|
|
@ -139,6 +157,8 @@ class PrismaDBExceptionHandler:
|
|||
return True
|
||||
if PrismaDBExceptionHandler.is_database_capacity_error(e):
|
||||
return True
|
||||
if PrismaDBExceptionHandler.is_database_not_accepting_connections_error(e):
|
||||
return True
|
||||
return isinstance(e, ProxyException) and e.type == ProxyErrorTypes.no_db_connection
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -289,6 +309,39 @@ class PrismaDBExceptionHandler:
|
|||
error_message: Final = str(e).lower()
|
||||
return any(phrase in error_message for phrase in _CONNECTION_CAPACITY_PHRASES)
|
||||
|
||||
@staticmethod
|
||||
def is_database_not_accepting_connections_error(e: Exception) -> bool:
|
||||
"""True iff Postgres refused the connection because it is starting up, shutting
|
||||
down or in crash recovery (SQLSTATE 57P03). The refusal comes before any statement
|
||||
runs, so nothing the write carried was applied."""
|
||||
import prisma
|
||||
|
||||
if not isinstance(e, _exception_types(prisma.errors.PrismaError)):
|
||||
return False
|
||||
if PrismaDBExceptionHandler.postgres_sqlstate(e) == _NOT_ACCEPTING_CONNECTIONS_SQLSTATE:
|
||||
return True
|
||||
error_message: Final = str(e).lower()
|
||||
return any(phrase in error_message for phrase in _NOT_ACCEPTING_CONNECTIONS_PHRASES)
|
||||
|
||||
@staticmethod
|
||||
def is_database_lifecycle_error(e: Exception) -> bool:
|
||||
"""True iff Postgres failed the statement because of a server lifecycle event:
|
||||
a refused connection during startup, shutdown or recovery (57P03), or a session
|
||||
the server terminated (57P01, 57P02, 57P05). The server comes back on its own,
|
||||
so the rows are not at fault. A terminated session may have committed before it
|
||||
was cut, so only ``is_database_not_accepting_connections_error`` proves a
|
||||
non-idempotent write never applied."""
|
||||
import prisma
|
||||
|
||||
if PrismaDBExceptionHandler.is_database_not_accepting_connections_error(e):
|
||||
return True
|
||||
if not isinstance(e, _exception_types(prisma.errors.PrismaError)):
|
||||
return False
|
||||
if PrismaDBExceptionHandler.postgres_sqlstate(e) in _SERVER_TERMINATED_SQLSTATES:
|
||||
return True
|
||||
error_message: Final = str(e).lower()
|
||||
return any(phrase in error_message for phrase in _SERVER_TERMINATED_PHRASES)
|
||||
|
||||
@staticmethod
|
||||
def is_lock_timeout_error(e: Exception) -> bool:
|
||||
"""True iff Postgres cancelled a statement for exceeding ``lock_timeout``
|
||||
|
|
@ -406,6 +459,8 @@ class PrismaDBExceptionHandler:
|
|||
return True
|
||||
if PrismaDBExceptionHandler.is_database_capacity_error(e):
|
||||
return True
|
||||
if PrismaDBExceptionHandler.is_database_lifecycle_error(e):
|
||||
return True
|
||||
if PrismaDBExceptionHandler.is_prisma_engine_internal_error(e):
|
||||
return True
|
||||
if "cached plan must not change result type" in str(e).lower():
|
||||
|
|
@ -680,3 +735,18 @@ def db_write_failure_is_requeue_safe(e: Exception) -> bool:
|
|||
return isinstance(e, DB_RETRY_SAFE_ERROR_TYPES)
|
||||
sqlstate: Final = PrismaDBExceptionHandler.postgres_sqlstate(e)
|
||||
return sqlstate is None or sqlstate[:2] not in _DATA_REJECTED_SQLSTATE_CLASSES
|
||||
|
||||
|
||||
def db_write_never_sent(e: Exception) -> bool:
|
||||
"""Whether the database could not take the write at all, because it was unreachable (P1001),
|
||||
refusing connections (57P03) or full (53300, P2024), so even a non-idempotent write can be sent again."""
|
||||
if PrismaDBExceptionHandler.is_database_not_accepting_connections_error(e):
|
||||
return True
|
||||
if PrismaDBExceptionHandler.is_database_capacity_error(e):
|
||||
return True
|
||||
if _prisma_error_code(e) == _UNREACHABLE_PRISMA_CODE:
|
||||
return True
|
||||
if not PrismaDBExceptionHandler.is_prisma_error(e):
|
||||
return False
|
||||
error_message: Final = str(e).lower()
|
||||
return any(phrase in error_message for phrase in _UNREACHABLE_PHRASES)
|
||||
|
|
|
|||
|
|
@ -23,7 +23,11 @@ from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import (
|
|||
ToolDiscoveryQueue,
|
||||
aggregate_tool_discoveries,
|
||||
)
|
||||
from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry, db_write_failure_is_requeue_safe
|
||||
from litellm.proxy.db.exception_handler import (
|
||||
call_with_db_reconnect_retry,
|
||||
db_write_failure_is_requeue_safe,
|
||||
db_write_never_sent,
|
||||
)
|
||||
from litellm.proxy.db.tool_table_bulk_upsert import build_tool_bulk_upsert
|
||||
from litellm.repositories.object_permission_repository import ObjectPermissionRepository
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
|
|
@ -134,8 +138,19 @@ async def flush_tool_discovery_queue(prisma_client: "PrismaClient", tool_discove
|
|||
return
|
||||
try:
|
||||
await batch_upsert_tools(prisma_client=prisma_client, items=items)
|
||||
except Exception as e: # noqa: BLE001 # a lost tool row only delays its discovery, the scheduler job must not fail
|
||||
verbose_proxy_logger.debug("flush_tool_discovery_queue error (non-blocking): %s", e)
|
||||
except Exception as e: # noqa: BLE001 # the rows go back to the queue or are dropped, the scheduler job must not fail
|
||||
if db_write_never_sent(e):
|
||||
tool_discovery_queue.requeue(items)
|
||||
verbose_proxy_logger.warning(
|
||||
"Tool discovery - %d tool rows went back to the queue for the next tick. Error: %s", len(items), str(e)
|
||||
)
|
||||
return
|
||||
verbose_proxy_logger.error(
|
||||
"Tool discovery - dropped %d tool rows: the failed statement may have applied or the database "
|
||||
"refused the data, so re-sending it is not safe. Error: %s",
|
||||
len(items),
|
||||
str(e),
|
||||
)
|
||||
|
||||
|
||||
async def commit_tool_discoveries(
|
||||
|
|
|
|||
|
|
@ -213,6 +213,7 @@ from litellm.proxy.db.db_url_settings import (
|
|||
from litellm.proxy.db.exception_handler import (
|
||||
PrismaDBExceptionHandler,
|
||||
call_with_db_reconnect_retry,
|
||||
db_write_never_sent,
|
||||
)
|
||||
from litellm.proxy.db.health_check_latest import (
|
||||
LatestHealthCheckRow,
|
||||
|
|
@ -8803,9 +8804,20 @@ async def update_daily_tag_spend(
|
|||
|
||||
def _rollup_batch_never_applied(e: Exception) -> bool:
|
||||
"""True when a drained rollup batch provably never reached its rows: Postgres had no
|
||||
connection for it (53300) or cancelled it under ``lock_timeout`` before it took the row
|
||||
lock (55P03). Either way re-sending the batch cannot double-count, so it is requeued."""
|
||||
return PrismaDBExceptionHandler.is_database_capacity_error(e) or PrismaDBExceptionHandler.is_lock_timeout_error(e)
|
||||
connection for it (53300), was unreachable or refusing connections (P1001, 57P03), or
|
||||
cancelled it under ``lock_timeout`` before it took the row lock (55P03). In each case
|
||||
re-sending the batch cannot double-count, so it is requeued."""
|
||||
return (
|
||||
PrismaDBExceptionHandler.is_database_capacity_error(e)
|
||||
or db_write_never_sent(e)
|
||||
or PrismaDBExceptionHandler.is_lock_timeout_error(e)
|
||||
)
|
||||
|
||||
|
||||
def _rollup_failure_stops_job(e: Exception) -> bool:
|
||||
"""True when the database cannot take any write, so the job stops instead of draining
|
||||
the queues after this one into the same failure."""
|
||||
return PrismaDBExceptionHandler.is_database_capacity_error(e) or db_write_never_sent(e)
|
||||
|
||||
|
||||
async def update_spend_logs_job(
|
||||
|
|
@ -8896,12 +8908,12 @@ async def _run_spend_logs_job(
|
|||
*prisma_client.tool_usage_transactions,
|
||||
]
|
||||
verbose_proxy_logger.warning(
|
||||
"Spend tracking - database out of connections or rollup row lock timed out during tool usage "
|
||||
"Spend tracking - database unavailable or rollup row lock timed out during tool usage "
|
||||
"flush, requeued %s tool usage transactions for the next flush: %s",
|
||||
len(tool_usage_to_process),
|
||||
tool_tracking_err,
|
||||
)
|
||||
if PrismaDBExceptionHandler.is_database_capacity_error(tool_tracking_err):
|
||||
if _rollup_failure_stops_job(tool_tracking_err):
|
||||
raise
|
||||
else:
|
||||
verbose_proxy_logger.error(
|
||||
|
|
@ -8928,12 +8940,12 @@ async def _run_spend_logs_job(
|
|||
*prisma_client.model_usage_transactions,
|
||||
]
|
||||
verbose_proxy_logger.warning(
|
||||
"Spend tracking - database out of connections or rollup row lock timed out during model usage "
|
||||
"Spend tracking - database unavailable or rollup row lock timed out during model usage "
|
||||
"flush, requeued %s model usage transactions for the next flush: %s",
|
||||
len(model_usage_to_process),
|
||||
model_usage_err,
|
||||
)
|
||||
if PrismaDBExceptionHandler.is_database_capacity_error(model_usage_err):
|
||||
if _rollup_failure_stops_job(model_usage_err):
|
||||
raise
|
||||
else:
|
||||
verbose_proxy_logger.error(
|
||||
|
|
@ -9110,6 +9122,7 @@ def _is_transient_spend_log_write_error(e: Exception) -> bool:
|
|||
return (
|
||||
PrismaDBExceptionHandler.is_database_transport_error(e)
|
||||
or PrismaDBExceptionHandler.is_prisma_engine_internal_error(e)
|
||||
or PrismaDBExceptionHandler.is_database_lifecycle_error(e)
|
||||
or PrismaDBExceptionHandler.is_deadlock_error(e)
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -2,7 +2,6 @@
|
|||
Unit tests for ToolDiscoveryQueue.
|
||||
"""
|
||||
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import (
|
||||
|
|
@ -63,3 +62,19 @@ def test_aggregate_sums_call_counts_across_workers_flushes():
|
|||
{"tool_name": "get_weather", "origin": "user_defined", "call_count": 5},
|
||||
{"tool_name": "search", "call_count": 1},
|
||||
)
|
||||
|
||||
|
||||
def test_requeued_rows_go_ahead_of_and_merge_with_sightings_since_the_flush(queue: ToolDiscoveryQueue) -> None:
|
||||
queue.add_update({"tool_name": "get_weather", "origin": "mcp_server"})
|
||||
queue.add_update({"tool_name": "search", "origin": "user_defined"})
|
||||
failed_write = queue.flush()
|
||||
queue.add_update({"tool_name": "translate", "origin": "user_defined"})
|
||||
queue.add_update({"tool_name": "search", "origin": "user_defined"})
|
||||
|
||||
queue.requeue(failed_write)
|
||||
|
||||
assert queue.flush() == [
|
||||
{"tool_name": "get_weather", "origin": "mcp_server", "call_count": 1},
|
||||
{"tool_name": "search", "origin": "user_defined", "call_count": 2},
|
||||
{"tool_name": "translate", "origin": "user_defined", "call_count": 1},
|
||||
]
|
||||
|
|
|
|||
|
|
@ -22,7 +22,7 @@ from prisma.errors import (
|
|||
|
||||
import litellm
|
||||
from litellm.proxy._types import ProxyErrorTypes, ProxyException
|
||||
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
|
||||
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler, db_write_never_sent
|
||||
|
||||
|
||||
# Test is_database_connection_error method
|
||||
|
|
@ -996,3 +996,137 @@ def test_connection_closure_code_survives_prisma_data_error_wrapping() -> None:
|
|||
)
|
||||
assert PrismaDBExceptionHandler.is_database_transport_error(error) is True
|
||||
assert PrismaDBExceptionHandler.is_permanent_database_fault(error) is False
|
||||
|
||||
|
||||
def _connector_error(message: str) -> DataError:
|
||||
return DataError(
|
||||
data={"user_facing_error": {"message": f"Error in connector: Error querying the database: FATAL: {message}"}}
|
||||
)
|
||||
|
||||
|
||||
def _batched_postgres_error(sqlstate: str, message: str) -> DataError:
|
||||
return DataError(
|
||||
data={
|
||||
"user_facing_error": {
|
||||
"message": "Error occurred during query execution: ConnectorError(ConnectorError { user_facing_error: "
|
||||
f'None, kind: QueryError(PostgresError {{ code: "{sqlstate}", message: "{message}", severity: "FATAL" }}) }})'
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _localized_error(sqlstate: str) -> DataError:
|
||||
return DataError(data={"user_facing_error": {"message": "FATAL: Datenbank", "meta": {"code": sqlstate}}})
|
||||
|
||||
|
||||
_NOT_ACCEPTING_CONNECTIONS: Final = (
|
||||
_connector_error("the database system is starting up"),
|
||||
_connector_error("the database system is shutting down"),
|
||||
_connector_error("the database system is not yet accepting connections"),
|
||||
_connector_error("the database system is not accepting connections"),
|
||||
_connector_error("the database system is in recovery mode"),
|
||||
_localized_error("57P03"),
|
||||
_batched_postgres_error("57P03", "the database system is starting up"),
|
||||
EngineRequestError(
|
||||
MagicMock(status=500),
|
||||
'{"is_panic":false,"message":"Error in connector: Error querying the database: FATAL: the database system is '
|
||||
'starting up","backtrace":null}',
|
||||
),
|
||||
RawQueryError(
|
||||
data={
|
||||
"user_facing_error": {
|
||||
"message": "Raw query failed. Code: `57P03`. Message: `db error: FATAL: the database system is starting up`",
|
||||
"meta": {"code": "57P03", "message": "FATAL: the database system is starting up"},
|
||||
"error_code": "P2010",
|
||||
}
|
||||
}
|
||||
),
|
||||
)
|
||||
_SERVER_TERMINATED: Final = (
|
||||
_connector_error("terminating connection due to administrator command"),
|
||||
_connector_error("terminating connection due to immediate shutdown command"),
|
||||
_connector_error("terminating connection because of unexpected SIGQUIT signal"),
|
||||
_connector_error("terminating connection because of crash of another server process"),
|
||||
_connector_error("terminating connection due to idle-session timeout"),
|
||||
_localized_error("57P01"),
|
||||
_localized_error("57P02"),
|
||||
_localized_error("57P05"),
|
||||
_batched_postgres_error("57P01", "terminating connection due to administrator command"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("error", "connection_error"),
|
||||
[
|
||||
*((error, True) for error in _NOT_ACCEPTING_CONNECTIONS),
|
||||
*((error, False) for error in _SERVER_TERMINATED),
|
||||
],
|
||||
)
|
||||
def test_postgres_lifecycle_errors_mean_the_database_is_unavailable(error: Exception, connection_error: bool) -> None:
|
||||
"""A server that is restarting, shutting down or cutting sessions is a transient outage:
|
||||
auth answers 503 rather than 401, and the spend-log writer must not take the rows for poison.
|
||||
Only a server refusing connections opens the ``allow_requests_on_db_unavailable`` fallback,
|
||||
one terminated session does not prove the database is down."""
|
||||
assert {
|
||||
"lifecycle": PrismaDBExceptionHandler.is_database_lifecycle_error(error),
|
||||
"service_unavailable": PrismaDBExceptionHandler.is_database_service_unavailable_error(error),
|
||||
"connection_error": PrismaDBExceptionHandler.is_database_connection_error(error),
|
||||
"permanent": PrismaDBExceptionHandler.is_permanent_database_fault(error),
|
||||
} == {
|
||||
"lifecycle": True,
|
||||
"service_unavailable": True,
|
||||
"connection_error": connection_error,
|
||||
"permanent": False,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("error", "never_sent"),
|
||||
[
|
||||
*((error, True) for error in _NOT_ACCEPTING_CONNECTIONS),
|
||||
*((error, False) for error in _SERVER_TERMINATED),
|
||||
(DataError(data={"user_facing_error": {"message": "Can't reach database server at `db`:`5432`"}}), True),
|
||||
(DataError(data={"user_facing_error": {"error_code": "P1001", "message": "unreachable"}}), True),
|
||||
(_capacity_error("sorry, too many clients already"), True),
|
||||
(_localized_error("53300"), True),
|
||||
(
|
||||
DataError(
|
||||
data={
|
||||
"user_facing_error": {
|
||||
"error_code": "P2024",
|
||||
"message": "Timed out fetching a new connection from the connection pool",
|
||||
}
|
||||
}
|
||||
),
|
||||
True,
|
||||
),
|
||||
(httpx.ReadTimeout("read timed out"), False),
|
||||
(_localized_error("57014"), False),
|
||||
],
|
||||
)
|
||||
def test_db_write_never_sent_only_for_failures_before_any_statement_ran(error: Exception, never_sent: bool) -> None:
|
||||
"""A refused, unreachable or full server provably applied nothing, while a terminated session
|
||||
or a timeout may have committed before the failure, so only the former lets a non-idempotent
|
||||
write be replayed."""
|
||||
assert db_write_never_sent(error) is never_sent
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"error",
|
||||
[
|
||||
_localized_error("57P04"),
|
||||
_localized_error("57014"),
|
||||
_localized_error("23505"),
|
||||
DataError(data={"user_facing_error": {"message": "invalid byte sequence for encoding UTF8: 0x00"}}),
|
||||
RuntimeError("the database system is starting up"),
|
||||
],
|
||||
)
|
||||
def test_other_failures_are_not_postgres_lifecycle_errors(error: Exception) -> None:
|
||||
assert PrismaDBExceptionHandler.is_database_lifecycle_error(error) is False
|
||||
assert db_write_never_sent(error) is False
|
||||
|
||||
|
||||
def test_lifecycle_error_is_seen_through_a_wrapping_exception() -> None:
|
||||
wrapped: Final = RuntimeError("spend flush failed")
|
||||
wrapped.__cause__ = _connector_error("the database system is starting up")
|
||||
assert PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(wrapped) is True
|
||||
|
|
|
|||
|
|
@ -69,9 +69,7 @@ def _make_prisma(
|
|||
prisma.db.litellm_tooltable = MagicMock()
|
||||
prisma.db.litellm_tooltable.upsert = AsyncMock(return_value=upsert_return)
|
||||
prisma.db.execute_raw = AsyncMock(return_value=1)
|
||||
prisma.db.litellm_tooltable.find_many = AsyncMock(
|
||||
return_value=find_many_rows if find_many_rows is not None else []
|
||||
)
|
||||
prisma.db.litellm_tooltable.find_many = AsyncMock(return_value=find_many_rows if find_many_rows is not None else [])
|
||||
prisma.db.litellm_tooltable.find_unique = AsyncMock(return_value=find_unique_row)
|
||||
prisma.db.litellm_verificationtoken = MagicMock()
|
||||
prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=list(key_rows))
|
||||
|
|
@ -227,9 +225,7 @@ async def test_get_tool_found():
|
|||
result = await get_tool(prisma, "my_tool")
|
||||
assert result is not None
|
||||
assert result.tool_name == "my_tool"
|
||||
prisma.db.litellm_tooltable.find_unique.assert_awaited_once_with(
|
||||
where={"tool_name": "my_tool"}
|
||||
)
|
||||
prisma.db.litellm_tooltable.find_unique.assert_awaited_once_with(where={"tool_name": "my_tool"})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -266,9 +262,7 @@ async def test_update_tool_policy_calls_upsert_then_get_tool():
|
|||
updated_by="admin",
|
||||
)
|
||||
prisma = _make_prisma(find_unique_row=row)
|
||||
result = await update_tool_policy(
|
||||
prisma, "my_tool", updated_by="admin", input_policy="blocked"
|
||||
)
|
||||
result = await update_tool_policy(prisma, "my_tool", updated_by="admin", input_policy="blocked")
|
||||
assert result is not None
|
||||
assert result.input_policy == "blocked"
|
||||
prisma.db.litellm_tooltable.upsert.assert_awaited_once()
|
||||
|
|
@ -567,3 +561,42 @@ async def test_failed_tool_table_write_requeues_the_rows_only_when_they_provably
|
|||
left_in_redis = [json.loads(entry) for entry in redis.lists.get(REDIS_TOOL_DISCOVERY_BUFFER_KEY, [])]
|
||||
expected = [[{"tool_name": "get_weather", "origin": "user_defined", "call_count": 1}]] if requeued else []
|
||||
assert left_in_redis == expected, f"requeued={requeued} but redis holds {left_in_redis}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("failure", "requeued"),
|
||||
[
|
||||
(_postgres_rejection("57P03"), True),
|
||||
(_postgres_rejection("53300"), True),
|
||||
(_postgres_rejection("57P01"), False),
|
||||
(_postgres_rejection("54000"), False),
|
||||
(_postgres_rejection("23502"), False),
|
||||
],
|
||||
ids=[
|
||||
"restarting_server_keeps_the_rows",
|
||||
"full_server_keeps_the_rows",
|
||||
"terminated_session_may_have_counted_so_the_rows_are_dropped",
|
||||
"oversized_value_is_dropped",
|
||||
"rejected_row_is_dropped",
|
||||
],
|
||||
)
|
||||
async def test_failed_write_without_redis_keeps_the_rows_in_the_queue_only_when_they_provably_did_not_apply(
|
||||
failure: Exception, requeued: bool
|
||||
):
|
||||
queue = ToolDiscoveryQueue()
|
||||
queue.add_update({"tool_name": "get_weather", "origin": "user_defined"})
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.execute_raw = AsyncMock(side_effect=failure)
|
||||
|
||||
await commit_tool_discoveries(
|
||||
prisma_client=prisma_client,
|
||||
tool_discovery_queue=queue,
|
||||
redis_update_buffer=RedisUpdateBuffer(redis_cache=_InMemoryRedis()),
|
||||
pod_lock_manager=PodLockManager(redis_cache=_InMemoryRedis()),
|
||||
use_redis_buffer=False,
|
||||
)
|
||||
queue.add_update({"tool_name": "get_weather", "origin": "user_defined"})
|
||||
|
||||
expected = [{"tool_name": "get_weather", "origin": "user_defined", "call_count": 2 if requeued else 1}]
|
||||
assert queue.flush() == expected
|
||||
|
|
|
|||
|
|
@ -23,9 +23,9 @@ import pytest
|
|||
from litellm.constants import REDIS_SPEND_LOGS_BUFFER_KEY
|
||||
from litellm.proxy.utils import (
|
||||
MAX_SPEND_LOG_DRAIN_ITERATIONS,
|
||||
drain_spend_logs_queue,
|
||||
monitor_spend_logs_queue,
|
||||
raise_failed_update_spend_exception,
|
||||
drain_spend_logs_queue,
|
||||
recover_parked_spend_logs,
|
||||
update_daily_tag_spend,
|
||||
update_spend,
|
||||
|
|
@ -985,6 +985,85 @@ async def test_update_spend_logs_job_requeues_whole_batch_when_postgres_is_out_o
|
|||
} == {"create_many_calls": 1, "queue_after": ["r1", "r2", "r3"], "backoff_sleeps": []}
|
||||
|
||||
|
||||
def _postgres_restarting() -> Exception:
|
||||
"""prisma's shape for Postgres SQLSTATE 57P03 while the server runs crash recovery after
|
||||
a restart: a base ``DataError`` carrying only the connector message."""
|
||||
from prisma.errors import DataError
|
||||
|
||||
return DataError(
|
||||
data={
|
||||
"user_facing_error": {
|
||||
"is_panic": False,
|
||||
"message": "Error in connector: Error querying the database: FATAL: the database system is starting up",
|
||||
"backtrace": None,
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _postgres_unreachable() -> Exception:
|
||||
from prisma.errors import DataError
|
||||
|
||||
return DataError(
|
||||
data={
|
||||
"user_facing_error": {
|
||||
"is_panic": False,
|
||||
"message": "Can't reach database server at `db.internal`:`5432`",
|
||||
"error_code": "P1001",
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _postgres_terminated_session() -> Exception:
|
||||
from prisma.errors import DataError
|
||||
|
||||
return DataError(
|
||||
data={
|
||||
"user_facing_error": {
|
||||
"is_panic": False,
|
||||
"message": "Error in connector: Error querying the database: FATAL: terminating connection due to "
|
||||
"administrator command",
|
||||
"meta": {"code": "57P01", "message": "terminating connection due to administrator command"},
|
||||
}
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("error", [_postgres_restarting(), _postgres_terminated_session()], ids=["57P03", "57P01"])
|
||||
async def test_update_spend_logs_job_keeps_every_row_while_postgres_restarts(
|
||||
error: Exception,
|
||||
mock_prisma_client: MagicMock,
|
||||
make_spend_log_row: Callable[..., dict[str, object]],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""A restarting server refuses or cuts every statement, so the batch is not bisected
|
||||
into per-row drops: it is retried whole and then goes back to the queue head."""
|
||||
monkeypatch.setattr(asyncio, "sleep", AsyncMock())
|
||||
proxy_logging: Final = MagicMock()
|
||||
proxy_logging.failure_handler = AsyncMock()
|
||||
mock_prisma_client.spend_log_transactions = [
|
||||
make_spend_log_row(request_id="r1"),
|
||||
make_spend_log_row(request_id="r2"),
|
||||
make_spend_log_row(request_id="r3"),
|
||||
make_spend_log_row(request_id="r4"),
|
||||
]
|
||||
mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=error)
|
||||
|
||||
with suppress(Exception):
|
||||
await update_spend_logs_job(
|
||||
prisma_client=mock_prisma_client,
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=proxy_logging,
|
||||
)
|
||||
|
||||
assert {
|
||||
"create_many_calls": mock_prisma_client.db.litellm_spendlogs.create_many.await_count,
|
||||
"queue_after": [row["request_id"] for row in mock_prisma_client.spend_log_transactions],
|
||||
} == {"create_many_calls": 4, "queue_after": ["r1", "r2", "r3", "r4"]}
|
||||
|
||||
|
||||
def _tool_usage_transaction(request_id: str) -> object:
|
||||
from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction
|
||||
|
||||
|
|
@ -1177,6 +1256,53 @@ async def test_model_usage_flush_requeues_and_stops_the_job_when_postgres_is_out
|
|||
} == {"queue_after": [first, second], "autorouter_queue_after": 1}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("error", [_postgres_restarting(), _postgres_unreachable()], ids=["57P03", "P1001"])
|
||||
async def test_rollups_requeue_and_stop_the_job_while_postgres_cannot_take_writes(
|
||||
error: Exception,
|
||||
mock_prisma_client: MagicMock,
|
||||
) -> None:
|
||||
"""A server that is restarting or unreachable never saw the rollup batch, so the tool
|
||||
usage increments are kept for the next flush, and the job stops before draining the
|
||||
auto-router queue into the same outage."""
|
||||
mock_prisma_client.db.litellm_spendlogtoolindex.create_many = AsyncMock(side_effect=error)
|
||||
mock_prisma_client.spend_log_transactions = []
|
||||
first, second = _tool_usage_transaction("r1"), _tool_usage_transaction("r2")
|
||||
mock_prisma_client.tool_usage_transactions = [first, second]
|
||||
mock_prisma_client.autorouter_turn_transactions = [MagicMock(name="autorouter_turn")]
|
||||
|
||||
with suppress(Exception):
|
||||
await update_spend_logs_job(
|
||||
prisma_client=mock_prisma_client,
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
assert {
|
||||
"queue_after": mock_prisma_client.tool_usage_transactions,
|
||||
"autorouter_queue_after": len(mock_prisma_client.autorouter_turn_transactions),
|
||||
} == {"queue_after": [first, second], "autorouter_queue_after": 1}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rollups_drop_a_batch_whose_session_the_server_terminated(
|
||||
mock_prisma_client: MagicMock,
|
||||
) -> None:
|
||||
"""A terminated session (57P01) may have committed before it was cut, and the rollup
|
||||
increments are not idempotent, so the batch is dropped rather than counted twice."""
|
||||
mock_prisma_client.db.litellm_spendlogtoolindex.create_many = AsyncMock(side_effect=_postgres_terminated_session())
|
||||
mock_prisma_client.spend_log_transactions = []
|
||||
mock_prisma_client.tool_usage_transactions = [_tool_usage_transaction("r1")]
|
||||
|
||||
await update_spend_logs_job(
|
||||
prisma_client=mock_prisma_client,
|
||||
db_writer_client=None,
|
||||
proxy_logging_obj=MagicMock(),
|
||||
)
|
||||
|
||||
assert mock_prisma_client.tool_usage_transactions == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_usage_flush_still_drops_a_batch_that_may_have_applied(
|
||||
mock_prisma_client: MagicMock,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue