mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-11 03:38:38 +00:00
fix(usage): match legacy NULL model_group through COALESCE instead of a backfill
The daily spend identity indexes and the writer's ON CONFLICT target use
COALESCE("model_group", ''), so rows stored with a NULL group before
model_group joined the key share one identity with the '' new writers store
and no historical row is rewritten. Prisma cannot declare an expression index,
so the three schema copies drop the @@unique and --use_prisma_db_push runs the
identity migration right after the push.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
Tin Chi Lo
parent
c6669babf4
commit
29607f8794
9 changed files with 85 additions and 29 deletions
|
|
@ -1,16 +1,16 @@
|
|||
BEGIN;
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyUserSpend_group_key" ON "LiteLLM_DailyUserSpend" ("user_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", "model_group");
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyUserSpend_group_key" ON "LiteLLM_DailyUserSpend" ("user_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", COALESCE("model_group", ''));
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyTeamSpend_group_key" ON "LiteLLM_DailyTeamSpend" ("team_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", "model_group");
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyTeamSpend_group_key" ON "LiteLLM_DailyTeamSpend" ("team_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", COALESCE("model_group", ''));
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyOrganizationSpend_group_key" ON "LiteLLM_DailyOrganizationSpend" ("organization_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", "model_group");
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyOrganizationSpend_group_key" ON "LiteLLM_DailyOrganizationSpend" ("organization_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", COALESCE("model_group", ''));
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyEndUserSpend_group_key" ON "LiteLLM_DailyEndUserSpend" ("end_user_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", "model_group");
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyEndUserSpend_group_key" ON "LiteLLM_DailyEndUserSpend" ("end_user_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", COALESCE("model_group", ''));
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyAgentSpend_group_key" ON "LiteLLM_DailyAgentSpend" ("agent_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", "model_group");
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyAgentSpend_group_key" ON "LiteLLM_DailyAgentSpend" ("agent_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", COALESCE("model_group", ''));
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyTagSpend_group_key" ON "LiteLLM_DailyTagSpend" ("tag", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", "model_group");
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_DailyTagSpend_group_key" ON "LiteLLM_DailyTagSpend" ("tag", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", COALESCE("model_group", ''));
|
||||
|
||||
DROP INDEX IF EXISTS "LiteLLM_DailyUserSpend_user_id_date_api_key_model_custom_ll_key";
|
||||
|
||||
|
|
|
|||
|
|
@ -873,7 +873,6 @@ model LiteLLM_DailyUserSpend {
|
|||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@unique([user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyUserSpend_group_key")
|
||||
@@index([date])
|
||||
@@index([user_id, date])
|
||||
@@index([api_key])
|
||||
|
|
@ -942,7 +941,6 @@ model LiteLLM_DailyOrganizationSpend {
|
|||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@unique([organization_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyOrganizationSpend_group_key")
|
||||
@@index([date])
|
||||
@@index([organization_id, date])
|
||||
@@index([api_key])
|
||||
|
|
@ -979,7 +977,6 @@ model LiteLLM_DailyEndUserSpend {
|
|||
timed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
@@unique([end_user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyEndUserSpend_group_key")
|
||||
@@index([date])
|
||||
@@index([end_user_id, date])
|
||||
@@index([api_key])
|
||||
|
|
@ -1016,7 +1013,6 @@ model LiteLLM_DailyAgentSpend {
|
|||
timed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
@@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyAgentSpend_group_key")
|
||||
@@index([date])
|
||||
@@index([agent_id, date])
|
||||
@@index([api_key])
|
||||
|
|
@ -1055,7 +1051,6 @@ model LiteLLM_DailyTeamSpend {
|
|||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@unique([team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyTeamSpend_group_key")
|
||||
@@index([date])
|
||||
@@index([team_id, date])
|
||||
@@index([api_key])
|
||||
|
|
@ -1094,7 +1089,6 @@ model LiteLLM_DailyTagSpend {
|
|||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@unique([tag, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyTagSpend_group_key")
|
||||
@@index([date])
|
||||
@@index([tag, date])
|
||||
@@index([api_key])
|
||||
|
|
|
|||
|
|
@ -202,6 +202,9 @@ def _max_migration_timestamp(names) -> int:
|
|||
return max(_migration_timestamp(n) for n in names)
|
||||
|
||||
|
||||
DAILY_SPEND_IDENTITY_MIGRATION: Final = "20260922000000_daily_spend_model_group_identity"
|
||||
|
||||
|
||||
def _get_prisma_command() -> str:
|
||||
"""Get the Prisma command to use, bypassing Python wrapper in offline mode."""
|
||||
if str_to_bool(os.getenv("PRISMA_OFFLINE_MODE")):
|
||||
|
|
@ -1027,6 +1030,7 @@ class ProxyExtrasDBManager:
|
|||
stdout=None,
|
||||
stderr=None,
|
||||
)
|
||||
ProxyExtrasDBManager._apply_daily_spend_identity(migrations_dir)
|
||||
return True
|
||||
except (
|
||||
subprocess.CalledProcessError,
|
||||
|
|
@ -1351,6 +1355,24 @@ class ProxyExtrasDBManager:
|
|||
thread.start()
|
||||
return thread
|
||||
|
||||
@staticmethod
|
||||
def _apply_daily_spend_identity(prisma_dir: str) -> None:
|
||||
"""`prisma db push` builds only what schema.prisma declares. The daily spend identity is a
|
||||
COALESCE expression index Prisma cannot declare, so the migration owning it runs after the push."""
|
||||
prisma_toolchain.run_prisma(
|
||||
[
|
||||
_get_prisma_command(),
|
||||
"db",
|
||||
"execute",
|
||||
"--file",
|
||||
f"{prisma_dir}/migrations/{DAILY_SPEND_IDENTITY_MIGRATION}/migration.sql",
|
||||
"--schema",
|
||||
f"{prisma_dir}/schema.prisma",
|
||||
],
|
||||
timeout=prisma_command_timeout(),
|
||||
env=_get_prisma_env(),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _run_migrations(use_migrate: bool, use_v2_resolver: bool) -> bool:
|
||||
if use_v2_resolver:
|
||||
|
|
@ -1584,6 +1606,7 @@ class ProxyExtrasDBManager:
|
|||
stderr=None,
|
||||
env=_get_prisma_env(),
|
||||
)
|
||||
ProxyExtrasDBManager._apply_daily_spend_identity(migrations_dir)
|
||||
return True
|
||||
except subprocess.TimeoutExpired:
|
||||
logger.warning(
|
||||
|
|
|
|||
|
|
@ -98,7 +98,7 @@ When the guard fires:
|
|||
|
||||
`20260922000000_daily_spend_model_group_identity` separates daily usage by the requested public model group. Schedule a maintenance window: the six replacement unique indexes scan existing daily spend tables and block writes while building
|
||||
|
||||
Drain incoming requests and all in-memory and shared Redis spend buffers, then stop every old proxy and spend collector before applying this migration. Start only upgraded writers after it completes. Old writers require the removed conflict target and cannot run alongside the expanded identity. The migration changes indexes only; previously merged attribution and legacy NULL groups remain as recorded
|
||||
Drain incoming requests and all in-memory and shared Redis spend buffers, then stop every old proxy and spend collector before applying this migration. Start only upgraded writers after it completes. Old writers require the removed conflict target and cannot run alongside the expanded identity. The migration changes indexes only and previously merged attribution stays as recorded. The new indexes and the writer's conflict target match `model_group` through `COALESCE("model_group", '')`, so rows stored with a NULL group before this change share one identity with the `''` new writers store and no row is rewritten. Prisma cannot declare an expression index, so `--use_prisma_db_push` deployments run this migration right after the push
|
||||
|
||||
PTU flat charges retain their deployment/day identity across public-model renames. New inference rows include the requested group in their identity. Rolling back requires stopping writers and reconciling group-separated rows before restoring the old unique indexes; simply starting an older image is unsupported
|
||||
|
||||
|
|
|
|||
|
|
@ -120,6 +120,17 @@ def _quoted(columns: Sequence[str]) -> str:
|
|||
return ", ".join(f'"{column}"' for column in columns)
|
||||
|
||||
|
||||
_CONFLICT_EXPRESSIONS: Final[Mapping[str, str]] = MappingProxyType({"model_group": "COALESCE(\"model_group\", '')"})
|
||||
|
||||
|
||||
def _conflict_target(table: DailySpendTable) -> str:
|
||||
"""The unique index's expression list. model_group joined the key after rows were stored with
|
||||
NULL in it, so the index matches it through COALESCE and those rows keep their identity unrewritten."""
|
||||
return ", ".join(
|
||||
_CONFLICT_EXPRESSIONS.get(column, f'"{column}"') for column in (table.entity_id_column, *_KEY_COLUMNS)
|
||||
)
|
||||
|
||||
|
||||
def _as_text(value: object) -> str:
|
||||
return "" if value is None else str(value)
|
||||
|
||||
|
|
@ -219,7 +230,7 @@ def build_bulk_upsert(
|
|||
sql: Final = (
|
||||
f'INSERT INTO {quoted_table} ({_quoted(columns)}, "updated_at")\n'
|
||||
f"VALUES {rows}\n"
|
||||
f"ON CONFLICT ({_quoted((table.entity_id_column, *_KEY_COLUMNS))}) DO UPDATE SET\n"
|
||||
f"ON CONFLICT ({_conflict_target(table)}) DO UPDATE SET\n"
|
||||
f" {increments}{request_id_update},\n"
|
||||
f" \"updated_at\" = (NOW() AT TIME ZONE 'UTC')"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -873,7 +873,6 @@ model LiteLLM_DailyUserSpend {
|
|||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@unique([user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyUserSpend_group_key")
|
||||
@@index([date])
|
||||
@@index([user_id, date])
|
||||
@@index([api_key])
|
||||
|
|
@ -942,7 +941,6 @@ model LiteLLM_DailyOrganizationSpend {
|
|||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@unique([organization_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyOrganizationSpend_group_key")
|
||||
@@index([date])
|
||||
@@index([organization_id, date])
|
||||
@@index([api_key])
|
||||
|
|
@ -979,7 +977,6 @@ model LiteLLM_DailyEndUserSpend {
|
|||
timed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
@@unique([end_user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyEndUserSpend_group_key")
|
||||
@@index([date])
|
||||
@@index([end_user_id, date])
|
||||
@@index([api_key])
|
||||
|
|
@ -1016,7 +1013,6 @@ model LiteLLM_DailyAgentSpend {
|
|||
timed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
@@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyAgentSpend_group_key")
|
||||
@@index([date])
|
||||
@@index([agent_id, date])
|
||||
@@index([api_key])
|
||||
|
|
@ -1055,7 +1051,6 @@ model LiteLLM_DailyTeamSpend {
|
|||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@unique([team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyTeamSpend_group_key")
|
||||
@@index([date])
|
||||
@@index([team_id, date])
|
||||
@@index([api_key])
|
||||
|
|
@ -1094,7 +1089,6 @@ model LiteLLM_DailyTagSpend {
|
|||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@unique([tag, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyTagSpend_group_key")
|
||||
@@index([date])
|
||||
@@index([tag, date])
|
||||
@@index([api_key])
|
||||
|
|
|
|||
|
|
@ -873,7 +873,6 @@ model LiteLLM_DailyUserSpend {
|
|||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@unique([user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyUserSpend_group_key")
|
||||
@@index([date])
|
||||
@@index([user_id, date])
|
||||
@@index([api_key])
|
||||
|
|
@ -942,7 +941,6 @@ model LiteLLM_DailyOrganizationSpend {
|
|||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@unique([organization_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyOrganizationSpend_group_key")
|
||||
@@index([date])
|
||||
@@index([organization_id, date])
|
||||
@@index([api_key])
|
||||
|
|
@ -979,7 +977,6 @@ model LiteLLM_DailyEndUserSpend {
|
|||
timed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
@@unique([end_user_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyEndUserSpend_group_key")
|
||||
@@index([date])
|
||||
@@index([end_user_id, date])
|
||||
@@index([api_key])
|
||||
|
|
@ -1016,7 +1013,6 @@ model LiteLLM_DailyAgentSpend {
|
|||
timed_requests BigInt @default(0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
@@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyAgentSpend_group_key")
|
||||
@@index([date])
|
||||
@@index([agent_id, date])
|
||||
@@index([api_key])
|
||||
|
|
@ -1055,7 +1051,6 @@ model LiteLLM_DailyTeamSpend {
|
|||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@unique([team_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyTeamSpend_group_key")
|
||||
@@index([date])
|
||||
@@index([team_id, date])
|
||||
@@index([api_key])
|
||||
|
|
@ -1094,7 +1089,6 @@ model LiteLLM_DailyTagSpend {
|
|||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
@@unique([tag, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name, endpoint, model_group], map: "LiteLLM_DailyTagSpend_group_key")
|
||||
@@index([date])
|
||||
@@index([tag, date])
|
||||
@@index([api_key])
|
||||
|
|
|
|||
|
|
@ -603,6 +603,45 @@ class TestResolveAllMigrationsLedger:
|
|||
assert self._executed_sql(calls) == _PARTITIONED_DRIFT_SQL
|
||||
|
||||
|
||||
class TestDbPushAppliesDailySpendIdentity:
|
||||
def _push(self, monkeypatch, tmp_path, run):
|
||||
import litellm_proxy_extras.utils as utils_module
|
||||
|
||||
calls = []
|
||||
|
||||
def fake_run(cmd, **kwargs):
|
||||
calls.append(cmd)
|
||||
return _FakeCompleted()
|
||||
|
||||
monkeypatch.setenv("LITELLM_MIGRATION_DIR", str(tmp_path))
|
||||
monkeypatch.setattr(utils_module.prisma_toolchain, "run_prisma", fake_run)
|
||||
monkeypatch.setattr(utils_module, "_get_prisma_env", lambda: {})
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "spend_logs_is_partitioned", staticmethod(lambda: False))
|
||||
monkeypatch.setattr(ProxyExtrasDBManager, "raise_if_lens_rename_pending", staticmethod(lambda: None))
|
||||
assert run() is True
|
||||
return calls
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"run",
|
||||
[
|
||||
lambda: ProxyExtrasDBManager._run_migrations(use_migrate=False, use_v2_resolver=False),
|
||||
lambda: ProxyExtrasDBManager._setup_database_v2(use_migrate=False),
|
||||
],
|
||||
ids=["v1", "v2"],
|
||||
)
|
||||
def test_the_identity_migration_runs_right_after_the_push(self, monkeypatch, tmp_path, run):
|
||||
import litellm_proxy_extras.utils as utils_module
|
||||
|
||||
calls = self._push(monkeypatch, tmp_path, run)
|
||||
|
||||
assert [cmd[1:3] for cmd in calls] == [["db", "push"], ["db", "execute"]]
|
||||
executed = calls[1]
|
||||
assert executed[executed.index("--file") + 1] == (
|
||||
f"{tmp_path}/migrations/{utils_module.DAILY_SPEND_IDENTITY_MIGRATION}/migration.sql"
|
||||
)
|
||||
assert executed[executed.index("--schema") + 1] == f"{tmp_path}/schema.prisma"
|
||||
|
||||
|
||||
class TestPartitionedSpendLogsPushGuard:
|
||||
def _forbid_subprocess(self, monkeypatch):
|
||||
import litellm_proxy_extras.utils as utils_module
|
||||
|
|
|
|||
|
|
@ -99,13 +99,14 @@ def test_one_statement_carries_every_row_in_the_batch():
|
|||
assert sql.count("(NOW() AT TIME ZONE 'UTC')") == 100 + 1
|
||||
|
||||
|
||||
def test_conflict_target_is_the_full_unique_constraint():
|
||||
def test_conflict_target_is_the_unique_index_with_model_group_matched_through_coalesce():
|
||||
sql, _ = build_bulk_upsert(TAG_TABLE, merge_by_conflict_key(TAG_TABLE, (tag_txn(),)))
|
||||
|
||||
conflict_target = re.search(r"ON CONFLICT \(([^)]*)\)", sql)
|
||||
conflict_target = re.search(r"ON CONFLICT \((.*)\) DO UPDATE", sql)
|
||||
assert conflict_target is not None
|
||||
assert conflict_target.group(1) == (
|
||||
'"tag", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", "model_group"'
|
||||
'"tag", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name", "endpoint", '
|
||||
"COALESCE(\"model_group\", '')"
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue