mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
fix(accounting): stop delivery metrics from scanning retained receipts on every flush (#50565)
* fix(accounting): stop delivery metrics from scanning retained receipts on every flush The backlog query now reads only pending and parked deliveries, so its cost follows the backlog instead of retained receipt history. Batch count and table size are measured once per check interval per pod, and the storage and delayed-delivery alerts repeat at most once per interval instead of every flush Fixes #50475 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test(accounting): prove the storage query reruns only when its interval is due Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * fix(accounting): read the check interval on each flush so YAML environment_variables apply The writer is built at proxy import, before the config's environment_variables load, so reading the interval there ignored values set in config.yaml Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * test(router): wait for router accounting instead of sleeping in user-config limit tests Each request ended with a 0.3s sleep and hoped the router's usage writes had landed, so under CI load a concurrent request could read counters mid-update. The tests now flush the logging worker and wait for every task the requests scheduled, and a retired router's counter flush lands before the next router reads Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
joshua-berri
GitHub
parent
7e6479b3a0
commit
82e43cd55a
8 changed files with 240 additions and 56 deletions
|
|
@ -1,7 +1,9 @@
|
|||
import time
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm.repositories.accounting_delivery_repository import DeliveryMetrics
|
||||
from litellm.repositories.accounting_delivery_repository import DeliveryBacklog, DeliveryStorage
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prometheus_client import Gauge, Histogram
|
||||
|
|
@ -44,12 +46,32 @@ def _instruments() -> _Instruments | None:
|
|||
_INSTRUMENTS: Final = _instruments()
|
||||
|
||||
|
||||
def record_deliveries(metrics: DeliveryMetrics, seconds: float) -> None:
|
||||
def record_backlog(backlog: DeliveryBacklog, seconds: float) -> None:
|
||||
if _INSTRUMENTS is None:
|
||||
return
|
||||
_INSTRUMENTS.pending.set(metrics.pending_count)
|
||||
_INSTRUMENTS.parked.set(metrics.parked_count)
|
||||
_INSTRUMENTS.oldest.set(max(0.0, metrics.oldest_pending_seconds))
|
||||
_INSTRUMENTS.pending.set(backlog.pending_count)
|
||||
_INSTRUMENTS.parked.set(backlog.parked_count)
|
||||
_INSTRUMENTS.oldest.set(max(0.0, backlog.oldest_pending_seconds))
|
||||
_INSTRUMENTS.batch_time.observe(seconds)
|
||||
_INSTRUMENTS.batch_count.set(metrics.batch_count)
|
||||
_INSTRUMENTS.table_bytes.set(metrics.table_bytes)
|
||||
|
||||
|
||||
def record_storage(storage: DeliveryStorage) -> None:
|
||||
if _INSTRUMENTS is None:
|
||||
return
|
||||
_INSTRUMENTS.batch_count.set(storage.batch_count)
|
||||
_INSTRUMENTS.table_bytes.set(storage.table_bytes)
|
||||
|
||||
|
||||
class Cadence:
|
||||
__slots__ = ("_clock", "_last")
|
||||
|
||||
def __init__(self, clock: Callable[[], float] = time.monotonic) -> None:
|
||||
self._clock = clock
|
||||
self._last: float | None = None
|
||||
|
||||
def due(self, interval_seconds: float) -> bool:
|
||||
now: Final = self._clock()
|
||||
if self._last is not None and now - self._last < interval_seconds:
|
||||
return False
|
||||
self._last = now
|
||||
return True
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ from typing import (
|
|||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.accounting.application import apply_deliveries
|
||||
from litellm.accounting.metrics import record_deliveries
|
||||
from litellm.accounting.metrics import record_backlog, record_storage
|
||||
from litellm.accounting.projection import Contribution
|
||||
from litellm.repositories.accounting_delivery_repository import (
|
||||
AccountingDeliveryRepository,
|
||||
|
|
@ -96,24 +96,37 @@ async def flush_accounting_deliveries(
|
|||
or time.monotonic() >= cleanup_deadline
|
||||
):
|
||||
break
|
||||
metrics: Final = await repository.metrics()
|
||||
record_deliveries(metrics, time.monotonic() - started)
|
||||
if metrics.batch_count >= int(
|
||||
os.environ.get("LITELLM_ACCOUNTING_BATCH_DELIVERY_ALERT_COUNT", "1000000")
|
||||
) or metrics.table_bytes >= int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_ALERT_BYTES", "1073741824")):
|
||||
await _alert(
|
||||
proxy_logging_obj,
|
||||
message=f"Accounting delivery storage requires review. Retained batches={metrics.batch_count}, table bytes={metrics.table_bytes}.",
|
||||
)
|
||||
if metrics.oldest_pending_seconds > 300:
|
||||
await _alert(
|
||||
proxy_logging_obj,
|
||||
message=f"Accounting deliveries are delayed. Pending={metrics.pending_count}, oldest={metrics.oldest_pending_seconds:.0f}s.",
|
||||
)
|
||||
await _report(writer, repository, proxy_logging_obj, time.monotonic() - started)
|
||||
finally:
|
||||
if acquired:
|
||||
await writer.pod_lock_manager.release_lock(cronjob_id=_JOB)
|
||||
|
||||
|
||||
async def _report(
|
||||
writer: "DBSpendUpdateWriter",
|
||||
repository: AccountingDeliveryRepository,
|
||||
proxy_logging_obj: _AccountingAlerts,
|
||||
seconds: float,
|
||||
) -> None:
|
||||
backlog: Final = await repository.backlog()
|
||||
record_backlog(backlog, seconds)
|
||||
interval: Final = float(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_CHECK_INTERVAL_SECONDS", "900"))
|
||||
if writer.accounting_storage_checks.due(interval):
|
||||
storage: Final = await repository.storage()
|
||||
record_storage(storage)
|
||||
if storage.batch_count >= int(
|
||||
os.environ.get("LITELLM_ACCOUNTING_BATCH_DELIVERY_ALERT_COUNT", "1000000")
|
||||
) or storage.table_bytes >= int(os.environ.get("LITELLM_ACCOUNTING_DELIVERY_ALERT_BYTES", "1073741824")):
|
||||
await _alert(
|
||||
proxy_logging_obj,
|
||||
message=f"Accounting delivery storage requires review. Retained batches={storage.batch_count}, table bytes={storage.table_bytes}.",
|
||||
)
|
||||
if backlog.oldest_pending_seconds > 300 and writer.accounting_delay_alerts.due(interval):
|
||||
await _alert(
|
||||
proxy_logging_obj,
|
||||
message=f"Accounting deliveries are delayed. Pending={backlog.pending_count}, oldest={backlog.oldest_pending_seconds:.0f}s.",
|
||||
)
|
||||
|
||||
|
||||
async def _alert(alerts: _AccountingAlerts, message: str) -> None:
|
||||
await alerts.alerting_handler(message=message, level="High", alert_type=AlertType.failed_tracking_spend)
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ from typing_extensions import LiteralString, ReadOnly, TypedDict
|
|||
import litellm
|
||||
from litellm._internal_context import service_target, with_service_target
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.accounting.metrics import Cadence
|
||||
from litellm.caching import RedisCache
|
||||
from litellm.constants import (
|
||||
BILLED_RETRY_ATTEMPT_METADATA_KEY,
|
||||
|
|
@ -682,8 +683,11 @@ class DBSpendUpdateWriter:
|
|||
def __init__(
|
||||
self,
|
||||
redis_cache: RedisCache | None = None,
|
||||
accounting_clock: Callable[[], float] = time.monotonic,
|
||||
):
|
||||
self.redis_cache = redis_cache
|
||||
self.accounting_storage_checks = Cadence(accounting_clock)
|
||||
self.accounting_delay_alerts = Cadence(accounting_clock)
|
||||
self.redis_update_buffer = RedisUpdateBuffer(redis_cache=self.redis_cache)
|
||||
self.pod_lock_manager = PodLockManager()
|
||||
self.spend_update_queue = SpendUpdateQueue()
|
||||
|
|
|
|||
|
|
@ -60,12 +60,17 @@ class DeliveryBatchResult:
|
|||
error: str | None = None
|
||||
|
||||
|
||||
class DeliveryMetrics(BaseModel):
|
||||
class DeliveryBacklog(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
pending_count: int
|
||||
parked_count: int
|
||||
oldest_pending_seconds: float
|
||||
|
||||
|
||||
class DeliveryStorage(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
batch_count: int
|
||||
table_bytes: int
|
||||
|
||||
|
|
@ -88,6 +93,14 @@ SET attempts = attempts + 1, last_error = $3,
|
|||
state = $4::"AccountingDeliveryState"
|
||||
WHERE delivery_id = $1 AND state = 'pending' AND attempts = $2::integer
|
||||
RETURNING delivery_id, attempts"""
|
||||
_BACKLOG: Final = """SELECT COUNT(*) FILTER (WHERE state = 'pending') AS pending_count,
|
||||
COUNT(*) FILTER (WHERE state = 'parked') AS parked_count,
|
||||
COALESCE(EXTRACT(EPOCH FROM CURRENT_TIMESTAMP - MIN(created_at) FILTER (WHERE state = 'pending')), 0)
|
||||
AS oldest_pending_seconds
|
||||
FROM "LiteLLM_SpendLogDelivery"
|
||||
WHERE state IN ('pending', 'parked')"""
|
||||
_STORAGE: Final = """SELECT (SELECT COUNT(*) FROM "LiteLLM_SpendLogDelivery" WHERE is_batch) AS batch_count,
|
||||
pg_total_relation_size('"LiteLLM_SpendLogDelivery"') AS table_bytes"""
|
||||
_UNDECODABLE: Final = """UPDATE "LiteLLM_SpendLogDelivery"
|
||||
SET state = 'parked', attempts = attempts + 1, last_error = 'UndecodablePayload'
|
||||
WHERE delivery_id = ANY($1::text[]) AND state = 'pending'"""
|
||||
|
|
@ -114,17 +127,13 @@ class AccountingDeliveryRepository:
|
|||
def __init__(self, database: LedgerDatabase) -> None:
|
||||
self._database = database
|
||||
|
||||
async def metrics(self) -> DeliveryMetrics:
|
||||
async def backlog(self) -> DeliveryBacklog:
|
||||
async with self._database.tx(timeout=_TIMEOUT) as transaction:
|
||||
rows: Final = await transaction.query_raw(
|
||||
"SELECT COUNT(*) FILTER (WHERE state = 'pending') AS pending_count, "
|
||||
"COUNT(*) FILTER (WHERE state = 'parked') AS parked_count, "
|
||||
"COALESCE(EXTRACT(EPOCH FROM CURRENT_TIMESTAMP - MIN(created_at) FILTER (WHERE state = 'pending')), 0) "
|
||||
"AS oldest_pending_seconds, COUNT(*) FILTER (WHERE is_batch) AS batch_count, "
|
||||
"pg_total_relation_size('\"LiteLLM_SpendLogDelivery\"') AS table_bytes "
|
||||
'FROM "LiteLLM_SpendLogDelivery"'
|
||||
)
|
||||
return DeliveryMetrics.model_validate(rows[0])
|
||||
return DeliveryBacklog.model_validate((await transaction.query_raw(_BACKLOG))[0])
|
||||
|
||||
async def storage(self) -> DeliveryStorage:
|
||||
async with self._database.tx(timeout=_TIMEOUT) as transaction:
|
||||
return DeliveryStorage.model_validate((await transaction.query_raw(_STORAGE))[0])
|
||||
|
||||
async def cleanup_applied(self, *, retention_days: int = 1, batch_size: int = 1000) -> int:
|
||||
if retention_days < 1 or batch_size < 1:
|
||||
|
|
|
|||
|
|
@ -373,15 +373,45 @@ async def test_receipt_cleanup_keeps_pending_and_parked_rows_and_reports_lag(
|
|||
)
|
||||
repository = AccountingDeliveryRepository(delivery_client)
|
||||
assert await repository.cleanup_applied() == 1
|
||||
metrics = await repository.metrics()
|
||||
assert metrics.pending_count == 1
|
||||
assert metrics.parked_count == 1
|
||||
assert metrics.oldest_pending_seconds >= 40 * 86400
|
||||
backlog = await repository.backlog()
|
||||
assert backlog.pending_count == 1
|
||||
assert backlog.parked_count == 1
|
||||
assert backlog.oldest_pending_seconds >= 40 * 86400
|
||||
assert delivery_database.execute(
|
||||
'SELECT request_id FROM "LiteLLM_SpendLogDelivery" ORDER BY request_id'
|
||||
).fetchall() == [("parked",), ("pending",)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_backlog_reads_only_unapplied_deliveries_however_many_receipts_are_retained(
|
||||
delivery_database: psycopg.Connection[tuple[object, ...]], delivery_client: Prisma
|
||||
) -> None:
|
||||
from litellm.repositories.accounting_delivery_repository import _BACKLOG
|
||||
|
||||
delivery_database.execute(
|
||||
'INSERT INTO "LiteLLM_SpendLogDelivery" (delivery_id, request_id, accounting_payload, is_batch, state, applied_at) '
|
||||
"SELECT 'applied-' || n, 'applied-' || n, jsonb_build_object('payload', repeat(md5(n::text), 32)), "
|
||||
"n % 10 = 0, 'applied', CURRENT_TIMESTAMP FROM generate_series(1, 20000) AS n"
|
||||
)
|
||||
_insert(delivery_database, tuple(_row(state) for state in ("pending", "parked")))
|
||||
delivery_database.execute(
|
||||
'UPDATE "LiteLLM_SpendLogDelivery" SET state = request_id::"AccountingDeliveryState" '
|
||||
"WHERE request_id IN ('pending', 'parked')"
|
||||
)
|
||||
delivery_database.execute('ANALYZE "LiteLLM_SpendLogDelivery"')
|
||||
explained = delivery_database.execute(f"EXPLAIN (ANALYZE, BUFFERS, FORMAT JSON) {_BACKLOG}").fetchone()
|
||||
assert explained is not None
|
||||
plan = explained[0][0]["Plan"]
|
||||
table_pages = delivery_database.execute(
|
||||
"SELECT relpages FROM pg_class WHERE oid = '\"LiteLLM_SpendLogDelivery\"'::regclass"
|
||||
).fetchone()
|
||||
assert table_pages is not None and table_pages[0] > 1000
|
||||
assert plan["Shared Hit Blocks"] + plan["Shared Read Blocks"] < table_pages[0] / 20
|
||||
|
||||
backlog = await AccountingDeliveryRepository(delivery_client).backlog()
|
||||
assert (backlog.pending_count, backlog.parked_count) == (1, 1)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("delivery_database", ["public", "ledger_tenant"], indirect=True)
|
||||
@pytest.mark.parametrize("state", [None, "pending", "parked", "applied"])
|
||||
def test_delivery_migration_replay_preserves_ledger_receipts(delivery_database, state):
|
||||
|
|
@ -467,9 +497,9 @@ async def test_retention_keeps_batch_receipts_indefinitely(delivery_database, de
|
|||
assert await repository.cleanup_applied() == 1
|
||||
assert await repository.cleanup_applied() == 0
|
||||
assert delivery_database.execute('SELECT request_id FROM "LiteLLM_SpendLogDelivery"').fetchall() == [("batch",)]
|
||||
metrics = await repository.metrics()
|
||||
assert metrics.batch_count == 1
|
||||
assert metrics.table_bytes > 0
|
||||
storage = await repository.storage()
|
||||
assert storage.batch_count == 1
|
||||
assert storage.table_bytes > 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("disabled", [False, True])
|
||||
|
|
|
|||
|
|
@ -127,6 +127,66 @@ async def test_worker_removes_more_expired_receipts_than_one_cleanup_batch_in_on
|
|||
await client.db.disconnect()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_worker_checks_storage_and_repeats_delay_alerts_once_per_interval(
|
||||
parity_database: str, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
from litellm.accounting.worker import flush_accounting_deliveries
|
||||
|
||||
class _Alerts:
|
||||
def __init__(self) -> None:
|
||||
self.messages: list[str] = [] # mutable-ok: alert sink the worker appends to
|
||||
|
||||
async def alerting_handler(self, message: str, level: str, alert_type: object) -> None:
|
||||
self.messages.append(message)
|
||||
|
||||
class _Clock:
|
||||
def __init__(self) -> None:
|
||||
self.now = 0.0
|
||||
|
||||
def __call__(self) -> float:
|
||||
return self.now
|
||||
|
||||
monkeypatch.setenv("DATABASE_URL", parity_database)
|
||||
monkeypatch.setenv("LITELLM_ACCOUNTING_DELIVERY_ALERT_BYTES", "1")
|
||||
with psycopg.connect(parity_database, autocommit=True) as connection:
|
||||
connection.execute(
|
||||
'INSERT INTO "LiteLLM_SpendLogDelivery" '
|
||||
"(delivery_id, request_id, accounting_payload, is_batch, attempts, created_at, next_attempt_at) "
|
||||
"VALUES ('stuck', 'stuck', '{}'::jsonb, false, 1, CURRENT_TIMESTAMP - INTERVAL '1 hour', "
|
||||
"CURRENT_TIMESTAMP + INTERVAL '1 hour')"
|
||||
)
|
||||
client: Final = PrismaClient(database_url=parity_database, proxy_logging_obj=proxy_server.proxy_logging_obj)
|
||||
clock: Final = _Clock()
|
||||
writer: Final = DBSpendUpdateWriter(accounting_clock=clock)
|
||||
alerts: Final = _Alerts()
|
||||
monkeypatch.setenv("LITELLM_ACCOUNTING_DELIVERY_CHECK_INTERVAL_SECONDS", "100")
|
||||
|
||||
async def flush_at(moment: float) -> tuple[str, ...]:
|
||||
clock.now = moment
|
||||
before: Final = len(alerts.messages)
|
||||
await flush_accounting_deliveries(writer, client, alerts)
|
||||
return tuple(message.split(",")[0] for message in alerts.messages[before:])
|
||||
|
||||
await client.db.connect()
|
||||
try:
|
||||
first: Final = await flush_at(0.0)
|
||||
with psycopg.connect(parity_database, autocommit=True) as connection:
|
||||
connection.execute(
|
||||
'INSERT INTO "LiteLLM_SpendLogDelivery" (delivery_id, request_id, accounting_payload, is_batch, state) '
|
||||
"VALUES ('batch', 'batch', '{}'::jsonb, true, 'applied')"
|
||||
)
|
||||
within: Final = (await flush_at(10.0), await flush_at(99.0))
|
||||
after: Final = await flush_at(100.0)
|
||||
finally:
|
||||
await client.db.disconnect()
|
||||
|
||||
delayed: Final = "Accounting deliveries are delayed. Pending=1"
|
||||
assert first == ("Accounting delivery storage requires review. Retained batches=0", delayed)
|
||||
assert within == ((), ())
|
||||
assert after == ("Accounting delivery storage requires review. Retained batches=1", delayed)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ledger", [False, True])
|
||||
@pytest.mark.parametrize("exclude_messages", [False, True])
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -5,12 +5,11 @@ import pytest
|
|||
from prometheus_client.metrics import MetricWrapperBase
|
||||
|
||||
from litellm.accounting import metrics
|
||||
from litellm.accounting.metrics import record_deliveries
|
||||
from litellm.repositories.accounting_delivery_repository import DeliveryMetrics
|
||||
from litellm.accounting.metrics import Cadence, record_backlog, record_storage
|
||||
from litellm.repositories.accounting_delivery_repository import DeliveryBacklog, DeliveryStorage
|
||||
|
||||
_DELIVERY_METRICS = DeliveryMetrics(
|
||||
pending_count=4, parked_count=1, oldest_pending_seconds=-3.0, batch_count=7, table_bytes=8192
|
||||
)
|
||||
_BACKLOG = DeliveryBacklog(pending_count=4, parked_count=1, oldest_pending_seconds=-3.0)
|
||||
_STORAGE = DeliveryStorage(batch_count=7, table_bytes=8192)
|
||||
|
||||
|
||||
def _sample(metric: MetricWrapperBase, suffix: str = "") -> float:
|
||||
|
|
@ -23,7 +22,8 @@ def test_delivery_metrics_publish_backlog_and_clamp_clock_skew() -> None:
|
|||
assert instruments is not None
|
||||
batches_before = _sample(instruments.batch_time, "_count")
|
||||
|
||||
record_deliveries(_DELIVERY_METRICS, 0.5)
|
||||
record_backlog(_BACKLOG, 0.5)
|
||||
record_storage(_STORAGE)
|
||||
|
||||
assert (
|
||||
_sample(instruments.pending),
|
||||
|
|
@ -46,4 +46,17 @@ def test_delivery_metrics_are_skipped_when_prometheus_client_is_not_installed(
|
|||
|
||||
spec.loader.exec_module(module)
|
||||
|
||||
assert module.record_deliveries(_DELIVERY_METRICS, 0.5) is None
|
||||
assert module.record_backlog(_BACKLOG, 0.5) is None
|
||||
assert module.record_storage(_STORAGE) is None
|
||||
|
||||
|
||||
def test_cadence_is_due_once_per_interval() -> None:
|
||||
now = [100.0]
|
||||
cadence = Cadence(lambda: now[0])
|
||||
|
||||
observed = []
|
||||
for moment in (100.0, 101.0, 999.9, 1000.0, 1000.0, 1899.0, 1900.0):
|
||||
now[0] = moment
|
||||
observed.append(cadence.due(900))
|
||||
|
||||
assert observed == [True, False, False, True, False, False, True]
|
||||
|
|
|
|||
|
|
@ -25,6 +25,7 @@ from fakeredis import TcpFakeServer
|
|||
from fastapi import HTTPException
|
||||
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.route_llm_request import (
|
||||
|
|
@ -2545,11 +2546,38 @@ def _mock_deployment(model_name: str, **deployment_limits: int) -> dict[str, obj
|
|||
}
|
||||
|
||||
|
||||
async def _settle_router_accounting() -> None:
|
||||
try:
|
||||
await GLOBAL_LOGGING_WORKER.flush()
|
||||
finally:
|
||||
await GLOBAL_LOGGING_WORKER.stop()
|
||||
while pending := {task for task in asyncio.all_tasks() - {asyncio.current_task()} if not _is_periodic_sync(task)}:
|
||||
_, still_running = await asyncio.wait(pending, timeout=5)
|
||||
assert not still_running, still_running
|
||||
|
||||
|
||||
def _is_periodic_sync(task: asyncio.Task[object]) -> bool:
|
||||
return task.get_coro().__qualname__.endswith(".periodic_sync_in_memory_spend_with_redis")
|
||||
|
||||
|
||||
async def _user_config_request_outcome(
|
||||
redis_port: int | None,
|
||||
routing_settings: Mapping[str, object],
|
||||
deployment_limits: Mapping[str, int],
|
||||
logging_obj: LiteLLMLoggingObj | None,
|
||||
) -> str:
|
||||
outcome: Final = await _unsettled_user_config_request_outcome(
|
||||
redis_port, routing_settings, deployment_limits, logging_obj
|
||||
)
|
||||
await _settle_router_accounting()
|
||||
return outcome
|
||||
|
||||
|
||||
async def _unsettled_user_config_request_outcome(
|
||||
redis_port: int | None,
|
||||
routing_settings: Mapping[str, object],
|
||||
deployment_limits: Mapping[str, int],
|
||||
logging_obj: LiteLLMLoggingObj | None,
|
||||
) -> str:
|
||||
data: Final = {
|
||||
"model": "rate-limited",
|
||||
|
|
@ -2566,8 +2594,6 @@ async def _user_config_request_outcome(
|
|||
await (await route_request(data, None, None, "acompletion"))
|
||||
except (litellm.RateLimitError, RouterRateLimitError):
|
||||
return "429"
|
||||
finally:
|
||||
await asyncio.sleep(0.3)
|
||||
return "200"
|
||||
|
||||
|
||||
|
|
@ -2606,7 +2632,7 @@ async def _proxy_router_outcomes(
|
|||
except (litellm.RateLimitError, RouterRateLimitError):
|
||||
return "429"
|
||||
finally:
|
||||
await asyncio.sleep(0.3)
|
||||
await _settle_router_accounting()
|
||||
return "200"
|
||||
|
||||
try:
|
||||
|
|
@ -2700,16 +2726,21 @@ async def test_concurrent_user_config_requests_each_count_once_against_the_tpm_l
|
|||
tpm_for_three_requests: Final = await _tpm_for_three_requests()
|
||||
|
||||
def request() -> Coroutine[None, None, str]:
|
||||
return _user_config_request_outcome(
|
||||
return _unsettled_user_config_request_outcome(
|
||||
fake_redis_port,
|
||||
{"routing_strategy": "usage-based-routing-v2"},
|
||||
{"tpm": tpm_for_three_requests},
|
||||
_proxy_logging_obj() if arrives_with_logging_obj else None,
|
||||
)
|
||||
|
||||
async def settled(*requests: Coroutine[None, None, str]) -> list[str]:
|
||||
outcomes: Final = await asyncio.gather(*requests)
|
||||
await _settle_router_accounting()
|
||||
return outcomes
|
||||
|
||||
await _wait_out_a_rate_limit_minute_rollover()
|
||||
concurrent_outcomes: Final = await asyncio.gather(request(), request())
|
||||
later_outcomes: Final = [await request() for _ in range(2)]
|
||||
concurrent_outcomes: Final = await settled(request(), request())
|
||||
later_outcomes: Final = [*(await settled(request())), *(await settled(request()))]
|
||||
|
||||
assert [*concurrent_outcomes, *later_outcomes] == ["200", "200", "200", "429"]
|
||||
|
||||
|
|
@ -2717,12 +2748,13 @@ async def test_concurrent_user_config_requests_each_count_once_against_the_tpm_l
|
|||
@pytest.mark.asyncio
|
||||
async def test_overlapping_user_config_requests_share_one_rpm_limit(fake_redis_port: int) -> None:
|
||||
def request() -> Coroutine[None, None, str]:
|
||||
return _user_config_request_outcome(
|
||||
return _unsettled_user_config_request_outcome(
|
||||
fake_redis_port, {"routing_strategy": "usage-based-routing-v2"}, {"rpm": 2}, _proxy_logging_obj()
|
||||
)
|
||||
|
||||
await _wait_out_a_rate_limit_minute_rollover()
|
||||
outcomes: Final = await asyncio.gather(*(request() for _ in range(4)))
|
||||
await _settle_router_accounting()
|
||||
|
||||
assert outcomes.count("200") <= 2, outcomes
|
||||
assert "429" in outcomes, outcomes
|
||||
|
|
@ -2747,13 +2779,14 @@ async def test_a_proxy_router_on_the_same_strategy_leaves_the_user_config_limit_
|
|||
async def outcomes() -> list[str]:
|
||||
first_instance: Final = [await request() for _ in range(2)]
|
||||
user_config_routers.clear()
|
||||
await _settle_router_accounting()
|
||||
second_instance: Final = [await request() for _ in range(3)]
|
||||
return [*first_instance, *second_instance]
|
||||
|
||||
await _wait_out_a_rate_limit_minute_rollover()
|
||||
alone: Final = await outcomes()
|
||||
user_config_routers.clear()
|
||||
await asyncio.sleep(0.3)
|
||||
await _settle_router_accounting()
|
||||
redis.Redis(host="127.0.0.1", port=fake_redis_port).flushall()
|
||||
proxy_router: Final = litellm.Router(
|
||||
routing_strategy=routing_strategy,
|
||||
|
|
@ -2781,7 +2814,7 @@ async def test_user_config_requests_stay_under_the_proxy_provider_budget(
|
|||
)
|
||||
try:
|
||||
await proxy_router.acompletion(model="proxy-model", messages=[{"role": "user", "content": "hi"}])
|
||||
await asyncio.sleep(0.3)
|
||||
await _settle_router_accounting()
|
||||
with pytest.raises(ValueError, match=RouterErrors.no_deployments_with_provider_budget_routing.value):
|
||||
await _user_config_request_outcome(None, user_config_settings, {}, _proxy_logging_obj())
|
||||
finally:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue