mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-18 04:44:37 +00:00
2cc0c79d07
`/global/spend/reset` and `/organization/update` with `spend` wrote only the DB `.spend` column, so with the
longer counter TTL a blocked key or organization stayed blocked until its counter expired. Both now drop the
affected counters (every `spend🔑*` and `spend:team:*` counter through a bounded Redis scan, the one
`spend:org:<id>` counter) so the reset takes effect on the next request.
`RedisCache.async_scan_iter` now returns decoded `str` keys and stops at `count` without a mutable
accumulator, and `delete_cache_keys` takes an `Iterable[str]`, so callers no longer see `bytes` keys.
The audit's integration cells for the TTL (16 counter scopes and surfaces, 7 settings edges, 4 chaos
cells) are added under tests/integration/spend. The leaked-reservation chaos cell asserts the shipped
behavior: a reservation left by a hard-killed worker stays on the counter for the counter TTL.
33 lines
942 B
Python
33 lines
942 B
Python
import os
|
|
from collections.abc import Iterator
|
|
from contextlib import contextmanager
|
|
from hashlib import sha256
|
|
from typing import Final
|
|
|
|
from redis import Redis
|
|
|
|
SPEND_COUNTER_TTL_SECONDS: Final = 3600
|
|
TTL_SLACK_SECONDS: Final = 600
|
|
|
|
|
|
@contextmanager
|
|
def rig_cache() -> Iterator[Redis]:
|
|
with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache:
|
|
yield cache
|
|
|
|
|
|
def key_counter(key: str) -> str:
|
|
return f"spend:key:{sha256(key.encode()).hexdigest()}"
|
|
|
|
|
|
def counter_ttl(cache: Redis, counter_key: str) -> int:
|
|
return int(str(cache.ttl(counter_key)))
|
|
|
|
|
|
def counter_value(cache: Redis, counter_key: str) -> float | None:
|
|
raw: Final = cache.get(counter_key)
|
|
return None if raw is None else float(str(raw, "utf-8") if isinstance(raw, bytes) else str(raw))
|
|
|
|
|
|
def outlives(ttl: int, expected: int = SPEND_COUNTER_TTL_SECONDS) -> bool:
|
|
return expected - TTL_SLACK_SECONDS < ttl <= expected
|