scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge pull request #36824 from BerriAI/litellm_/concurrent-view-creation

Browse source 
fix(proxy): tolerate a concurrent creator when creating spend views
This commit is contained in:
yuneng-jiang 2026-08-13 15:20:11 -07:00 • committed by GitHub
commit 0c1355d54a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 119 additions and 30 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

75
litellm/proxy/db/create_views.py
View file

@ -1,15 +1,50 @@
from typing import Any, Final
from typing import Any, Final, Protocol
from litellm import verbose_logger
_db = Any
class SupportsExecuteRaw(Protocol):
"""The one database operation create_view_tolerating_race needs.
Narrower than the `_db = Any` the rest of this module still uses, so the
helper's contract is checkable at its call sites without retyping every
function here.
"""
async def execute_raw(self, query: str, *args: object) -> int: ...
# Markers that indicate a view/relation does not yet exist in the database.
# Keeping these in one place avoids repeating the check across all view blocks
# and prevents overly broad matches (e.g. bare 'undefined' would also match
# 'undefined function' or 'column undefined_col referenced in query').
_VIEW_NOT_FOUND_MARKERS: Final = ("does not exist", "no such table", "undefined table")
# Markers for the inverse condition: another replica created the view between
# our existence probe and our CREATE.
_VIEW_ALREADY_EXISTS_MARKERS: Final = ("already exists", "duplicate object", "duplicate table")
async def create_view_tolerating_race(db: SupportsExecuteRaw, view_name: str, ddl: str) -> None:
"""
Create a view, treating "a concurrent creator won" as success.
Every replica booting against the same fresh database observes the view as
absent and issues the CREATE; Postgres fails all but one with a
duplicate-object error. The desired end state is still reached, so losing
that race is success. Without this, the loser's exception propagates out of
a detached startup task and the remaining views are never created.
"""
try:
await db.execute_raw(ddl)
verbose_logger.debug("%s Created!", view_name)
except Exception as e:
if not any(marker in str(e).lower() for marker in _VIEW_ALREADY_EXISTS_MARKERS):
raise
verbose_logger.debug("%s already created by a concurrent replica", view_name)
async def create_missing_views(db: _db):
"""
@ -34,7 +69,10 @@ async def create_missing_views(db: _db):
if not any(marker in error_msg for marker in _VIEW_NOT_FOUND_MARKERS):
raise
# If an error occurs, the view does not exist, so create it
await db.execute_raw("""
await create_view_tolerating_race(
db,
"LiteLLM_VerificationTokenView",
"""
CREATE VIEW "LiteLLM_VerificationTokenView" AS
SELECT
v.*,
@ -46,9 +84,8 @@ async def create_missing_views(db: _db):
FROM "LiteLLM_VerificationToken" v
LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id = t.team_id
LEFT JOIN "LiteLLM_ProjectTable" p ON v.project_id = p.project_id;
""")
verbose_logger.debug("LiteLLM_VerificationTokenView Created!")
""",
)
try:
await db.query_raw("""SELECT 1 FROM "MonthlyGlobalSpend" LIMIT 1""")
@ -69,9 +106,7 @@ async def create_missing_views(db: _db):
GROUP BY
DATE("startTime");
"""
await db.execute_raw(query=sql_query)
verbose_logger.debug("MonthlyGlobalSpend Created!")
await create_view_tolerating_race(db, "MonthlyGlobalSpend", sql_query)
try:
await db.query_raw("""SELECT 1 FROM "Last30dKeysBySpend" LIMIT 1""")
@ -100,9 +135,7 @@ async def create_missing_views(db: _db):
ORDER BY
total_spend DESC;
"""
await db.execute_raw(query=sql_query)
verbose_logger.debug("Last30dKeysBySpend Created!")
await create_view_tolerating_race(db, "Last30dKeysBySpend", sql_query)
try:
await db.query_raw("""SELECT 1 FROM "Last30dModelsBySpend" LIMIT 1""")
@ -126,9 +159,7 @@ async def create_missing_views(db: _db):
ORDER BY
total_spend DESC;
"""
await db.execute_raw(query=sql_query)
verbose_logger.debug("Last30dModelsBySpend Created!")
await create_view_tolerating_race(db, "Last30dModelsBySpend", sql_query)
try:
await db.query_raw("""SELECT 1 FROM "MonthlyGlobalSpendPerKey" LIMIT 1""")
verbose_logger.debug("MonthlyGlobalSpendPerKey Exists!")
@ -150,9 +181,7 @@ async def create_missing_views(db: _db):
DATE("startTime"),
api_key;
"""
await db.execute_raw(query=sql_query)
verbose_logger.debug("MonthlyGlobalSpendPerKey Created!")
await create_view_tolerating_race(db, "MonthlyGlobalSpendPerKey", sql_query)
try:
await db.query_raw("""SELECT 1 FROM "MonthlyGlobalSpendPerUserPerKey" LIMIT 1""")
verbose_logger.debug("MonthlyGlobalSpendPerUserPerKey Exists!")
@ -176,9 +205,7 @@ async def create_missing_views(db: _db):
"user",
api_key;
"""
await db.execute_raw(query=sql_query)
verbose_logger.debug("MonthlyGlobalSpendPerUserPerKey Created!")
await create_view_tolerating_race(db, "MonthlyGlobalSpendPerUserPerKey", sql_query)
try:
await db.query_raw("""SELECT 1 FROM "DailyTagSpend" LIMIT 1""")
@ -197,9 +224,7 @@ async def create_missing_views(db: _db):
FROM "LiteLLM_SpendLogs" s
GROUP BY individual_request_tag, DATE(s."startTime");
"""
await db.execute_raw(query=sql_query)
verbose_logger.debug("DailyTagSpend Created!")
await create_view_tolerating_race(db, "DailyTagSpend", sql_query)
try:
await db.query_raw("""SELECT 1 FROM "Last30dTopEndUsersSpend" LIMIT 1""")
@ -218,9 +243,7 @@ async def create_missing_views(db: _db):
ORDER BY total_spend DESC
LIMIT 100;
"""
await db.execute_raw(query=sql_query)
verbose_logger.debug("Last30dTopEndUsersSpend Created!")
await create_view_tolerating_race(db, "Last30dTopEndUsersSpend", sql_query)
async def should_create_missing_views(db: _db) -> bool:

