mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
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].
This commit is contained in:
parent
93b5dd7e9f
commit
311bf593f4
2 changed files with 43 additions and 4 deletions
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue