fix(proxy): use cursor pagination for reset budget scans

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

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