From b82fe539022a861f2027243b6c97e0d7d9135459 Mon Sep 17 00:00:00 2001 From: Thijmen Stavenuiter Date: Sat, 10 Oct 2026 13:07:41 +0200 Subject: [PATCH 1/3] feat(proxy): report when each per-model budget window resets model_max_budget_usage on /key/info, /user/info and /team/info now carries budget_reset_at per model, the open window's start plus its period, or null when no window is open. The window start is read in the same batched lookup as the spend counter --- .../proxy/hooks/model_max_budget_limiter.py | 94 +++++++-- .../key_management_endpoints.py | 3 +- ...test_unit_test_max_model_budget_limiter.py | 178 +++++++++++++++++- .../test_internal_user_endpoints.py | 2 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 3 +- 5 files changed, 250 insertions(+), 30 deletions(-) diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 8bca51784f2..d0b39d98dd1 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -3,6 +3,7 @@ import json import time from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass +from datetime import datetime, timezone from types import MappingProxyType from typing import Final @@ -198,17 +199,28 @@ def _bedrock_candidates(model: str) -> tuple[str, ...]: return (base_model, without_vendor) if without_vendor else (base_model,) +@dataclass(frozen=True, slots=True) +class _WindowReading: + current_spend: float + started_at: float | None + + async def build_model_max_budget_usage( entity_type: Litellm_EntityType, entity_id: str | None, model_max_budget: Mapping[str, object] | None, cache: DualCache | None, + now: datetime | None = None, ) -> dict[str, dict[str, object]]: """Current-window spend per configured budget model, as `/key/info` and `/user/info` report it. `cache` must be the DualCache the limiter writes the counters to; callers read it off the limiter rather than re-deriving it, so a scope that is being blocked can never report zero usage. + + `budget_reset_at` is when the open window ends and its spend counter + expires, which is the window's start plus its period. It is None when no + window is open, because nothing has been charged since the last one ended. """ if cache is None or entity_id is None or not model_max_budget: return {} @@ -221,41 +233,80 @@ async def build_model_max_budget_usage( ) if not budgets: return {} - spend_keys: Final = tuple( - model_budget_spend_cache_key( - entity_type=entity_type, - entity_id=entity_id, - budget_model=budget_model, - budget_duration=budget_config.budget_duration, - ) - for budget_model, budget_config in budgets + windows: Final = await _read_windows( + cache=cache, + spend_keys=tuple( + model_budget_spend_cache_key( + entity_type=entity_type, + entity_id=entity_id, + budget_model=budget_model, + budget_duration=budget_config.budget_duration, + ) + for budget_model, budget_config in budgets + ), + start_time_keys=tuple( + model_budget_start_time_cache_key( + entity_type=entity_type, + entity_id=entity_id, + budget_model=budget_model, + budget_duration=budget_config.budget_duration, + ) + for budget_model, budget_config in budgets + ), ) - current_spends: Final = await _current_window_spends(cache=cache, spend_keys=spend_keys) + current_time: Final = (now or datetime.now(timezone.utc)).timestamp() return { budget_model: { - "current_spend": round(current_spend, 4), + "current_spend": round(window.current_spend, 4), "budget_limit": budget_config.max_budget, "time_period": budget_config.budget_duration, + "budget_reset_at": _window_reset_at( + started_at=window.started_at, + budget_duration=budget_config.budget_duration, + now=current_time, + ), } - for (budget_model, budget_config), current_spend in zip(budgets, current_spends, strict=True) + for (budget_model, budget_config), window in zip(budgets, windows, strict=True) } @with_service_target("model_budgets") -async def _current_window_spends(cache: DualCache, spend_keys: Sequence[str]) -> tuple[float, ...]: - """Redis holds the window total across replicas; the in-memory copy is one replica's share.""" - keys: Final = list(spend_keys) +async def _read_windows( + cache: DualCache, spend_keys: Sequence[str], start_time_keys: Sequence[str] +) -> tuple[_WindowReading, ...]: + """Every window's spend and start time in one batched read, however many models are budgeted. + + Redis holds the window total across replicas; the in-memory copy is one replica's share. + """ + keys: Final = [*spend_keys, *start_time_keys] redis_cache: Final = cache.redis_cache if redis_cache is not None: shared: Final = await redis_cache.async_batch_get_cache(key_list=keys) - return tuple(_as_spend(shared.get(key)) for key in keys) + return _readings(tuple(_as_float(shared.get(key)) for key in keys), window_count=len(spend_keys)) # async_batch_get_cache returns None if it fails internally, and its result is # index-aligned with `keys` otherwise. An unusable result reads as a miss, # which is what a never-written counter already reads as. batched: Final = await cache.async_batch_get_cache(keys=keys) if not isinstance(batched, list) or len(batched) != len(keys): - return (0.0,) * len(keys) - return tuple(_as_spend(current_spend) for current_spend in batched) + return tuple(_WindowReading(current_spend=0.0, started_at=None) for _ in spend_keys) + return _readings(tuple(_as_float(value) for value in batched), window_count=len(spend_keys)) + + +def _readings(values: Sequence[float | None], window_count: int) -> tuple[_WindowReading, ...]: + """`values` holds every window's spend followed by every window's start, in the same order.""" + return tuple( + _WindowReading(current_spend=spend or 0.0, started_at=started_at) + for spend, started_at in zip(values[:window_count], values[window_count:], strict=True) + ) + + +def _window_reset_at(started_at: float | None, budget_duration: str | None, now: float) -> str | None: + if started_at is None or budget_duration is None: + return None + ends_at: Final = started_at + duration_in_seconds(budget_duration) + if ends_at <= now: + return None + return datetime.fromtimestamp(ends_at, tz=timezone.utc).isoformat() def _usable_budget_config(raw_budget_config: object) -> BudgetConfig | None: @@ -276,6 +327,15 @@ def _as_spend(current_spend: object) -> float: return 0.0 +def _as_float(value: object) -> float | None: + if value is None: + return None + try: + return float(value) # pyright: ignore[reportArgumentType] # non-numeric falls to the except + except (TypeError, ValueError): + return None + + def _resolve_entity_model_budgets( model: str, entity_budgets: Iterable[tuple[Litellm_EntityType, str | None, object]], diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 6e226868918..f4210f9cf00 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -4612,7 +4612,8 @@ async def info_key_fn( hour), so subtracting budget_duration from it does not give the window's start - model_max_budget: dict - Per-model budgets, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}} - model_max_budget_usage: dict | None - Current-window spend per model, present only when - the key has per-model budgets + the key has per-model budgets. Each entry's budget_reset_at is when that model's open + window ends and its spend resets, or None when no window is open - budget_limits: list | None - Concurrent budget windows, exactly as stored - budget_limits_usage: dict | None - Current-window spend per budget window, e.g. {"1h": {"current_spend": 0.0009}}, present only when the key has budget windows diff --git a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py index dc9dac9a57b..b17b5c34e37 100644 --- a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py +++ b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py @@ -8,7 +8,8 @@ import pytest import litellm from litellm.caching.caching import DualCache from litellm.caching.redis_cache import RedisCache -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone +from typing import Final from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import Litellm_EntityType @@ -16,6 +17,8 @@ from litellm.proxy.hooks.model_max_budget_limiter import ( _budget_model_candidates, PROXY_VirtualKeyModelMaxBudgetLimiter, build_model_max_budget_usage, + model_budget_spend_cache_key, + model_budget_start_time_cache_key, resolve_model_budget, ) from litellm.proxy._types import UserAPIKeyAuth @@ -23,6 +26,17 @@ from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.utils import BudgetConfig as GenericBudgetInfo +async def _reset_at_of_open_window( + cache: DualCache, entity_type: Litellm_EntityType, entity_id: str, budget_model: str, period: str +) -> str: + started_at = await cache.async_get_cache( + key=model_budget_start_time_cache_key( + entity_type=entity_type, entity_id=entity_id, budget_model=budget_model, budget_duration=period + ) + ) + return datetime.fromtimestamp(started_at + duration_in_seconds(period), tz=timezone.utc).isoformat() + + # Test class setup @pytest.fixture def budget_limiter(): @@ -680,6 +694,9 @@ async def test_logged_spend_is_visible_to_key_info_usage_and_enforcement(request "current_spend": 0.75, "budget_limit": 1.0, "time_period": "1d", + "budget_reset_at": await _reset_at_of_open_window( + dual_cache, Litellm_EntityType.KEY, key_hash, "gpt-4", "1d" + ), } } @@ -747,7 +764,16 @@ async def test_user_model_budget_is_tracked_and_enforced(): entity_id=user_id, model_max_budget=user_model_max_budget, cache=dual_cache, - ) == {"gpt-4": {"current_spend": 1.5, "budget_limit": 1.0, "time_period": "1mo"}} + ) == { + "gpt-4": { + "current_spend": 1.5, + "budget_limit": 1.0, + "time_period": "1mo", + "budget_reset_at": await _reset_at_of_open_window( + dual_cache, Litellm_EntityType.USER, user_id, "gpt-4", "1mo" + ), + } + } with pytest.raises(litellm.BudgetExceededError) as exc: await limiter.is_user_within_model_budget( @@ -857,7 +883,9 @@ async def test_build_model_max_budget_usage_skips_unusable_entries(): }, cache=dual_cache, ) - assert usage == {"gpt-4": {"current_spend": 3.0, "budget_limit": 10.0, "time_period": "1d"}} + assert usage == { + "gpt-4": {"current_spend": 3.0, "budget_limit": 10.0, "time_period": "1d", "budget_reset_at": None} + } @pytest.mark.asyncio @@ -896,6 +924,9 @@ async def test_bedrock_traffic_charges_the_bare_family_name_budget(): "current_spend": 1.5, "budget_limit": 1.0, "time_period": "18h", + "budget_reset_at": await _reset_at_of_open_window( + dual_cache, Litellm_EntityType.KEY, key_hash, "claude-opus-4-8", "18h" + ), } } @@ -963,7 +994,16 @@ async def test_user_model_budget_window_resets_when_the_period_elapses(): entity_id=user_id, model_max_budget=user_model_max_budget, cache=dual_cache, - ) == {"gpt-4": {"current_spend": 1.5, "budget_limit": 1.0, "time_period": "1mo"}} + ) == { + "gpt-4": { + "current_spend": 1.5, + "budget_limit": 1.0, + "time_period": "1mo", + "budget_reset_at": await _reset_at_of_open_window( + dual_cache, Litellm_EntityType.USER, user_id, "gpt-4", "1mo" + ), + } + } @pytest.mark.asyncio @@ -995,7 +1035,7 @@ async def test_a_zero_dollar_cap_is_reported_as_a_cap_not_as_absent(): entity_id="hash-zero", model_max_budget={"gpt-4": {"budget_limit": 0, "time_period": "1d"}}, cache=DualCache(), - ) == {"gpt-4": {"current_spend": 0.0, "budget_limit": 0.0, "time_period": "1d"}} + ) == {"gpt-4": {"current_spend": 0.0, "budget_limit": 0.0, "time_period": "1d", "budget_reset_at": None}} @pytest.mark.asyncio @@ -1033,7 +1073,7 @@ async def test_usage_report_reads_every_counter_in_one_batched_lookup(): budget = {f"model-{i}": {"budget_limit": 1.0, "time_period": "1d"} for i in range(50)} with ( - patch.object(dual_cache, "async_batch_get_cache", new=AsyncMock(return_value=[None] * 50)) as batched, + patch.object(dual_cache, "async_batch_get_cache", new=AsyncMock(return_value=[None] * 100)) as batched, patch.object(dual_cache, "async_get_cache", new=AsyncMock()) as single, ): usage = await build_model_max_budget_usage( @@ -1044,7 +1084,8 @@ async def test_usage_report_reads_every_counter_in_one_batched_lookup(): ) assert batched.await_count == 1 - assert len(batched.await_args.kwargs["keys"]) == 50 + # One spend counter and one window start per budgeted model, all in the same round trip. + assert len(batched.await_args.kwargs["keys"]) == 100 assert single.await_count == 0 assert len(usage) == 50 @@ -1063,7 +1104,75 @@ async def test_usage_report_survives_a_batch_lookup_that_returns_nothing(): entity_id="hash-none", model_max_budget={"gpt-4": {"budget_limit": 1.0, "time_period": "1d"}}, cache=dual_cache, - ) == {"gpt-4": {"current_spend": 0.0, "budget_limit": 1.0, "time_period": "1d"}} + ) == {"gpt-4": {"current_spend": 0.0, "budget_limit": 1.0, "time_period": "1d", "budget_reset_at": None}} + + +_WINDOW_OPENED_AT: Final = datetime(2026, 10, 10, 12, 0, tzinfo=timezone.utc) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "seconds_since_window_opened,expected_reset_at", + [ + (0, "2026-10-10T13:00:00+00:00"), + (600, "2026-10-10T13:00:00+00:00"), + (3599, "2026-10-10T13:00:00+00:00"), + (3600, None), + (7200, None), + ], + ids=["just_opened", "mid_window", "last_second", "ended_exactly_now", "long_ended"], +) +async def test_usage_reports_when_the_open_window_resets(seconds_since_window_opened, expected_reset_at): + """ + The counter expires at the window's start plus its period, and that is when the + model is admitted again. A window that has already run its period is not open, + even if its start time is still sitting in the cache, so no reset is reported. + """ + dual_cache = DualCache() + key = dict(entity_type=Litellm_EntityType.KEY, entity_id="vk-hash", budget_model="gpt-4", budget_duration="1h") + await dual_cache.async_set_cache(key=model_budget_spend_cache_key(**key), value=0.4) + await dual_cache.async_set_cache(key=model_budget_start_time_cache_key(**key), value=_WINDOW_OPENED_AT.timestamp()) + + usage = await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id="vk-hash", + model_max_budget={"gpt-4": {"budget_limit": 1.0, "time_period": "1h"}}, + cache=dual_cache, + now=_WINDOW_OPENED_AT + timedelta(seconds=seconds_since_window_opened), + ) + + assert usage["gpt-4"]["budget_reset_at"] == expected_reset_at + + +@pytest.mark.asyncio +async def test_each_budgeted_model_reports_its_own_window_reset(): + """Two models on one key own separate windows, so one model's reset must never be reported for the other.""" + dual_cache = DualCache() + hourly = dict(entity_type=Litellm_EntityType.KEY, entity_id="vk-hash", budget_model="gpt-4", budget_duration="1h") + daily = dict(entity_type=Litellm_EntityType.KEY, entity_id="vk-hash", budget_model="claude-3", budget_duration="1d") + await dual_cache.async_set_cache( + key=model_budget_start_time_cache_key(**hourly), value=_WINDOW_OPENED_AT.timestamp() + ) + await dual_cache.async_set_cache( + key=model_budget_start_time_cache_key(**daily), + value=(_WINDOW_OPENED_AT - timedelta(hours=5)).timestamp(), + ) + + usage = await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id="vk-hash", + model_max_budget={ + "gpt-4": {"budget_limit": 1.0, "time_period": "1h"}, + "claude-3": {"budget_limit": 1.0, "time_period": "1d"}, + "never-used": {"budget_limit": 1.0, "time_period": "1d"}, + }, + cache=dual_cache, + now=_WINDOW_OPENED_AT, + ) + + assert usage["gpt-4"]["budget_reset_at"] == "2026-10-10T13:00:00+00:00" + assert usage["claude-3"]["budget_reset_at"] == "2026-10-11T07:00:00+00:00" + assert usage["never-used"]["budget_reset_at"] is None @pytest.mark.asyncio @@ -1097,7 +1206,16 @@ async def test_one_malformed_scope_does_not_abort_the_other_scopes(): entity_id="hash-mixed", model_max_budget=key_budget, cache=dual_cache, - ) == {"gpt-4": {"current_spend": 0.25, "budget_limit": 10.0, "time_period": "1d"}} + ) == { + "gpt-4": { + "current_spend": 0.25, + "budget_limit": 10.0, + "time_period": "1d", + "budget_reset_at": await _reset_at_of_open_window( + dual_cache, Litellm_EntityType.KEY, "hash-mixed", "gpt-4", "1d" + ), + } + } @pytest.mark.asyncio @@ -1441,6 +1559,37 @@ async def test_spend_logged_on_one_replica_is_enforced_and_reported_on_another() await replica_c.is_key_within_model_budget(user_api_key, "gpt-4") +@pytest.mark.asyncio +async def test_a_replica_reports_the_reset_of_a_window_another_replica_opened(): + """ + The window start is written by whichever replica served the first charge, so a + replica that has never seen this key must read it from the shared Redis, or it + would report no reset for a model it is actively refusing. + """ + shared_redis = _SharedFakeRedis() + opening_replica = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache(redis_cache=shared_redis)) + other_replica = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache(redis_cache=shared_redis)) + model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1h"}} + + await _log_spend(opening_replica, key_hash="vk-shared", model_max_budget=model_max_budget, response_cost=0.5) + + usage_on_other: Final = await build_model_max_budget_usage( + entity_type=Litellm_EntityType.KEY, + entity_id="vk-shared", + model_max_budget=model_max_budget, + cache=other_replica.dual_cache, + ) + started_at: Final = shared_redis._store[ + model_budget_start_time_cache_key( + entity_type=Litellm_EntityType.KEY, entity_id="vk-shared", budget_model="gpt-4", budget_duration="1h" + ) + ] + assert ( + usage_on_other["gpt-4"]["budget_reset_at"] + == (datetime.fromtimestamp(started_at, tz=timezone.utc) + timedelta(hours=1)).isoformat() + ) + + def _log_success(limiter, **kwargs): return limiter.async_log_success_event( _success_kwargs(**kwargs), response_obj=None, start_time=None, end_time=None @@ -1496,7 +1645,16 @@ async def test_team_model_budget_is_shared_by_every_key_without_an_override(requ entity_id="team-1", model_max_budget=team_model_max_budget, cache=dual_cache, - ) == {"gpt-4": {"current_spend": pytest.approx(1.2), "budget_limit": 1.0, "time_period": "1d"}} + ) == { + "gpt-4": { + "current_spend": pytest.approx(1.2), + "budget_limit": 1.0, + "time_period": "1d", + "budget_reset_at": await _reset_at_of_open_window( + dual_cache, Litellm_EntityType.TEAM, "team-1", "gpt-4", "1d" + ), + } + } @pytest.mark.asyncio diff --git a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py index 39265b37b95..c789421b3d6 100644 --- a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py @@ -3499,7 +3499,7 @@ async def test_user_info_v2_response_shape(mocker): # these two, so dropping them makes a save replace the user's budgets. assert response_dict["model_max_budget"] == {"gpt-3.5-turbo": {"budget_limit": 5.0, "time_period": "30d"}} assert response_dict["model_max_budget_usage"] == { - "gpt-3.5-turbo": {"current_spend": 0.0, "budget_limit": 5.0, "time_period": "30d"} + "gpt-3.5-turbo": {"current_spend": 0.0, "budget_limit": 5.0, "time_period": "30d", "budget_reset_at": None} } # Verify teams is a list of strings (team IDs), not team objects diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 53ecf9b2d75..86a41afceda 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -8503,7 +8503,8 @@ export interface paths { * hour), so subtracting budget_duration from it does not give the window's start * - model_max_budget: dict - Per-model budgets, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}} * - model_max_budget_usage: dict | None - Current-window spend per model, present only when - * the key has per-model budgets + * the key has per-model budgets. Each entry's budget_reset_at is when that model's open + * window ends and its spend resets, or None when no window is open * - budget_limits: list | None - Concurrent budget windows, exactly as stored * - budget_limits_usage: dict | None - Current-window spend per budget window, e.g. * {"1h": {"current_spend": 0.0009}}, present only when the key has budget windows From 25d52df97b284324582a09d0de33d333fe2ca856 Mon Sep 17 00:00:00 2001 From: Thijmen Stavenuiter Date: Sat, 10 Oct 2026 17:04:24 +0200 Subject: [PATCH 2/3] fix(proxy): count a monthly budget window from its own start when reporting its reset duration_in_seconds measures Nmo from now, so a window opened on January 15 and read on February 10 reported February 12, then null while the counter still refused requests until February 15. duration_in_seconds_from measures the month from the window start, and duration_in_seconds now shares that month arithmetic --- litellm/litellm_core_utils/duration_parser.py | 54 +++++++-------- .../proxy/hooks/model_max_budget_limiter.py | 6 +- .../test_duration_parser.py | 29 +++++++++ ...test_unit_test_max_model_budget_limiter.py | 65 +++++++++++++++++++ 4 files changed, 121 insertions(+), 33 deletions(-) diff --git a/litellm/litellm_core_utils/duration_parser.py b/litellm/litellm_core_utils/duration_parser.py index fe8b71e7623..9e6ad7493f6 100644 --- a/litellm/litellm_core_utils/duration_parser.py +++ b/litellm/litellm_core_utils/duration_parser.py @@ -74,38 +74,29 @@ def duration_in_seconds(duration: str) -> int: elif unit == "w": return value * 604800 elif unit == "mo": - now: Final = time_module.time() - current_time: Final = datetime.fromtimestamp(now) - - # Calculate target month and year, handling overflow past December - total_months: Final = current_time.month - 1 + value # 0-indexed months - target_year: Final = current_time.year + total_months // 12 - target_month: Final = total_months % 12 + 1 # back to 1-indexed - - # Determine the day to set for next month - target_day = current_time.day - last_day_of_target_month: Final = get_last_day_of_month(target_year, target_month) - - target_day = min(target_day, last_day_of_target_month) - - next_month: Final = datetime( - year=target_year, - month=target_month, - day=target_day, - hour=current_time.hour, - minute=current_time.minute, - second=current_time.second, - microsecond=current_time.microsecond, - ) - - # Calculate the duration until the first day of the next month - duration_until_next_month: Final = next_month - current_time - return int(duration_until_next_month.total_seconds()) + return _months_in_seconds(start=datetime.fromtimestamp(time_module.time()), months=value) else: raise ValueError(f"Unsupported duration unit, passed duration: {duration}") +def duration_in_seconds_from(start: datetime, duration: str) -> int: + """How long a period of `duration` lasts when it begins at `start`. + + `duration_in_seconds` measures `Nmo` from now, so a monthly window that opened in a + month of a different length needs its own start to know when it ends. Every other + unit is a fixed length and ignores `start`. + """ + value, unit = _extract_from_regex(duration=_normalize_duration(duration)) + if unit == "mo": + return _months_in_seconds(start=start, months=value) + return duration_in_seconds(duration) + + +def _months_in_seconds(start: datetime, months: int) -> int: + return int((_shift_months(start, months) - start).total_seconds()) + + def get_next_standardized_reset_time( duration: str, current_time: datetime, @@ -166,8 +157,9 @@ def get_next_standardized_reset_time( return base_midnight + timedelta(days=1) -def _subtract_months(moment: datetime, months: int) -> datetime: - total_months: Final = moment.year * 12 + moment.month - 1 - months +def _shift_months(moment: datetime, months: int) -> datetime: + """`moment` moved by `months` calendar months, back when negative, clamped to the target month's last day.""" + total_months: Final = moment.year * 12 + moment.month - 1 + months year, month_index = divmod(total_months, 12) month: Final = month_index + 1 return moment.replace(year=year, month=month, day=min(moment.day, get_last_day_of_month(year, month))) @@ -182,9 +174,9 @@ def get_budget_window_start(duration: str, reset_at: datetime) -> datetime: return reset_at - timedelta(days=1) match unit: case "mo": - return _subtract_months(reset_at, value) + return _shift_months(reset_at, -value) case "d": - return _subtract_months(reset_at, 1) if value == 30 else reset_at - timedelta(days=value) + return _shift_months(reset_at, -1) if value == 30 else reset_at - timedelta(days=value) case "w": return reset_at - timedelta(weeks=value) case "h": diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index d0b39d98dd1..a518d60de05 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -14,7 +14,7 @@ from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import Span -from litellm.litellm_core_utils.duration_parser import duration_in_seconds +from litellm.litellm_core_utils.duration_parser import duration_in_seconds, duration_in_seconds_from from litellm.llms.bedrock.common_utils import get_bedrock_base_model from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth from litellm.router_strategy.budget_limiter import RouterBudgetLimiting @@ -303,7 +303,9 @@ def _readings(values: Sequence[float | None], window_count: int) -> tuple[_Windo def _window_reset_at(started_at: float | None, budget_duration: str | None, now: float) -> str | None: if started_at is None or budget_duration is None: return None - ends_at: Final = started_at + duration_in_seconds(budget_duration) + ends_at: Final = started_at + duration_in_seconds_from( + start=datetime.fromtimestamp(started_at), duration=budget_duration + ) if ends_at <= now: return None return datetime.fromtimestamp(ends_at, tz=timezone.utc).isoformat() diff --git a/tests/unit/litellm_core_utils/test_duration_parser.py b/tests/unit/litellm_core_utils/test_duration_parser.py index f5e50d67e71..9212da633e0 100644 --- a/tests/unit/litellm_core_utils/test_duration_parser.py +++ b/tests/unit/litellm_core_utils/test_duration_parser.py @@ -7,6 +7,7 @@ from zoneinfo import ZoneInfo import litellm.litellm_core_utils.duration_parser as duration_parser from litellm.litellm_core_utils.duration_parser import ( duration_in_seconds, + duration_in_seconds_from, get_budget_window_start, get_next_standardized_reset_time, ) @@ -373,5 +374,33 @@ class TestGetBudgetWindowStart(unittest.TestCase): ) +class TestDurationInSecondsFrom(unittest.TestCase): + def test_a_month_lasts_as_long_as_the_month_it_starts_in(self): + cases: Final = ( + (datetime(2026, 1, 15, 12), "1mo", 31), + (datetime(2026, 2, 10, 12), "1mo", 28), + (datetime(2026, 1, 31, 12), "1mo", 28), + (datetime(2028, 1, 31, 12), "1mo", 29), + (datetime(2026, 12, 15, 12), "1mo", 31), + (datetime(2026, 1, 15, 12), "2mo", 59), + ) + for start, duration, days in cases: + with self.subTest(start=start, duration=duration): + self.assertEqual(duration_in_seconds_from(start, duration), days * 86400) + + def test_fixed_length_units_do_not_depend_on_the_start(self): + for duration in ("45s", "30m", "1h", "1d", "7d", "2w", "30d", "daily", "monthly"): + with self.subTest(duration=duration): + self.assertEqual( + duration_in_seconds_from(datetime(2026, 2, 10, 12), duration), + duration_in_seconds(duration), + ) + + def test_a_month_starting_now_matches_duration_in_seconds(self): + now: Final = datetime(2026, 1, 31, 9, 30, 0, 250000) + with patch.object(duration_parser.time_module, "time", return_value=now.timestamp()): + self.assertEqual(duration_in_seconds("1mo"), duration_in_seconds_from(now, "1mo")) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py index b17b5c34e37..b523e5fced0 100644 --- a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py +++ b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py @@ -1144,6 +1144,71 @@ async def test_usage_reports_when_the_open_window_resets(seconds_since_window_op assert usage["gpt-4"]["budget_reset_at"] == expected_reset_at +_DAY: Final = 86400 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "opened_at,read_after_days,window_days", + [ + (datetime(2026, 1, 15, 12), 26, 31), + (datetime(2026, 1, 15, 12), 29, 31), + (datetime(2026, 1, 31, 12), 20, 28), + (datetime(2026, 12, 15, 12), 30, 31), + ], + ids=["read_in_a_shorter_month", "read_past_the_shorter_month_length", "clamped_to_february", "across_new_year"], +) +async def test_a_monthly_window_resets_a_calendar_month_after_it_opened(opened_at, read_after_days, window_days): + """ + A month is only as long as the month the window opened in, which is the lifetime the + counter was given when it opened. Measuring the month from the read instead made a + window opened on January 15 and read on February 10 report February 12, and then null + while the counter still refused requests until February 15. + """ + dual_cache = DualCache() + started_at: Final = opened_at.timestamp() + await dual_cache.async_set_cache( + key=model_budget_start_time_cache_key( + entity_type=Litellm_EntityType.USER, entity_id="user-1", budget_model="gpt-4", budget_duration="1mo" + ), + value=started_at, + ) + + usage = await build_model_max_budget_usage( + entity_type=Litellm_EntityType.USER, + entity_id="user-1", + model_max_budget={"gpt-4": {"budget_limit": 1.0, "time_period": "1mo"}}, + cache=dual_cache, + now=datetime.fromtimestamp(started_at + read_after_days * _DAY, tz=timezone.utc), + ) + + assert usage["gpt-4"]["budget_reset_at"] == ( + datetime.fromtimestamp(started_at + window_days * _DAY, tz=timezone.utc).isoformat() + ) + + +@pytest.mark.asyncio +async def test_a_monthly_window_reports_no_reset_once_its_month_has_passed(): + dual_cache = DualCache() + started_at: Final = datetime(2026, 1, 15, 12).timestamp() + await dual_cache.async_set_cache( + key=model_budget_start_time_cache_key( + entity_type=Litellm_EntityType.USER, entity_id="user-1", budget_model="gpt-4", budget_duration="1mo" + ), + value=started_at, + ) + + usage = await build_model_max_budget_usage( + entity_type=Litellm_EntityType.USER, + entity_id="user-1", + model_max_budget={"gpt-4": {"budget_limit": 1.0, "time_period": "1mo"}}, + cache=dual_cache, + now=datetime.fromtimestamp(started_at + 31 * _DAY, tz=timezone.utc), + ) + + assert usage["gpt-4"]["budget_reset_at"] is None + + @pytest.mark.asyncio async def test_each_budgeted_model_reports_its_own_window_reset(): """Two models on one key own separate windows, so one model's reset must never be reported for the other.""" From fcdd92c2bcde031bce78668470559d5fce7aa449 Mon Sep 17 00:00:00 2001 From: Thijmen Stavenuiter Date: Sat, 10 Oct 2026 17:29:17 +0200 Subject: [PATCH 3/3] fix(proxy): measure months from timezone-aware datetimes and assert whole usage values datetime.fromtimestamp without a tz tripped DTZ006. The month arithmetic now runs on an aware local datetime, which keeps the same wall-clock result. The duration tests move to parametrized functions, since subTest payloads carrying datetimes cannot cross the xdist worker boundary, and the new usage tests compare the whole entry --- litellm/litellm_core_utils/duration_parser.py | 4 +- .../proxy/hooks/model_max_budget_limiter.py | 2 +- .../test_duration_parser.py | 51 ++++++++------- ...test_unit_test_max_model_budget_limiter.py | 65 +++++++++++++------ 4 files changed, 76 insertions(+), 46 deletions(-) diff --git a/litellm/litellm_core_utils/duration_parser.py b/litellm/litellm_core_utils/duration_parser.py index 9e6ad7493f6..d52716951df 100644 --- a/litellm/litellm_core_utils/duration_parser.py +++ b/litellm/litellm_core_utils/duration_parser.py @@ -74,7 +74,9 @@ def duration_in_seconds(duration: str) -> int: elif unit == "w": return value * 604800 elif unit == "mo": - return _months_in_seconds(start=datetime.fromtimestamp(time_module.time()), months=value) + return _months_in_seconds( + start=datetime.fromtimestamp(time_module.time(), tz=timezone.utc).astimezone(), months=value + ) else: raise ValueError(f"Unsupported duration unit, passed duration: {duration}") diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index a518d60de05..f562956f77a 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -304,7 +304,7 @@ def _window_reset_at(started_at: float | None, budget_duration: str | None, now: if started_at is None or budget_duration is None: return None ends_at: Final = started_at + duration_in_seconds_from( - start=datetime.fromtimestamp(started_at), duration=budget_duration + start=datetime.fromtimestamp(started_at, tz=timezone.utc).astimezone(), duration=budget_duration ) if ends_at <= now: return None diff --git a/tests/unit/litellm_core_utils/test_duration_parser.py b/tests/unit/litellm_core_utils/test_duration_parser.py index 9212da633e0..f14015effe6 100644 --- a/tests/unit/litellm_core_utils/test_duration_parser.py +++ b/tests/unit/litellm_core_utils/test_duration_parser.py @@ -4,6 +4,8 @@ from typing import Final from unittest.mock import patch from zoneinfo import ZoneInfo +import pytest + import litellm.litellm_core_utils.duration_parser as duration_parser from litellm.litellm_core_utils.duration_parser import ( duration_in_seconds, @@ -374,32 +376,33 @@ class TestGetBudgetWindowStart(unittest.TestCase): ) -class TestDurationInSecondsFrom(unittest.TestCase): - def test_a_month_lasts_as_long_as_the_month_it_starts_in(self): - cases: Final = ( - (datetime(2026, 1, 15, 12), "1mo", 31), - (datetime(2026, 2, 10, 12), "1mo", 28), - (datetime(2026, 1, 31, 12), "1mo", 28), - (datetime(2028, 1, 31, 12), "1mo", 29), - (datetime(2026, 12, 15, 12), "1mo", 31), - (datetime(2026, 1, 15, 12), "2mo", 59), - ) - for start, duration, days in cases: - with self.subTest(start=start, duration=duration): - self.assertEqual(duration_in_seconds_from(start, duration), days * 86400) +@pytest.mark.parametrize( + "start,duration,days", + [ + (datetime(2026, 1, 15, 12, tzinfo=timezone.utc), "1mo", 31), + (datetime(2026, 2, 10, 12, tzinfo=timezone.utc), "1mo", 28), + (datetime(2026, 1, 31, 12, tzinfo=timezone.utc), "1mo", 28), + (datetime(2028, 1, 31, 12, tzinfo=timezone.utc), "1mo", 29), + (datetime(2026, 12, 15, 12, tzinfo=timezone.utc), "1mo", 31), + (datetime(2026, 1, 15, 12, tzinfo=timezone.utc), "2mo", 59), + ], + ids=["january", "february", "clamped_to_february", "leap_february", "across_new_year", "two_months"], +) +def test_a_month_lasts_as_long_as_the_month_it_starts_in(start, duration, days): + assert duration_in_seconds_from(start, duration) == days * 86400 - def test_fixed_length_units_do_not_depend_on_the_start(self): - for duration in ("45s", "30m", "1h", "1d", "7d", "2w", "30d", "daily", "monthly"): - with self.subTest(duration=duration): - self.assertEqual( - duration_in_seconds_from(datetime(2026, 2, 10, 12), duration), - duration_in_seconds(duration), - ) - def test_a_month_starting_now_matches_duration_in_seconds(self): - now: Final = datetime(2026, 1, 31, 9, 30, 0, 250000) - with patch.object(duration_parser.time_module, "time", return_value=now.timestamp()): - self.assertEqual(duration_in_seconds("1mo"), duration_in_seconds_from(now, "1mo")) +@pytest.mark.parametrize("duration", ["45s", "30m", "1h", "1d", "7d", "2w", "30d", "daily", "monthly"]) +def test_fixed_length_units_do_not_depend_on_the_start(duration): + assert duration_in_seconds_from(datetime(2026, 2, 10, 12, tzinfo=timezone.utc), duration) == duration_in_seconds( + duration + ) + + +def test_a_month_starting_now_matches_duration_in_seconds(): + now: Final = datetime(2026, 1, 31, 9, 30, 0, 250000, tzinfo=timezone.utc).astimezone() + with patch.object(duration_parser.time_module, "time", return_value=now.timestamp()): + assert duration_in_seconds("1mo") == duration_in_seconds_from(now, "1mo") if __name__ == "__main__": diff --git a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py index b523e5fced0..98170a406ea 100644 --- a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py +++ b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py @@ -1141,7 +1141,9 @@ async def test_usage_reports_when_the_open_window_resets(seconds_since_window_op now=_WINDOW_OPENED_AT + timedelta(seconds=seconds_since_window_opened), ) - assert usage["gpt-4"]["budget_reset_at"] == expected_reset_at + assert usage == { + "gpt-4": {"current_spend": 0.4, "budget_limit": 1.0, "time_period": "1h", "budget_reset_at": expected_reset_at} + } _DAY: Final = 86400 @@ -1151,10 +1153,10 @@ _DAY: Final = 86400 @pytest.mark.parametrize( "opened_at,read_after_days,window_days", [ - (datetime(2026, 1, 15, 12), 26, 31), - (datetime(2026, 1, 15, 12), 29, 31), - (datetime(2026, 1, 31, 12), 20, 28), - (datetime(2026, 12, 15, 12), 30, 31), + (datetime(2026, 1, 15, 12, tzinfo=timezone.utc), 26, 31), + (datetime(2026, 1, 15, 12, tzinfo=timezone.utc), 29, 31), + (datetime(2026, 1, 31, 12, tzinfo=timezone.utc), 20, 28), + (datetime(2026, 12, 15, 12, tzinfo=timezone.utc), 30, 31), ], ids=["read_in_a_shorter_month", "read_past_the_shorter_month_length", "clamped_to_february", "across_new_year"], ) @@ -1179,23 +1181,28 @@ async def test_a_monthly_window_resets_a_calendar_month_after_it_opened(opened_a entity_id="user-1", model_max_budget={"gpt-4": {"budget_limit": 1.0, "time_period": "1mo"}}, cache=dual_cache, - now=datetime.fromtimestamp(started_at + read_after_days * _DAY, tz=timezone.utc), + now=opened_at + timedelta(days=read_after_days), ) - assert usage["gpt-4"]["budget_reset_at"] == ( - datetime.fromtimestamp(started_at + window_days * _DAY, tz=timezone.utc).isoformat() - ) + assert usage == { + "gpt-4": { + "current_spend": 0.0, + "budget_limit": 1.0, + "time_period": "1mo", + "budget_reset_at": (opened_at + timedelta(days=window_days)).isoformat(), + } + } @pytest.mark.asyncio async def test_a_monthly_window_reports_no_reset_once_its_month_has_passed(): dual_cache = DualCache() - started_at: Final = datetime(2026, 1, 15, 12).timestamp() + opened_at: Final = datetime(2026, 1, 15, 12, tzinfo=timezone.utc) await dual_cache.async_set_cache( key=model_budget_start_time_cache_key( entity_type=Litellm_EntityType.USER, entity_id="user-1", budget_model="gpt-4", budget_duration="1mo" ), - value=started_at, + value=opened_at.timestamp(), ) usage = await build_model_max_budget_usage( @@ -1203,10 +1210,12 @@ async def test_a_monthly_window_reports_no_reset_once_its_month_has_passed(): entity_id="user-1", model_max_budget={"gpt-4": {"budget_limit": 1.0, "time_period": "1mo"}}, cache=dual_cache, - now=datetime.fromtimestamp(started_at + 31 * _DAY, tz=timezone.utc), + now=opened_at + timedelta(days=31), ) - assert usage["gpt-4"]["budget_reset_at"] is None + assert usage == { + "gpt-4": {"current_spend": 0.0, "budget_limit": 1.0, "time_period": "1mo", "budget_reset_at": None} + } @pytest.mark.asyncio @@ -1235,9 +1244,21 @@ async def test_each_budgeted_model_reports_its_own_window_reset(): now=_WINDOW_OPENED_AT, ) - assert usage["gpt-4"]["budget_reset_at"] == "2026-10-10T13:00:00+00:00" - assert usage["claude-3"]["budget_reset_at"] == "2026-10-11T07:00:00+00:00" - assert usage["never-used"]["budget_reset_at"] is None + assert usage == { + "gpt-4": { + "current_spend": 0.0, + "budget_limit": 1.0, + "time_period": "1h", + "budget_reset_at": "2026-10-10T13:00:00+00:00", + }, + "claude-3": { + "current_spend": 0.0, + "budget_limit": 1.0, + "time_period": "1d", + "budget_reset_at": "2026-10-11T07:00:00+00:00", + }, + "never-used": {"current_spend": 0.0, "budget_limit": 1.0, "time_period": "1d", "budget_reset_at": None}, + } @pytest.mark.asyncio @@ -1649,10 +1670,14 @@ async def test_a_replica_reports_the_reset_of_a_window_another_replica_opened(): entity_type=Litellm_EntityType.KEY, entity_id="vk-shared", budget_model="gpt-4", budget_duration="1h" ) ] - assert ( - usage_on_other["gpt-4"]["budget_reset_at"] - == (datetime.fromtimestamp(started_at, tz=timezone.utc) + timedelta(hours=1)).isoformat() - ) + assert usage_on_other == { + "gpt-4": { + "current_spend": 0.5, + "budget_limit": 1.0, + "time_period": "1h", + "budget_reset_at": (datetime.fromtimestamp(started_at, tz=timezone.utc) + timedelta(hours=1)).isoformat(), + } + } def _log_success(limiter, **kwargs):