From 0c07a26c8f7434de8d3f478ff6fb8a977e28796a Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Thu, 20 Aug 2026 13:45:54 -0400 Subject: [PATCH] fix(rate-limiting): fix routing-group bucket mismatch and unpinned token accounting task Two real findings from Veria AI and Bugbot, both independently caught by both bots: Success accounting checked a different bucket than admission (High/Low): resolve_any's dedup stamps resolved_group from whichever member frozenset(candidate_model_names) yields first, but success accounting for tokens/dollars only passed the one deployment that actually served as its sole candidate -- a trivial single-candidate dedup that resolves to that deployment's own name, which can differ from whichever member admission's full-group view picked. Success accounting now reconstructs the full routing-group candidate set via Router._get_routing_group_deployments, so it lands on the identical bucket admission checked regardless of which member actually served. Token/dollar accounting task not retained (Medium/Low): the same GC-before-running gap the previous commit fixed for concurrency release also applied to this hook's other fire-and-forget task -- token/dollar usage accounting, fired per cache partition with a bare asyncio.create_task and no strong reference. Renamed _BACKGROUND_RELEASE_TASKS to the more general _BACKGROUND_TASKS and wired this task through it too. Also investigated Bugbot's "concurrency TTL never refreshes" finding (TAG_RL_CHECK_AND_INCR_SCRIPT only sets EXPIRE when Redis reports TTL -1, so a bucket that already has a countdown running never gets it extended by a later reservation). Confirmed real and traces to the very first commit introducing this file, predating this session entirely. A correct fix needs the shared atomic check-and-increment script to distinguish concurrency's "extend the TTL on every new reservation" semantics from requests' "never extend, let the fixed window expire on schedule" semantics, since both units share this same script -- flagging as a follow-up rather than rushing a change to shared, security-sensitive admission logic. --- litellm/proxy/hooks/tag_rate_limiter.py | 52 ++++-- .../proxy/hooks/test_tag_rate_limiter.py | 148 ++++++++++++++++-- 2 files changed, 173 insertions(+), 27 deletions(-) diff --git a/litellm/proxy/hooks/tag_rate_limiter.py b/litellm/proxy/hooks/tag_rate_limiter.py index 2e41ba6271a..a20ca97baf4 100644 --- a/litellm/proxy/hooks/tag_rate_limiter.py +++ b/litellm/proxy/hooks/tag_rate_limiter.py @@ -79,14 +79,16 @@ _EMPTY_MAPPING: Final[Mapping[str, object]] = MappingProxyType({}) # function, to avoid a task disappearing mid-execution. The event loop only # keeps weak references to tasks. A task that isn't referenced elsewhere may # get garbage collected at any time, even before it's done." The success path -# below deliberately fires-and-forgets its release (unlike the failure/ -# disconnect paths, which await it directly) to keep the hot success-response -# path from waiting on a Redis round trip; by the time that background task -# would run, its keys have already been popped out of model_call_details, so -# a collected task's release is unrecoverable, not just delayed. Holding a -# strong reference here until the task's own completion callback discards it -# is the standard fix. -_BACKGROUND_RELEASE_TASKS: Final[set["asyncio.Task[None]"]] = set() # mutable-ok: see comment above +# deliberately fires-and-forgets its concurrency release and its token/dollar +# accounting increment (unlike the failure/disconnect paths, which await +# concurrency release directly) to keep the hot success-response path from +# waiting on a Redis round trip; by the time either background task would +# run, the state it needs (popped pending keys, or the request's own usage +# figures) is only available in that task's own closure, so a collected +# task's work is unrecoverable, not just delayed. Holding a strong reference +# here until each task's own completion callback discards it is the standard +# fix, shared by every fire-and-forget task this hook creates. +_BACKGROUND_TASKS: Final[set["asyncio.Task[None]"]] = set() # mutable-ok: see comment above # Single-key atomic check-and-increment. Deliberately one key per script call # (never a batch of differently-hash-tagged keys in one call): every tag_rl @@ -1214,8 +1216,8 @@ class _PROXY_TagRateLimiter( # pyright: ignore[reportUnusedClass] # only refer release_keys: Final = self._pop_pending_concurrency_keys(kwargs) if release_keys: release_task: Final = asyncio.create_task(self._release_keys(release_keys)) - _BACKGROUND_RELEASE_TASKS.add(release_task) # mutable-ok: see _BACKGROUND_RELEASE_TASKS's own docstring - release_task.add_done_callback(_BACKGROUND_RELEASE_TASKS.discard) + _BACKGROUND_TASKS.add(release_task) # mutable-ok: see _BACKGROUND_TASKS's own docstring + release_task.add_done_callback(_BACKGROUND_TASKS.discard) if self.llm_router is None: return @@ -1233,13 +1235,33 @@ class _PROXY_TagRateLimiter( # pyright: ignore[reportUnusedClass] # only refer key_hash: Final = standard_logging_metadata.get("user_api_key_hash") # model_group is the caller-visible name, which Router deliberately # keeps distinct from the serving deployment's own model_name for a - # routing-group call (see resolve_any's docstring); fall back to the - # one deployment that actually served this hop. + # routing-group call (see resolve_any's docstring). Passing only the + # one deployment that actually served this hop as the sole candidate + # would make resolve_any's dedup independently re-derive a + # *different* resolved_group than admission did whenever the group + # has more than one member: admission sees every member and picks + # whichever one frozenset(candidate_model_names) yields first for a + # shared signature, so success accounting must reconstruct that same + # full candidate set to land on the identical bucket, not just + # whichever deployment happened to serve -- otherwise a token/dollar + # limit is checked against one bucket at admission and accounted + # against a different one on success, letting usage silently bypass + # the configured limit. Falls back to the serving deployment alone + # only when `model_group` isn't a routing group at all (a plain + # single-model_name chain, where resolve() already matches directly + # and this candidate set is never actually consulted). deployment_id: Final = standard_logging_object.get("model_id") serving_deployment: Final = ( self.llm_router.get_deployment(deployment_id) if isinstance(deployment_id, str) else None ) - candidate_model_names: Final = (serving_deployment.model_name,) if serving_deployment is not None else () + routing_group_deployments: Final = self.llm_router._get_routing_group_deployments( # pyright: ignore[reportPrivateUsage] # reused across module boundaries, matching resolve_any's own reliance on this method + model=model_group, team_id=team_id + ) + candidate_model_names: Final = ( + tuple(dep["model_name"] for dep in routing_group_deployments) + if routing_group_deployments is not None + else ((serving_deployment.model_name,) if serving_deployment is not None else ()) + ) configured: Final = self._index.get(self.llm_router).resolve_any(model_group, team_id, candidate_model_names) if not configured: return @@ -1293,9 +1315,11 @@ class _PROXY_TagRateLimiter( # pyright: ignore[reportUnusedClass] # only refer parent_otel_span: Final = _get_parent_otel_span_from_kwargs(kwargs) for partition_key, group_operations in operations_by_partition.items(): partition = await self._partition_for(partition_key) # not Final: rebound each loop iteration - asyncio.create_task( + accounting_task = asyncio.create_task( # not Final: rebound each loop iteration partition.v3.async_increment_tokens_with_ttl_preservation( pipeline_operations=tuple(group_operations), parent_otel_span=parent_otel_span, ) ) + _BACKGROUND_TASKS.add(accounting_task) # mutable-ok: see _BACKGROUND_TASKS's own docstring + accounting_task.add_done_callback(_BACKGROUND_TASKS.discard) diff --git a/tests/test_litellm/proxy/hooks/test_tag_rate_limiter.py b/tests/test_litellm/proxy/hooks/test_tag_rate_limiter.py index 1e869ab71f9..dadc6aedd0a 100644 --- a/tests/test_litellm/proxy/hooks/test_tag_rate_limiter.py +++ b/tests/test_litellm/proxy/hooks/test_tag_rate_limiter.py @@ -24,14 +24,14 @@ from litellm.proxy.hooks.tag_rate_limiter import ( _extract_key_hash, _extract_team_id, _fixed_length_identity, - _BACKGROUND_RELEASE_TASKS, + _BACKGROUND_TASKS, _inflight_key, _partition_key, _PENDING_CONCURRENCY_KEYS_FIELD, _PROXY_TagRateLimiter, _queue_pending_concurrency_reservations, ) -from litellm.types.router import TagRateLimitEntry +from litellm.types.router import RoutingGroup, TagRateLimitEntry class TimeController: @@ -871,6 +871,74 @@ async def test_log_success_event_falls_back_to_serving_deployment_model_name_for ) +@pytest.mark.asyncio +async def test_log_success_event_accounts_against_the_same_bucket_admission_checked(time_controller): + """ + resolve_any dedups an identical signature across a routing group's + members into one shared entry, stamped with resolved_group from + whichever member frozenset(candidate_model_names) yields first (see + resolve_any's own docstring). Success accounting for tokens/dollars only + learns the one deployment that actually served this hop; passing just + that single name as resolve_any's sole candidate would make its dedup + trivially resolve to that deployment's own name -- which can differ from + whichever member admission's full-group view picked, silently + accounting usage against a bucket admission never checked and letting a + token/dollar limit be bypassed. Success accounting must reconstruct the + full routing-group candidate set so it lands on the identical bucket + regardless of which member actually served. + """ + token_limits = { + "token_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 500000, "period_seconds": 86400}] + } + } + router = litellm.Router( + model_list=[ + _deployment("backend-a", "dep-a", token_limits), + _deployment("backend-b", "dep-b", token_limits), + ], + routing_groups=[ + RoutingGroup(group_name="my-group", models=["backend-a", "backend-b"], routing_strategy="simple-shuffle") + ], + ) + limiter = _make_limiter(time_controller) + limiter.update_variables(llm_router=router) + + # What admission would check: it sees every member, and resolve_any's + # dedup picks whichever one frozenset yields first for the shared entry. + admitted = limiter._index.get(router).resolve_any( + "my-group", team_id=None, candidate_model_names=("backend-a", "backend-b") + ) + assert len(admitted) == 1 + admission_bucket_group = admitted[0].resolved_group + + # Force the deployment that actually serves to be the *other* member -- + # deterministic regardless of which one frozenset happened to pick above, + # so this test always exercises the mismatch the fix guards against. + serving_model_name = "backend-b" if admission_bucket_group == "backend-a" else "backend-a" + serving_deployment_id = "dep-b" if serving_model_name == "backend-b" else "dep-a" + + kwargs = { + "metadata": {"tags": ["end_user_id:u1"]}, + "standard_logging_object": { + "model_group": "my-group", + "model_id": serving_deployment_id, + "total_tokens": 42, + "response_cost": 0.01, + }, + } + await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + now = time_controller.now().timestamp() + token_key = _expected_bucket_key( + "my-group", "tokens", "daily", "end_user_id", "u1", 86400, now, resolved_group=admission_bucket_group + ) + assert ( + float(await limiter.internal_usage_cache.async_get_cache(key=token_key, litellm_parent_otel_span=None)) == 42.0 + ) + + # --------------------------------------------------------------------------- # concurrency limits -- reserve at admission, release on success/failure # --------------------------------------------------------------------------- @@ -1125,39 +1193,39 @@ async def test_background_release_tasks_registry_holds_a_reference_until_done(): one with no other referrer can be garbage collected before it runs -- and by the time it would run here, its keys are already popped out of model_call_details, so a collected task's release is unrecoverable, not - merely delayed. _BACKGROUND_RELEASE_TASKS exists to hold a strong + merely delayed. _BACKGROUND_TASKS exists to hold a strong reference for exactly as long as the task is pending, then release it via the task's own done-callback -- exercised directly here (an Event gate gives a deterministic pending window; going through the real async_log_success_event doesn't, since its own further awaits let a fast in-memory release resolve before a test could ever observe it pending). """ - assert len(_BACKGROUND_RELEASE_TASKS) == 0 + assert len(_BACKGROUND_TASKS) == 0 gate = asyncio.Event() async def _pending_release(): await gate.wait() task = asyncio.create_task(_pending_release()) - _BACKGROUND_RELEASE_TASKS.add(task) - task.add_done_callback(_BACKGROUND_RELEASE_TASKS.discard) + _BACKGROUND_TASKS.add(task) + task.add_done_callback(_BACKGROUND_TASKS.discard) - assert task in _BACKGROUND_RELEASE_TASKS + assert task in _BACKGROUND_TASKS gate.set() await task # The done-callback removes it -- the registry doesn't grow unbounded # across requests. - assert task not in _BACKGROUND_RELEASE_TASKS - assert len(_BACKGROUND_RELEASE_TASKS) == 0 + assert task not in _BACKGROUND_TASKS + assert len(_BACKGROUND_TASKS) == 0 @pytest.mark.asyncio async def test_success_event_release_is_wired_through_the_background_registry(time_controller): """ End-to-end check that async_log_success_event's fire-and-forget release - is genuinely wired through _BACKGROUND_RELEASE_TASKS, not a bare + is genuinely wired through _BACKGROUND_TASKS, not a bare unreferenced asyncio.create_task -- the registry must be empty again once the (fast, in-memory) release has had a chance to run, and the release itself must have actually happened. @@ -1181,11 +1249,11 @@ async def test_success_event_release_is_wired_through_the_background_registry(ti await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) # The registry was actually populated: proves the release ran through - # _BACKGROUND_RELEASE_TASKS, not a bare unreferenced asyncio.create_task + # _BACKGROUND_TASKS, not a bare unreferenced asyncio.create_task # (which would never touch this set at all, and an "empty at the end" # check alone can't tell the two apart -- an empty registry throughout # would satisfy that just as well as one that filled and drained). - assert len(_BACKGROUND_RELEASE_TASKS) == 1 + assert len(_BACKGROUND_TASKS) == 1 # Two ticks: one for the release task itself to finish (it may already be # done by the time async_log_success_event returns, given that method's @@ -1194,7 +1262,7 @@ async def test_success_event_release_is_wired_through_the_background_registry(ti await asyncio.sleep(0) await asyncio.sleep(0) - assert len(_BACKGROUND_RELEASE_TASKS) == 0 + assert len(_BACKGROUND_TASKS) == 0 result = await limiter.async_filter_deployments( model="grp", @@ -1205,6 +1273,60 @@ async def test_success_event_release_is_wired_through_the_background_registry(ti assert result == healthy +@pytest.mark.asyncio +async def test_success_event_token_accounting_is_wired_through_the_background_registry(time_controller): + """ + Same gap as the concurrency release above, in a second fire-and-forget + task on the same success path: token/dollar accounting is also fired + via a bare asyncio.create_task per cache partition, with no strong + reference of its own. A collected task here drops a usage increment + that can never be recovered (the figures it needed only exist in that + task's own closure), silently under-counting a caller's token/dollar + usage against its configured limit. Must be tracked the same way. + """ + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "token_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 500000, "period_seconds": 86400}] + } + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + + kwargs = { + "metadata": {"tags": ["end_user_id:u1"]}, + "standard_logging_object": { + "model_group": "grp", + "model_id": "dep-1", + "total_tokens": 42, + "response_cost": 0.01, + }, + } + await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + + # The registry was actually populated: proves accounting ran through + # _BACKGROUND_TASKS, not a bare unreferenced asyncio.create_task. + assert len(_BACKGROUND_TASKS) == 1 + + await asyncio.sleep(0) + await asyncio.sleep(0) + + assert len(_BACKGROUND_TASKS) == 0 + + now = time_controller.now().timestamp() + token_key = _expected_bucket_key("grp", "tokens", "daily", "end_user_id", "u1", 86400, now) + assert ( + float(await limiter.internal_usage_cache.async_get_cache(key=token_key, litellm_parent_otel_span=None)) == 42.0 + ) + + @pytest.mark.asyncio async def test_concurrency_slot_released_on_disconnect_frees_capacity(time_controller): """