fix(proxy): paginate reset_budget_windows scans

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

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
 ):