scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 154

fix(router): enforce fusion dependency budgets

Browse source
This commit is contained in:
moe-berri 2026-09-03 15:27:45 -07:00
parent 079065bb39
commit 0a26e0cc51
8 changed files with 218 additions and 19 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

23
litellm/proxy/auth/auth_checks.py
View file

@ -4284,7 +4284,7 @@ async def can_key_call_resolved_model(
llm_model_list: list | None,
valid_token: UserAPIKeyAuth,
llm_router: litellm.Router | None,
) -> None:
) -> tuple[str, ...]:
from litellm.proxy.proxy_server import (
prisma_client,
proxy_logging_obj,
@ -4355,8 +4355,9 @@ async def can_key_call_resolved_model(
proxy_logging_obj=proxy_logging_obj,
)
project_object: LiteLLM_ProjectTableCachedObj | None = None
if valid_token.project_id is not None:
project_object: Final = await get_project_object(
project_object = await get_project_object(
project_id=valid_token.project_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
@ -4369,6 +4370,24 @@ async def can_key_call_resolved_model(
llm_router=llm_router,
)
matched_model_access_groups: Final = await collect_matched_model_access_groups(
model=model,
valid_token=valid_token,
team_object=team_object,
project_object=project_object,
llm_router=llm_router,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
if matched_model_access_groups:
await _model_access_group_max_budget_check(
matched_model_access_groups=matched_model_access_groups,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
)
return matched_model_access_groups
def can_org_access_model(
model: str,

18
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

@ -741,10 +741,18 @@ async def _update_database_and_spend_counters(
)
raise
if defer_budget_counter_update:
return
try:
if defer_budget_counter_update:
if model_access_groups:
from litellm.proxy.proxy_server import increment_fusion_model_access_group_spend_counters
await increment_fusion_model_access_group_spend_counters(
model_access_groups=model_access_groups,
response_cost=response_cost,
budget_reservation=budget_reservation,
)
return
await increment_spend_counters(
token=user_api_key,
team_id=team_id,
@ -756,6 +764,10 @@ async def _update_database_and_spend_counters(
tags=request_tags,
request_started_at=start_time,
model_access_groups=model_access_groups,
# Global scopes reconcile the whole logical Fusion request. Access
# groups are deployment-specific, so the final provider call must
# be charged only its own cost; hidden calls were charged above.
model_access_group_response_cost=response_cost,
)
except Exception:
if budget_reservation is not None:

29
litellm/proxy/proxy_server.py
View file

@ -2670,6 +2670,7 @@ async def increment_spend_counters(
tags: list[str] | None = None,
request_started_at: datetime | None = None,
model_access_groups: Sequence[str] | None = None,
model_access_group_response_cost: float | None = None,
):
"""
Atomically increment spend counters for budget enforcement.
@ -2692,6 +2693,9 @@ async def increment_spend_counters(
return
cost: Final[float] = response_cost
model_access_group_cost: Final = (
cost if model_access_group_response_cost is None else model_access_group_response_cost
)
async def _key_scope(key_token: str) -> None:
# key_token arrives pre-hashed from metadata["user_api_key"] (auth flow
@ -2825,10 +2829,10 @@ async def increment_spend_counters(
else None,
_increment_model_access_group_spend_counters(
model_access_groups=model_access_groups,
response_cost=cost,
response_cost=model_access_group_cost,
reserved_counter_keys=reserved_counter_keys,
)
if model_access_groups
if model_access_groups and model_access_group_cost != 0
else None,
_increment_org_spend_counter(
org_id=org_id,
@ -2853,6 +2857,27 @@ async def increment_spend_counters(
budget_reservation["finalized"] = True
async def increment_fusion_model_access_group_spend_counters(
model_access_groups: Sequence[str],
response_cost: float,
budget_reservation: dict | None, # mutable-ok: shared reservation ledger is read here to avoid duplicate charges
) -> None:
"""Charge one deferred Fusion provider call to only its serving access groups.
Fusion defers the shared key/team/user reservation until its final outer call,
but model access groups are deployment-specific. Updating those counters per
provider call preserves attribution while skipping any group already reserved
for the virtual Fusion model itself.
"""
from litellm.proxy.spend_tracking.budget_reservation import get_reserved_counter_keys
await _increment_model_access_group_spend_counters(
model_access_groups=model_access_groups,
response_cost=response_cost,
reserved_counter_keys=get_reserved_counter_keys(budget_reservation=budget_reservation),
)
async def _reconcile_budget_reservation_for_counter_update(
budget_reservation: dict | None,
response_cost: float | None,

68
litellm/router.py
View file

@ -21,7 +21,16 @@ import time
import traceback
import weakref
from collections import defaultdict
from collections.abc import AsyncGenerator, AsyncIterator, Callable, Generator, Iterator, Mapping, Sequence
from collections.abc import (
AsyncGenerator,
AsyncIterator,
Callable,
Generator,
Iterator,
Mapping,
MutableMapping,
Sequence,
)
from functools import lru_cache, partial
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypeVar, Union, cast
@ -68,6 +77,7 @@ from litellm.litellm_core_utils.coroutine_checker import coroutine_checker
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.litellm_core_utils.dd_tracing import tracer
from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider
from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.litellm_core_utils.ptu_pricing import (
PTU_COST_ATTRIBUTION_ENV_VAR,
@ -9470,9 +9480,11 @@ class Router:
async def _authorize_fusion_dependencies(
self,
fusion_router: FusionRouter,
request_kwargs: Mapping[str, object],
request_kwargs: MutableMapping[ # mutable-ok: request-local carrier is enriched before hidden calls dispatch
str, object
],
) -> None:
"""Apply the originating proxy caller's model access to every hidden call."""
"""Apply the originating proxy caller's model access and group budgets to every hidden call."""
metadata_values: Final = tuple(request_kwargs.get(key) for key in ("litellm_metadata", "metadata"))
raw_user_api_key_auth: Final = next(
(
@ -9512,13 +9524,51 @@ class Router:
)
)
)
for dependency_model in dependency_models:
await can_key_call_resolved_model(
model=dependency_model,
llm_model_list=self.model_list,
valid_token=user_api_key_auth,
llm_router=self,
matched_dependency_groups: Final = await asyncio.gather(
*(
can_key_call_resolved_model(
model=dependency_model,
llm_model_list=self.model_list,
valid_token=user_api_key_auth,
llm_router=self,
)
for dependency_model in dependency_models
)
)
dependency_access_groups: Final = frozenset(
group for matched_groups in matched_dependency_groups for group in matched_groups
)
if not dependency_access_groups:
return
# Keep the group gating the virtual Fusion model and add every group
# gating a hidden dependency. The normal spend writer narrows this
# authorization upper bound to groups serving each provider call.
all_groups: Final = tuple(
dict.fromkeys(
(
*(user_api_key_auth.matched_model_access_groups or ()),
*sorted(dependency_access_groups),
)
)
)
user_api_key_auth.matched_model_access_groups = list( # mutable-ok: auth schema requires a list carrier
all_groups
)
for metadata_key in ("litellm_metadata", "metadata"):
metadata = request_kwargs.get(metadata_key)
if not isinstance(metadata, Mapping):
continue
mutable_metadata = (
metadata
if isinstance(metadata, dict)
else dict(metadata) # mutable-ok: SDK metadata boundary requires a native mapping
)
mutable_metadata[MODEL_ACCESS_GROUP_METADATA_KEY] = list( # mutable-ok: metadata JSON requires a list
all_groups
)
mutable_metadata["user_api_key_auth"] = user_api_key_auth
request_kwargs[metadata_key] = mutable_metadata # rebind-ok: enrich request-local metadata for dispatch
async def _fusion_asearch( # kwargs-ok: bridge preserves the Router.asearch keyword surface
self,

29
tests/test_litellm/proxy/auth/test_model_access_group_budgets.py
View file

@ -25,6 +25,7 @@ from litellm.proxy._types import (
)
from litellm.proxy.auth.auth_checks import (
_model_access_group_max_budget_check,
can_key_call_resolved_model,
collect_matched_model_access_groups,
common_checks,
stamp_matched_model_access_groups,
@ -381,6 +382,34 @@ async def test_group_exactly_at_its_max_budget_blocks_the_request():
assert exc_info.value.current_cost == 10.0
@pytest.mark.asyncio
async def test_resolved_model_authorization_enforces_its_access_group_budget(
monkeypatch: pytest.MonkeyPatch,
):
from litellm.proxy import proxy_server
cache = await _cache()
prisma = _RecordingPrismaClient(_MagBudgetRow("tier-a", spend=10.0, max_budget=10.0))
router = Router(model_list=MODEL_LIST)
valid_token = UserAPIKeyAuth(api_key="hashed", models=["tier-a"])
read, seen = _spend_reader({MODEL_ACCESS_GROUP_COUNTER_KEY: 10.0})
monkeypatch.setattr(proxy_server, "prisma_client", prisma)
monkeypatch.setattr(proxy_server, "user_api_key_cache", cache)
monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(user_api_key_cache=cache))
monkeypatch.setattr(proxy_server, "get_current_spend", read)
with pytest.raises(litellm.BudgetExceededError):
await can_key_call_resolved_model(
model="gpt-4o",
llm_model_list=MODEL_LIST,
valid_token=valid_token,
llm_router=router,
)
assert seen == [MODEL_ACCESS_GROUP_COUNTER_KEY]
@pytest.mark.asyncio
async def test_group_just_under_its_max_budget_passes():
"""Asserting the counter was read is what keeps this honest: a group that got skipped entirely,

2
tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py
View file

@ -607,6 +607,7 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda
tags=["tag-a"],
request_started_at=start_time,
model_access_groups=("premium",),
model_access_group_response_cost=0.2,
)
@ -730,6 +731,7 @@ async def test_fusion_hidden_costs_accumulate_then_continuation_reconciles_once(
increment.assert_awaited_once()
assert increment.await_args.kwargs["response_cost"] == pytest.approx(0.7)
assert increment.await_args.kwargs["model_access_group_response_cost"] == pytest.approx(0.4)
assert increment.await_args.kwargs["budget_reservation"] is reservation

49
tests/test_litellm/proxy/test_budget_reservation.py
View file

@ -1192,7 +1192,7 @@ def test_fusion_reservation_does_not_return_a_partial_additive_estimate() -> Non
def child_estimate(*, model: str, **_: object) -> float | None:
return None if model == "panel" else 1.0
with patch(
with patch( # test-quality-ok: isolates one unavailable dependency estimate to verify all-or-nothing Fusion admission
"litellm.proxy.spend_tracking.budget_reservation._estimate_request_max_cost_for_model",
side_effect=child_estimate,
):
@ -3445,6 +3445,53 @@ async def test_model_access_group_counter_accumulates_across_calls_without_a_res
assert counter_cache.in_memory_cache.get_cache(key=model_access_group_spend_counter_key("")) is None
@pytest.mark.asyncio
async def test_model_access_group_can_use_its_provider_call_cost_during_fusion_reconciliation(spend_counter_state):
counter_cache, key_cache = spend_counter_state
await _cache_model_access_group_budget(key_cache, "panel", spend=1.0, max_budget=25.0)
from litellm.proxy.proxy_server import increment_spend_counters
await increment_spend_counters(
token=None,
team_id=None,
user_id=None,
response_cost=0.7,
model_access_groups=["panel"],
model_access_group_response_cost=0.2,
)
assert counter_cache.in_memory_cache.get_cache(
key=model_access_group_spend_counter_key("panel")
) == pytest.approx(1.2)
@pytest.mark.asyncio
async def test_deferred_fusion_access_group_cost_skips_a_group_reserved_by_the_virtual_model(spend_counter_state):
counter_cache, key_cache = spend_counter_state
await _cache_model_access_group_budget(key_cache, "fusion", spend=1.0, max_budget=25.0)
await _cache_model_access_group_budget(key_cache, "panel", spend=2.0, max_budget=25.0)
await counter_cache.async_set_cache(key=model_access_group_spend_counter_key("fusion"), value=1.5)
reservation = {
"entries": [{"counter_key": model_access_group_spend_counter_key("fusion")}],
}
from litellm.proxy.proxy_server import increment_fusion_model_access_group_spend_counters
await increment_fusion_model_access_group_spend_counters(
model_access_groups=["fusion", "panel"],
response_cost=0.2,
budget_reservation=reservation,
)
assert counter_cache.in_memory_cache.get_cache(
key=model_access_group_spend_counter_key("fusion")
) == pytest.approx(1.5)
assert counter_cache.in_memory_cache.get_cache(
key=model_access_group_spend_counter_key("panel")
) == pytest.approx(2.2)
@pytest.mark.asyncio
async def test_reserved_model_access_group_is_not_charged_twice(spend_counter_state):
"""The reservation already wrote this counter, so the post-call pass has to skip it."""

19
tests/test_litellm/test_fusion_router.py
View file

@ -17,6 +17,7 @@ from litellm.fusion_router import (
fusion_router_dependencies,
validate_fusion_router_write,
)
from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY
from litellm.router import Router
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import ModelResponseStream
@ -683,18 +684,32 @@ async def test_proxy_fusion_authorizes_every_hidden_model(monkeypatch: pytest.Mo
)
model_list[-1]["litellm_params"]["fusion_router_config"]["analyst_model"] = "analyst"
router = Router(model_list=model_list)
authorize = AsyncMock(return_value=None)
authorize = AsyncMock(side_effect=[("outer-budget",), ("panel-budget",), ("analyst-budget",)])
monkeypatch.setattr(auth_checks, "can_key_call_resolved_model", authorize)
auth = UserAPIKeyAuth(models=["*"], matched_model_access_groups=["fusion-budget"])
metadata: dict[str, object] = {"user_api_key_auth": auth}
await router._authorize_fusion_dependencies( # pyright: ignore[reportPrivateUsage]
fusion_router=router.fusion_routers["fusion/test"],
request_kwargs={
"metadata": {"user_api_key_auth": UserAPIKeyAuth(models=["*"])},
"metadata": metadata,
"proxy_server_request": {"body": {"model": "fusion/test"}},
},
)
assert [call.kwargs["model"] for call in authorize.await_args_list] == ["outer", "panel-a", "analyst"]
assert metadata[MODEL_ACCESS_GROUP_METADATA_KEY] == [
"fusion-budget",
"analyst-budget",
"outer-budget",
"panel-budget",
]
assert auth.matched_model_access_groups == [
"fusion-budget",
"analyst-budget",
"outer-budget",
"panel-budget",
]
@pytest.mark.asyncio