From 311bf593f4bb7e5c87342b42e4609f92a68f0caf Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Tue, 25 Aug 2026 12:31:04 -0400 Subject: [PATCH] fix(rate-limiting): fold policy scoping fields into the partition key _partition_key only distinguished entries by tag_id/name/limit/period_seconds/ scope_by_key_hash/max_in_memory_cache_size, but two entries can share all of those while disagreeing on enabled_for/disabled_for/apply_to_key_alias (the same class of collision _DedupSignature and _policy_fingerprint already guard against for dedup and bucket keys). Two such entries setting the same max_in_memory_cache_size to get their own dedicated partition collided onto one shared partition instead, letting one entry's high-cardinality traffic evict the other's active counters. Folds _policy_fingerprint into the partition key, keeping max_in_memory_cache_size as the trailing element since _partition_for reads it via partition_key[-1]. --- .../hooks/model_based_tag_rate_limits_hook.py | 16 +++++++--- .../test_model_based_tag_rate_limits_hook.py | 31 +++++++++++++++++++ 2 files changed, 43 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py index 2149e254e9a..fcb4ce98502 100644 --- a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py @@ -883,7 +883,17 @@ def _resolve_max_in_memory_cache_size() -> int | None: # always resolves to the same signature across index rebuilds, which is what # keeps _PROXY_ModelBasedTagRateLimitsHook._partitions from leaking a fresh partition # every time _TagRateLimitIndex rebuilds and reconstructs `_ConfiguredLimit`s. -_PartitionKey: TypeAlias = tuple[str, str, float, int, bool, int] | None +# The str before the trailing int is `_policy_fingerprint(entry)`: two +# entries can share tag_id/name (see +# test_bucket_key_differs_for_same_named_entries_with_different_scoping_only) +# while disagreeing on limit/period_seconds/scope_by_key_hash/enabled_for/ +# disabled_for/apply_to_key_alias -- _DedupSignature already treats that as +# two distinct policies, so a shared max_in_memory_cache_size must not route +# them onto the same partition either, or one entry's high-cardinality +# traffic can evict the other's active counters from a cache neither entry +# asked to share. max_in_memory_cache_size stays the trailing element: +# `partition_key[-1]` reads it directly to size the partition's cache. +_PartitionKey: TypeAlias = tuple[str, str, str, int] | None # Grouping type for async_log_success_event's per-partition tokens/dollars # pipeline dispatch -- named only so the declaration fits on one line; see # that method for why the grouping is needed. @@ -896,9 +906,7 @@ def _partition_key(entry: TagRateLimitEntry) -> _PartitionKey: return ( entry.tag_id, entry.name, - entry.limit, - entry.period_seconds, - entry.scope_by_key_hash, + _policy_fingerprint(entry), entry.max_in_memory_cache_size, ) diff --git a/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py b/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py index 1ab22959358..5d924add74e 100644 --- a/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py +++ b/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py @@ -3748,6 +3748,37 @@ def test_partition_key_distinguishes_entries_that_differ_only_by_scope_by_key_ha assert _partition_key(unscoped) != _partition_key(scoped) +def test_partition_key_distinguishes_entries_that_differ_only_by_scoping_fields(): + """ + A plain, unscoped entry and a scoped override can legitimately share + name/tag_id/limit/period_seconds/scope_by_key_hash (see + test_bucket_key_differs_for_same_named_entries_with_different_scoping_only) + while disagreeing on enabled_for/disabled_for/apply_to_key_alias -- + _DedupSignature and _policy_fingerprint already treat that as two + distinct policies, so a shared max_in_memory_cache_size must not route + them onto the same in-memory partition either, or one entry's + high-cardinality traffic can evict the other's active counters from a + cache neither entry asked to share. + """ + base_kwargs = {"name": "daily", "tag_id": "end_user_id", "limit": 100, "period_seconds": 86400, "max_in_memory_cache_size": 50} + unscoped = TagRateLimitEntry(**base_kwargs) + enabled_for_scoped = TagRateLimitEntry( + **base_kwargs, enabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",)) + ) + disabled_for_scoped = TagRateLimitEntry( + **base_kwargs, disabled_for=TagRateLimitScope(tag_id="company_id", values=("1032",)) + ) + alias_scoped = TagRateLimitEntry(**base_kwargs, apply_to_key_alias=("premium-key",)) + + keys = { + _partition_key(unscoped), + _partition_key(enabled_for_scoped), + _partition_key(disabled_for_scoped), + _partition_key(alias_scoped), + } + assert len(keys) == 4 + + @pytest.mark.asyncio async def test_scope_by_key_hash_composes_with_max_in_memory_cache_size_and_key_ttl_seconds_overrides(time_controller): """