diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 6c769b0ff54..681e8439716 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -26,6 +26,8 @@ class ResetBudgetJob: self.proxy_logging_obj: ProxyLogging = proxy_logging_obj self.prisma_client: PrismaClient = prisma_client + _RESET_BUDGET_WINDOWS_BATCH_SIZE = 500 + async def reset_budget( self, ): @@ -634,26 +636,37 @@ class ResetBudgetJob: # --- Keys --- try: - all_keys = await self.prisma_client.db.litellm_verificationtoken.find_many() - 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] - ) + 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 except Exception as e: verbose_proxy_logger.exception( "Failed to reset budget windows for keys: %s", e @@ -661,26 +674,37 @@ class ResetBudgetJob: # --- Teams --- try: - all_teams = await self.prisma_client.db.litellm_teamtable.find_many() - 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] - ) + 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 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 17eb882ea21..a083e1826e6 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 @@ -36,9 +36,13 @@ class MockLiteLLMVerificationToken: def set_find_many_results(self, results: List[Any]): self._find_many_results = results - async def find_many(self, where: Dict[str, Any] | None = None) -> List[Any]: - self.find_many_calls.append({"where": where}) - return self._find_many_results + async def find_many(self, **kwargs) -> List[Any]: + self.find_many_calls.append(kwargs) + skip = kwargs.get("skip", 0) + take = kwargs.get("take") + if take is None: + return self._find_many_results[skip:] + return self._find_many_results[skip : skip + take] async def update(self, where: Dict[str, Any], data: Dict[str, Any]) -> Dict[str, Any]: self.update_calls.append({"where": where, "data": data}) @@ -60,9 +64,13 @@ class MockLiteLLMTeamTable: def set_find_many_results(self, results: List[Any]): self._find_many_results = results - async def find_many(self, where: Dict[str, Any] | None = None) -> List[Any]: - self.find_many_calls.append({"where": where}) - return self._find_many_results + async def find_many(self, **kwargs) -> List[Any]: + self.find_many_calls.append(kwargs) + skip = kwargs.get("skip", 0) + take = kwargs.get("take") + if take is None: + return self._find_many_results[skip:] + return self._find_many_results[skip : skip + take] async def update(self, where: Dict[str, Any], data: Dict[str, Any]) -> Dict[str, Any]: self.update_calls.append({"where": where, "data": data}) @@ -700,10 +708,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 == [ - {"where": None} + {"skip": 0, "take": 500, "order": {"token": "asc"}} ] assert mock_prisma_client.db.litellm_teamtable.find_many_calls == [ - {"where": None} + {"skip": 0, "take": 500, "order": {"team_id": "asc"}} ] key_updates = mock_prisma_client.db.litellm_verificationtoken.update_calls @@ -715,6 +723,46 @@ def test_reset_budget_windows_filters_null_budget_limits_in_python( assert team_updates[0]["where"] == {"team_id": "team-with-windows"} +def test_reset_budget_windows_paginates_large_tables( + 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) + ] + teams = [ + type( + "LiteLLM_TeamTable", + (), + {"team_id": f"team-{i}", "budget_limits": None}, + ) + for i in range(5) + ] + + 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()): + 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"}}, + ] + 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"}}, + ] + + def test_reset_budget_resets_endusers_with_null_budget_id( reset_budget_job, mock_prisma_client ):