diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 681e8439716..49235afdb6d 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -28,6 +28,30 @@ class ResetBudgetJob: _RESET_BUDGET_WINDOWS_BATCH_SIZE = 500 + async def _iter_budget_window_rows(self, table: Any, id_field: str): + cursor_value = None + + while True: + query_kwargs = { + "take": self._RESET_BUDGET_WINDOWS_BATCH_SIZE, + "order": {id_field: "asc"}, + } + if cursor_value is not None: + query_kwargs["cursor"] = {id_field: cursor_value} + query_kwargs["skip"] = 1 + + rows = await table.find_many(**query_kwargs) + if len(rows) == 0: + break + + for row in rows: + yield row + + if len(rows) < self._RESET_BUDGET_WINDOWS_BATCH_SIZE: + break + + cursor_value = getattr(rows[-1], id_field) + async def reset_budget( self, ): @@ -636,37 +660,27 @@ class ResetBudgetJob: # --- Keys --- try: - skip = 0 - while True: - all_keys = await self.prisma_client.db.litellm_verificationtoken.find_many( - skip=skip, - take=self._RESET_BUDGET_WINDOWS_BATCH_SIZE, - order={"token": "asc"}, - ) - if len(all_keys) == 0: - break - for key in all_keys: - raw = key.budget_limits # type: ignore[attr-defined] - if not raw: - continue - windows: list = raw if isinstance(raw, list) else json.loads(raw) - changed = False - for window in windows: - counter_key = ( - f"spend:key:{key.token}:window:{window['budget_duration']}" - ) - if await ResetBudgetJob._reset_expired_window( - window, counter_key, spend_counter_cache, now - ): - changed = True - if changed: - await self.prisma_client.db.litellm_verificationtoken.update( - where={"token": key.token}, - data={"budget_limits": json.dumps(windows)}, # type: ignore[arg-type] - ) - if len(all_keys) < self._RESET_BUDGET_WINDOWS_BATCH_SIZE: - break - skip += self._RESET_BUDGET_WINDOWS_BATCH_SIZE + async for key in self._iter_budget_window_rows( + self.prisma_client.db.litellm_verificationtoken, "token" + ): + raw = key.budget_limits # type: ignore[attr-defined] + if not raw: + continue + windows: list = raw if isinstance(raw, list) else json.loads(raw) + changed = False + for window in windows: + counter_key = ( + f"spend:key:{key.token}:window:{window['budget_duration']}" + ) + if await ResetBudgetJob._reset_expired_window( + window, counter_key, spend_counter_cache, now + ): + changed = True + if changed: + await self.prisma_client.db.litellm_verificationtoken.update( + where={"token": key.token}, + data={"budget_limits": json.dumps(windows)}, # type: ignore[arg-type] + ) except Exception as e: verbose_proxy_logger.exception( "Failed to reset budget windows for keys: %s", e @@ -674,37 +688,27 @@ class ResetBudgetJob: # --- Teams --- try: - skip = 0 - while True: - all_teams = await self.prisma_client.db.litellm_teamtable.find_many( - skip=skip, - take=self._RESET_BUDGET_WINDOWS_BATCH_SIZE, - order={"team_id": "asc"}, - ) - if len(all_teams) == 0: - break - for team in all_teams: - raw = team.budget_limits # type: ignore[attr-defined] - if not raw: - continue - windows = raw if isinstance(raw, list) else json.loads(raw) - changed = False - for window in windows: - counter_key = ( - f"spend:team:{team.team_id}:window:{window['budget_duration']}" - ) - if await ResetBudgetJob._reset_expired_window( - window, counter_key, spend_counter_cache, now - ): - changed = True - if changed: - await self.prisma_client.db.litellm_teamtable.update( - where={"team_id": team.team_id}, - data={"budget_limits": json.dumps(windows)}, # type: ignore[arg-type] - ) - if len(all_teams) < self._RESET_BUDGET_WINDOWS_BATCH_SIZE: - break - skip += self._RESET_BUDGET_WINDOWS_BATCH_SIZE + async for team in self._iter_budget_window_rows( + self.prisma_client.db.litellm_teamtable, "team_id" + ): + raw = team.budget_limits # type: ignore[attr-defined] + if not raw: + continue + windows = raw if isinstance(raw, list) else json.loads(raw) + changed = False + for window in windows: + counter_key = ( + f"spend:team:{team.team_id}:window:{window['budget_duration']}" + ) + if await ResetBudgetJob._reset_expired_window( + window, counter_key, spend_counter_cache, now + ): + changed = True + if changed: + await self.prisma_client.db.litellm_teamtable.update( + where={"team_id": team.team_id}, + data={"budget_limits": json.dumps(windows)}, # type: ignore[arg-type] + ) except Exception as e: verbose_proxy_logger.exception( "Failed to reset budget windows for teams: %s", e diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index a083e1826e6..f3dc506b5ce 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -38,11 +38,21 @@ class MockLiteLLMVerificationToken: async def find_many(self, **kwargs) -> List[Any]: self.find_many_calls.append(kwargs) - skip = kwargs.get("skip", 0) take = kwargs.get("take") + cursor = kwargs.get("cursor") + skip = kwargs.get("skip", 0) + results = self._find_many_results + if cursor is not None: + token = cursor.get("token") + start = next( + index for index, item in enumerate(results) if item.token == token + ) + results = results[start + skip :] + elif skip: + results = results[skip:] if take is None: - return self._find_many_results[skip:] - return self._find_many_results[skip : skip + take] + return results + return results[:take] async def update(self, where: Dict[str, Any], data: Dict[str, Any]) -> Dict[str, Any]: self.update_calls.append({"where": where, "data": data}) @@ -66,11 +76,21 @@ class MockLiteLLMTeamTable: async def find_many(self, **kwargs) -> List[Any]: self.find_many_calls.append(kwargs) - skip = kwargs.get("skip", 0) take = kwargs.get("take") + cursor = kwargs.get("cursor") + skip = kwargs.get("skip", 0) + results = self._find_many_results + if cursor is not None: + team_id = cursor.get("team_id") + start = next( + index for index, item in enumerate(results) if item.team_id == team_id + ) + results = results[start + skip :] + elif skip: + results = results[skip:] if take is None: - return self._find_many_results[skip:] - return self._find_many_results[skip : skip + take] + return results + return results[:take] async def update(self, where: Dict[str, Any], data: Dict[str, Any]) -> Dict[str, Any]: self.update_calls.append({"where": where, "data": data}) @@ -190,6 +210,10 @@ async def run_async_test(coro): return await coro +async def _collect_async_iter(async_iterable): + return [item async for item in async_iterable] + + # Tests def test_reset_budget_for_key(reset_budget_job, mock_prisma_client): # Setup test data with timezone-aware datetime @@ -708,10 +732,10 @@ def test_reset_budget_windows_filters_null_budget_limits_in_python( asyncio.run(reset_budget_job.reset_budget_windows()) assert mock_prisma_client.db.litellm_verificationtoken.find_many_calls == [ - {"skip": 0, "take": 500, "order": {"token": "asc"}} + {"take": 500, "order": {"token": "asc"}} ] assert mock_prisma_client.db.litellm_teamtable.find_many_calls == [ - {"skip": 0, "take": 500, "order": {"team_id": "asc"}} + {"take": 500, "order": {"team_id": "asc"}} ] key_updates = mock_prisma_client.db.litellm_verificationtoken.update_calls @@ -752,17 +776,104 @@ def test_reset_budget_windows_paginates_large_tables( asyncio.run(reset_budget_job.reset_budget_windows()) assert mock_prisma_client.db.litellm_verificationtoken.find_many_calls == [ - {"skip": 0, "take": 2, "order": {"token": "asc"}}, - {"skip": 2, "take": 2, "order": {"token": "asc"}}, - {"skip": 4, "take": 2, "order": {"token": "asc"}}, + {"take": 2, "order": {"token": "asc"}}, + {"take": 2, "order": {"token": "asc"}, "cursor": {"token": "sk-1"}, "skip": 1}, + {"take": 2, "order": {"token": "asc"}, "cursor": {"token": "sk-3"}, "skip": 1}, ] assert mock_prisma_client.db.litellm_teamtable.find_many_calls == [ - {"skip": 0, "take": 2, "order": {"team_id": "asc"}}, - {"skip": 2, "take": 2, "order": {"team_id": "asc"}}, - {"skip": 4, "take": 2, "order": {"team_id": "asc"}}, + {"take": 2, "order": {"team_id": "asc"}}, + { + "take": 2, + "order": {"team_id": "asc"}, + "cursor": {"team_id": "team-1"}, + "skip": 1, + }, + { + "take": 2, + "order": {"team_id": "asc"}, + "cursor": {"team_id": "team-3"}, + "skip": 1, + }, ] +def test_iter_budget_window_rows_uses_cursor_pagination(reset_budget_job, mock_prisma_client): + reset_budget_job._RESET_BUDGET_WINDOWS_BATCH_SIZE = 2 + + keys = [ + type( + "LiteLLM_VerificationToken", + (), + {"token": f"sk-{i}", "budget_limits": None}, + ) + for i in range(5) + ] + mock_prisma_client.db.litellm_verificationtoken.set_find_many_results(keys) + + rows = asyncio.run( + _collect_async_iter( + reset_budget_job._iter_budget_window_rows( + mock_prisma_client.db.litellm_verificationtoken, "token" + ) + ) + ) + + assert [row.token for row in rows] == [f"sk-{i}" for i in range(5)] + assert mock_prisma_client.db.litellm_verificationtoken.find_many_calls == [ + {"take": 2, "order": {"token": "asc"}}, + {"take": 2, "order": {"token": "asc"}, "cursor": {"token": "sk-1"}, "skip": 1}, + {"take": 2, "order": {"token": "asc"}, "cursor": {"token": "sk-3"}, "skip": 1}, + ] + + +def test_reset_budget_windows_updates_items_across_cursor_pages( + reset_budget_job, mock_prisma_client +): + reset_budget_job._RESET_BUDGET_WINDOWS_BATCH_SIZE = 1 + + keys = [ + type( + "LiteLLM_VerificationToken", + (), + { + "token": f"sk-{i}", + "budget_limits": '[{"budget_duration": "24h", "reset_at": null}]', + }, + ) + for i in range(2) + ] + teams = [ + type( + "LiteLLM_TeamTable", + (), + { + "team_id": f"team-{i}", + "budget_limits": [{"budget_duration": "7d", "reset_at": None}], + }, + ) + for i in range(2) + ] + + mock_prisma_client.db.litellm_verificationtoken.set_find_many_results(keys) + mock_prisma_client.db.litellm_teamtable.set_find_many_results(teams) + + with patch( + "litellm.proxy.proxy_server.spend_counter_cache", new=MagicMock() + ), patch.object( + ResetBudgetJob, + "_reset_expired_window", + new=AsyncMock(return_value=True), + ): + asyncio.run(reset_budget_job.reset_budget_windows()) + + assert [ + call["where"] for call in mock_prisma_client.db.litellm_verificationtoken.update_calls + ] == [{"token": "sk-0"}, {"token": "sk-1"}] + assert [ + call["where"] for call in mock_prisma_client.db.litellm_teamtable.update_calls + ] == [{"team_id": "team-0"}, {"team_id": "team-1"}] + + def test_reset_budget_resets_endusers_with_null_budget_id( reset_budget_job, mock_prisma_client ):