refactor(proxy): trim the invalidation docstrings and inject the page read failure

Cuts the new docstrings back to the parts a reader cannot get from the code,
and fixes a stale reference: the walk this one is modelled on is
_reset_windows_for, not _reset_windows_for_source.

The truncation test reached in and replaced MockTable.find_many. The mock takes
a scheduled read failure instead, the way it already takes canned rows.

litellm/caching/dual_cache.py

@@ -522,13 +522,8 @@ class DualCache(BaseCache):
             await self.redis_cache.async_delete_cache(key)
 
     async def async_delete_cache_keys(self, keys: Sequence[str]) -> None:
-        """Batch twin of ``async_delete_cache``: one Redis round trip per chunk
-        instead of one per key.
-
-        Chunked because Redis takes the whole list as a single DELETE command,
-        and a caller holding a population-sized list would otherwise build one
-        command out of it.
-        """
+        """Batch twin of ``async_delete_cache``, chunked because Redis takes the
+        whole list as one DELETE command."""
         if not keys:
             return
         for key in keys:

litellm/proxy/common_utils/reset_budget_job.py

@@ -202,10 +202,8 @@ def _budget_link_where(
 def _enduser_invalidation_where(budget_ids: Sequence[str]) -> dict[str, object]:
     """Customers whose cached spend a committed reset of these tiers invalidated.
 
-    Mirrors ``_queue_enduser_resets``: the link, plus the NULL-budget_id rows
-    that ride the default tier when that tier is one of the expiring ones. The
-    write's ``spend > 0`` filter has no twin here because the commit already
-    zeroed those rows, so post-commit it would match nobody.
+    Mirrors ``_queue_enduser_resets`` without its ``spend > 0`` filter, which
+    post-commit would match nobody.
     """
     linked: Final = _budget_link_where(budget_ids)
     default_budget_id: Final = litellm.max_end_user_budget_id
@@ -279,9 +277,8 @@ class _BudgetCascade:
 
 @dataclass(frozen=True, slots=True)
 class _EndUserWalk:
-    """Where the post-commit customer walk stands: the keyset cursor its next
-    page resumes from, None once there is no next page, how many customers it
-    has reached, and whether a failed page read cut it short of the tail."""
+    """Where the customer walk stands. ``cursor`` is None once it is done, and
+    ``truncated`` says a failed page read cut it short of the tail."""
 
     cursor: str | None = ""
     invalidated: int = 0
@@ -306,8 +303,6 @@ class _BudgetCascadeFailed:
 
 _EMPTY_CASCADE: Final = _BudgetCascade()
 
-#: Which of the proxy's two caches a batch of keys belongs to. ``spend_counter_cache``
-#: holds the live running spend; ``user_api_key_cache`` holds the cached management rows.
 _InvalidatedCache = Literal["spend counter", "user_api_key_cache"]
 
 
@@ -623,20 +618,15 @@ class ResetBudgetJob:
     @staticmethod
     async def _invalidate_caches(counter_keys: Sequence[str], cache_keys: Sequence[str]) -> None:
         """Batch twin of ``_invalidate_spend_counter`` and
-        ``_invalidate_user_api_key_cache_entry``, carrying the same
-        after-the-commit requirement as both.
-
-        One round trip per chunk rather than one per key: a tier's dependent
-        population is unbounded, and awaiting each key in turn makes the last
-        dependent wait out every dependent ahead of it.
-        """
+        ``_invalidate_user_api_key_cache_entry``, after the commit like both:
+        one round trip per chunk where a tier's dependents are unbounded."""
         await ResetBudgetJob._invalidate_cache("spend counter", counter_keys)
         await ResetBudgetJob._invalidate_cache("user_api_key_cache", cache_keys)
 
     @staticmethod
     async def _invalidate_cache(cache: _InvalidatedCache, keys: Sequence[str]) -> None:
-        """One cache's share of a batch, awaited separately from the other's so a
-        failure against either still leaves the other one invalidated."""
+        """One cache's share of a batch, awaited separately so either failing
+        still leaves the other invalidated."""
         if not keys:
             return
         try:
@@ -679,21 +669,9 @@ class ResetBudgetJob:
     async def _invalidate_enduser_caches(self, budget_ids: Sequence[str]) -> _EndUserWalk:
         """Drop the cached spend of every customer the committed tier reset zeroed.
 
-        Walked a page at a time with a keyset cursor, for the same reason
-        ``_reset_windows_for_source`` is: the customers sharing one tier are
-        unbounded, so reading them into one result set puts a
-        customer-count-sized list in the proxy's heap on every tick, and a
-        deployment large enough turns that into an OOM rather than a slow tick.
-
-        No per-run page cap, also for that walk's reason: the position cannot
-        survive the run, so a cap would restart at the first customer every tick
-        and never reach the tail. The cursor strictly advances, so this
-        terminates on its own.
-
-        A page that fails to read stops the walk short of the tail. The window is
-        already advanced by then, so no later tick comes back for the customers
-        past it, which is why the walk reports that it was cut short instead of
-        passing the part it managed off as the whole.
+        Paged like ``_reset_windows_for``, and capless for its reason too: the
+        customers on one tier are unbounded, and a cap cannot keep its position
+        across pod elections, so it would restart at the first customer forever.
         """
         if not budget_ids:
             return _ENDUSER_WALK_DONE

litellm/proxy/common_utils/user_api_key_cache.py

@@ -223,13 +223,11 @@ class UserApiKeyCache(DualCache):
         await super().async_delete_cache(key)
 
     async def async_delete_cache_keys(self, keys: Sequence[str]) -> None:
-        """Batch twin of ``async_delete_cache``, partitioned the way
-        ``async_set_cache_pipeline`` partitions its writes.
+        """Batch twin of ``async_delete_cache``, partitioned like
+        ``async_set_cache_pipeline``.
 
-        Both partitions are cleared even when one of them raises: a caller
-        batching these has already committed the rows they cache, so a partition
-        left holding pre-reset spend goes on being authorized against until the
-        entry expires. The first failure is re-raised for the caller to report.
+        Both partitions are cleared even when one raises, because a caller
+        batching these has already committed the rows they cache.
         """
         key_object_keys: Final = tuple(key for key in keys if is_user_key_cache_key(key))
         other_keys: Final = tuple(key for key in keys if not is_user_key_cache_key(key))

tests/test_litellm/proxy/common_utils/test_reset_budget_job.py

@@ -32,10 +32,16 @@ class MockTable:
         self.find_many_calls: List[Dict[str, Any]] = []
         self.update_many_calls: List[Dict[str, Any]] = []
         self._find_many_results: List[Any] = []
+        self._find_many_error: Optional[tuple[int, Exception]] = None
 
     def set_find_many_results(self, results: List[Any]):
         self._find_many_results = results
 
+    def set_find_many_error(self, after_reads: int, error: Exception):
+        """Fail every read past the first ``after_reads``, the way a connection
+        dropping partway through a paged walk does."""
+        self._find_many_error = (after_reads, error)
+
     async def find_many(
         self,
         where: Dict[str, Any],
@@ -45,6 +51,8 @@ class MockTable:
         """Replays canned rows, honouring the keyset cursor + ``take`` a paged
         caller relies on: without that a paged walk never advances and the
         test would hang instead of failing."""
+        if self._find_many_error is not None and len(self.find_many_calls) >= self._find_many_error[0]:
+            raise self._find_many_error[1]
         paging = {k: v for k, v in (("order", order), ("take", take)) if v is not None}
         self.find_many_calls.append({"where": where, **paging})
         rows = list(self._find_many_results)
@@ -1686,14 +1694,7 @@ def test_enduser_invalidation_reports_a_page_read_failure_instead_of_a_clean_fin
             for i in range(RESET_BUDGET_JOB_BATCH_SIZE + 3)
         ]
     )
-    read_page: Final = endusers.find_many
-
-    async def fail_after_the_first_page(**kwargs):
-        if endusers.find_many_calls:
-            raise RuntimeError("connection reset while paging customers")
-        return await read_page(**kwargs)
-
-    endusers.find_many = fail_after_the_first_page
+    endusers.set_find_many_error(1, RuntimeError("connection reset while paging customers"))
     logging_obj: Final = RecordingProxyLogging()
     job: Final = ResetBudgetJob(proxy_logging_obj=logging_obj, prisma_client=mock_prisma_client)