11
litellm/proxy/utils.py
View file

@ -106,6 +106,7 @@ from litellm.proxy.common_utils.config_sync_pubsub import publish_config_param_c
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.db.create_views import (
create_missing_views,
create_view_tolerating_race,
should_create_missing_views,
)
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
@ -3273,7 +3274,10 @@ class PrismaClient:
## check if required view exists ##
if ret[0]["view_names"] and required_view not in ret[0]["view_names"]:
await self.health_check() # make sure we can connect to db
await self.db.execute_raw("""
await create_view_tolerating_race(
self.db,
"LiteLLM_VerificationTokenView",
"""
CREATE VIEW "LiteLLM_VerificationTokenView" AS
SELECT
v.*,
@ -3283,9 +3287,8 @@ class PrismaClient:
t.rpm_limit AS team_rpm_limit
FROM "LiteLLM_VerificationToken" v
LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id = t.team_id;
""")
verbose_proxy_logger.info("LiteLLM_VerificationTokenView Created in DB!")
""",
)
else:
should_create_views: Final = await should_create_missing_views(db=self.db)
if should_create_views:

63
tests/test_litellm/proxy/db/test_create_views.py
View file

@ -189,3 +189,66 @@ async def test_create_views_creates_view_on_undefined_table_error():
await create_missing_views(mock_db)
mock_db.execute_raw.assert_called_once()
# Every view create_missing_views is responsible for. Hard-coded rather than
# derived from the module, so adding a view without guarding it fails here.
EXPECTED_VIEW_COUNT = 8
@pytest.mark.asyncio
async def test_create_views_tolerates_a_concurrent_creator_on_every_view():
"""A replica that loses the CREATE race must attempt every view regardless.
Regression: two proxy pods booting on a fresh DB both see every view as
absent and both issue the CREATE, and Postgres fails the loser with a
duplicate-object error on whichever views the winner got to first. Any
creation site still calling execute_raw unguarded re-raises that error and
aborts the rest of the function.
Every CREATE loses here, which is what pins the guard to all of them: an
earlier version of this fix converted only the first and the last site and
still died on MonthlyGlobalSpend against a real Postgres. Counting the
attempts is the assertion, because a partial fix simply stops early.
"""
from litellm.proxy.db.create_views import create_missing_views
mock_db = MagicMock()
mock_db.query_raw = AsyncMock(side_effect=Exception("relation does not exist"))
mock_db.execute_raw = AsyncMock(
side_effect=Exception('relation "some_view" already exists')
)
await create_missing_views(mock_db)
assert mock_db.execute_raw.await_count == EXPECTED_VIEW_COUNT, (
f"every view must still be attempted when the replica loses every race; "
f"got {mock_db.execute_raw.await_count} of {EXPECTED_VIEW_COUNT}, so a "
f"creation site is still unguarded and aborted the rest"
)
@pytest.mark.asyncio
async def test_create_views_reraises_genuine_ddl_error():
"""An already-exists guard must not swallow real DDL failures."""
from litellm.proxy.db.create_views import create_missing_views
mock_db = MagicMock()
mock_db.query_raw = AsyncMock(side_effect=Exception("relation does not exist"))
mock_db.execute_raw = AsyncMock(side_effect=Exception("syntax error at or near"))
with pytest.raises(Exception, match="syntax error"):
await create_missing_views(mock_db)
@pytest.mark.asyncio
async def test_create_view_tolerating_race_swallows_only_already_exists():
from litellm.proxy.db.create_views import create_view_tolerating_race
mock_db = MagicMock()
mock_db.execute_raw = AsyncMock(side_effect=Exception("duplicate object"))
await create_view_tolerating_race(mock_db, "SomeView", "CREATE VIEW ...")
mock_db.execute_raw = AsyncMock(side_effect=Exception("permission denied"))
with pytest.raises(Exception, match="permission denied"):
await create_view_tolerating_race(mock_db, "SomeView", "CREATE VIEW ...")