mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(cache_warming): derive a record's partition from the caller's full identity
Anything stored under a partition is also attributed to that caller, so the partition has to cover every dimension the attribution does. It did not. The attribution carried seven identity fields while the partition returned the first of four, and the gap between those two lists was a cross-tenant collision channel: principals differing only in an omitted dimension shared one record, and whoever wrote last owned the payload and the attribution while state accumulated under it stayed. End users are the reachable case, since they share a virtual key by design and a key hash cannot tell them apart. One end user could seed a session's touched-model set, the next end user's turn would overwrite the payload and attribution while inheriting those models, and the replays then sent that caller's prompt to a model they had never used and billed it to their end-user budget. Projects were the same bug, unreported. The partition is now composed from every dimension rather than the strongest one, and both it and the attribution read the same list, taken from the type that declares those fields rather than spelled out in either place, so they cannot drift apart again. That also removes the need for a rule about state surviving an attribution change: a different attribution is a different partition, so there is no shared state to clear. Values are hashed rather than concatenated. End-user and project identifiers are customer-chosen strings that would otherwise sit in Redis key names and in every log line quoting one, and hashing identifiers before they become cache keys is what DeploymentAffinityCheck already does. Operational note: this changes partition strings, so existing cache-warming records and complexity-router session-affinity pins are not read after deploy. Both are caches with TTLs; pins re-derive on the next turn and records re-capture.
This commit is contained in:
parent
18dc38f9f6
commit
4ff2ecd4c8
3 changed files with 59 additions and 37 deletions
|
|
@ -1,7 +1,10 @@
|
|||
# What is this?
|
||||
## Helper utilities
|
||||
import copy
|
||||
import hashlib
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from functools import lru_cache
|
||||
from typing import TYPE_CHECKING, Any, Iterable, List, Literal, Optional, Union
|
||||
|
||||
import httpx
|
||||
|
|
@ -229,25 +232,43 @@ def iter_request_metadata_dicts(request_kwargs: Mapping[str, object]) -> tuple[M
|
|||
)
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def proxy_identity_fields() -> "tuple[str, ...]":
|
||||
"""Every tenant dimension the proxy stamps, read off the type that defines them rather than listed here,
|
||||
so a dimension litellm adds is picked up instead of silently omitted."""
|
||||
from litellm.types.utils import StandardLoggingUserAPIKeyMetadata
|
||||
|
||||
return ("user_api_key_hash",) + tuple(
|
||||
sorted(
|
||||
field
|
||||
for field in StandardLoggingUserAPIKeyMetadata.__annotations__
|
||||
if field.startswith("user_api_key_") and field.endswith("_id")
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def get_caller_scope(request_kwargs: Mapping[str, object]) -> str:
|
||||
"""A cache/affinity partition key for the authenticated caller, scoped by the strongest identity the
|
||||
request actually carries. Virtual keys carry a key hash; JWT and other keyless proxy principals do not,
|
||||
so they scope by the tenancy the proxy stamped instead, mirroring how the v3 rate limiter builds an
|
||||
api_key descriptor only when api_key is present and separate descriptors per user, team and organization
|
||||
(parallel_request_limiter_v3.py:1988). Collapsing keyless callers onto one shared literal would merge
|
||||
distinct tenants that reuse a session id into a single partition. "unscoped" is returned only when the
|
||||
request carries no proxy identity at all, which is direct SDK use where there is one tenant by
|
||||
"""A cache/affinity partition for the authenticated caller, derived from the caller's COMPLETE identity.
|
||||
|
||||
Anything stored under this partition is also attributed to the caller, so the partition has to cover
|
||||
every dimension the attribution does. A partition built from a subset is a collision channel: two
|
||||
principals that differ only in an omitted dimension share one entry, and whoever writes last owns the
|
||||
payload and the attribution while any accumulated state stays. End users are the reachable case, since
|
||||
they share a virtual key by design and the key hash alone cannot tell them apart, but the rule is what
|
||||
matters rather than that instance. Composing every dimension also means a tenancy change (a key removed
|
||||
from a team, say) lands on a new partition instead of inheriting the old one's state.
|
||||
|
||||
The dimensions come from the type that declares them, so this cannot drift from what the proxy stamps.
|
||||
Values are hashed rather than concatenated: end-user and project identifiers are customer-chosen strings
|
||||
that would otherwise appear in Redis key names and in every log line quoting one, and hashing identifiers
|
||||
before they become cache keys is what DeploymentAffinityCheck already does. "unscoped" is returned only
|
||||
when the request carries no proxy identity at all, which is direct SDK use with one tenant by
|
||||
construction."""
|
||||
for prefix, field in (
|
||||
("key", "user_api_key_hash"),
|
||||
("user", "user_api_key_user_id"),
|
||||
("team", "user_api_key_team_id"),
|
||||
("org", "user_api_key_org_id"),
|
||||
):
|
||||
value = get_request_metadata_field(request_kwargs, field)
|
||||
if value is not None:
|
||||
return f"{prefix}:{value}"
|
||||
return "unscoped"
|
||||
identity = tuple((field, get_request_metadata_field(request_kwargs, field)) for field in proxy_identity_fields())
|
||||
if all(value is None for _, value in identity):
|
||||
return "unscoped"
|
||||
digest = hashlib.sha256(json.dumps(identity, sort_keys=True, default=str).encode("utf-8")).hexdigest()
|
||||
return f"id:{digest[:32]}"
|
||||
|
||||
|
||||
def get_request_metadata_field(request_kwargs: Mapping[str, object], field: str) -> str | None:
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ from litellm.litellm_core_utils.core_helpers import (
|
|||
get_caller_scope,
|
||||
get_request_metadata_field,
|
||||
iter_request_metadata_dicts,
|
||||
proxy_identity_fields,
|
||||
)
|
||||
from litellm.router_strategy.complexity_router.cache_warming.eligibility import (
|
||||
min_prompt_cache_tokens_for_warm_set,
|
||||
|
|
@ -27,15 +28,11 @@ if TYPE_CHECKING:
|
|||
from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter
|
||||
|
||||
_MAX_UNCOMPRESSED_RATIO = 8
|
||||
_ATTRIBUTION_KEYS = (
|
||||
"user_api_key",
|
||||
"user_api_key_hash",
|
||||
"user_api_key_user_id",
|
||||
"user_api_key_team_id",
|
||||
"user_api_key_org_id",
|
||||
"user_api_key_project_id",
|
||||
"user_api_key_end_user_id",
|
||||
)
|
||||
# The partition a record is stored under and the identity it is attributed to must cover the same
|
||||
# dimensions, or two principals differing only in an omitted one share a record. Both are derived from
|
||||
# proxy_identity_fields() for that reason; user_api_key is the key-state lookup handle rather than a
|
||||
# tenant dimension, so it rides along here and not in the partition.
|
||||
_ATTRIBUTION_KEYS = ("user_api_key", *proxy_identity_fields())
|
||||
|
||||
|
||||
@lru_cache(maxsize=64)
|
||||
|
|
|
|||
|
|
@ -94,7 +94,9 @@ async def test_second_turn_overwrites_payload_and_preserves_other_model_warmth()
|
|||
redis = FakeRedisCache()
|
||||
router = _complexity_router(redis)
|
||||
await router._capture_session(_kwargs(), MESSAGES, "claude-sonnet-4-5")
|
||||
key = CacheWarmingStore.record_key("smart-router", "key:hash-1", "sess-1")
|
||||
from litellm.litellm_core_utils.core_helpers import get_caller_scope
|
||||
|
||||
key = CacheWarmingStore.record_key("smart-router", get_caller_scope(_kwargs()), "sess-1")
|
||||
first = json.loads(redis.hashes[SESSIONS_KEY][key])
|
||||
other_stamp = WarmthStamp(at=123.0, warmed=True).model_dump()
|
||||
redis.data[CacheWarmingStore.warmth_key(key, "gpt-5-mini")] = json.dumps(other_stamp)
|
||||
|
|
@ -166,15 +168,17 @@ async def test_warming_does_not_accept_a_forged_billing_identity():
|
|||
assert record["attribution"]["user_api_key_hash"] == "hash-1"
|
||||
|
||||
|
||||
def test_keyless_tenants_reusing_a_session_id_do_not_share_a_record():
|
||||
"""JWT and other keyless proxy principals carry no key hash. Collapsing them onto one literal would merge
|
||||
distinct tenants that reuse a session id into a single record, so the last writer's payload and
|
||||
attribution would win and later replays could spend under the wrong team."""
|
||||
from litellm.litellm_core_utils.core_helpers import get_caller_scope
|
||||
def test_two_tenants_sharing_a_session_id_never_share_a_partition():
|
||||
"""Whatever is stored under a partition is also attributed to that caller, so the partition has to cover
|
||||
every dimension the attribution does. A partition built from a subset is a collision channel: principals
|
||||
differing only in an omitted dimension share one record, and the last writer owns the payload and the
|
||||
attribution while any state accumulated under it stays. Every dimension is driven independently here,
|
||||
including the two the old strongest-identity partition dropped, end user and project."""
|
||||
from litellm.litellm_core_utils.core_helpers import get_caller_scope, proxy_identity_fields
|
||||
|
||||
tenant_a = {"litellm_metadata": {"session_id": "shared", "user_api_key_team_id": "team-a"}}
|
||||
tenant_b = {"litellm_metadata": {"session_id": "shared", "user_api_key_team_id": "team-b"}}
|
||||
assert get_caller_scope(tenant_a) != get_caller_scope(tenant_b)
|
||||
base = {field: "same" for field in proxy_identity_fields()}
|
||||
scopes = {get_caller_scope({"metadata": {**base, field: "other"}}) for field in proxy_identity_fields()}
|
||||
assert len(scopes) == len(proxy_identity_fields()), "each identity dimension must move the partition"
|
||||
assert get_caller_scope({"metadata": base}) not in scopes
|
||||
assert get_caller_scope({"metadata": {}}) == "unscoped"
|
||||
keyed = {"metadata": {"user_api_key_hash": "hash-1", "user_api_key_team_id": "team-a"}}
|
||||
assert get_caller_scope(keyed) == "key:hash-1"
|
||||
assert "same" not in get_caller_scope({"metadata": base}), "identifiers must not appear in the key"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue