mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
fix(router): bound fusion timeout accounting and config access
This commit is contained in:
parent
dd40897f24
commit
563107b497
7 changed files with 126 additions and 6 deletions
|
|
@ -43,7 +43,7 @@ Call `fusion/general` exactly like any other model. `invocation: auto` lets the
|
|||
|
||||
`panel_timeout_seconds` bounds each panel member and the complete analyst phase, including any Search Tool loop. If the analyst times out, the outer model still receives the successful raw panel responses.
|
||||
|
||||
The outer model must support function calling. Panel and analyst models only need function calling when a Search Tool is configured. Granting access to the Fusion model lets the request use its administrator-configured model and search dependencies; the panel query and private research are sent to those deployments under their normal provider data policies.
|
||||
The outer model must support function calling. Panel and analyst models only need function calling when a Search Tool is configured. Granting access to the Fusion model lets the request use its administrator-configured model and search dependencies; the panel query and private research are sent to those deployments under their normal provider data policies. Only proxy administrators can create or edit DB-backed Fusion models; config-file definitions are controlled by the proxy operator.
|
||||
|
||||
## Operational behavior
|
||||
|
||||
|
|
@ -51,6 +51,7 @@ The outer model must support function calling. Panel and analyst models only nee
|
|||
- Initial outer, panel, analyst, continuation, and search calls are marked separately in spend logs. They inherit the caller identity and remain part of one logical Fusion request.
|
||||
- Client-visible `usage` describes the outer response returned to that client. Hidden panel, analyst, and search usage remains in its separately tagged spend-log rows; budget reconciliation includes the cost of every hidden call rather than merging heterogeneous model tokens into one public token count.
|
||||
- Admission control reserves the worst-case model-call cost. Hidden calls accumulate against that shared reservation, and the direct initial response or final continuation reconciles it once. This keeps concurrent requests from spending the same remaining budget while Fusion is still running.
|
||||
- A deliberately timed-out panel or analyst call does not charge the worst-case cost of the entire Fusion request. Completed calls retain their known spend; the missing-callback fallback remains fail-closed for calls that completed without reporting cost.
|
||||
- Chat-completion streaming is buffered until LiteLLM knows whether the private tool was invoked. A direct response is replayed as a normal stream; a Fusion invocation suppresses the private tool-call stream and exposes only the final outer-model stream.
|
||||
- A request-level `tool_choice: required` is considered satisfied when Fusion runs. The continuation changes it to `auto` when client tools exist, or removes it when they do not, so the outer model can finish instead of being forced into a second tool call.
|
||||
- A client tool named `litellm_fusion` is rejected because that name is reserved for the private server tool.
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ from litellm.constants import (
|
|||
INTERNAL_CALL_ORIGIN_METADATA_KEY,
|
||||
)
|
||||
from litellm.litellm_core_utils.fusion_budget import (
|
||||
cancel_fusion_budget_call,
|
||||
complete_fusion_budget_call,
|
||||
register_fusion_budget_call,
|
||||
wait_for_fusion_budget_calls,
|
||||
|
|
@ -909,7 +910,7 @@ class FusionRouter:
|
|||
_fusion_proxy_auth_required=isinstance(request_kwargs.get("proxy_server_request"), Mapping),
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
complete_fusion_budget_call(metadata, cost_known=False)
|
||||
cancel_fusion_budget_call(metadata)
|
||||
raise
|
||||
except Exception:
|
||||
complete_fusion_budget_call(metadata, cost_known=True)
|
||||
|
|
@ -967,7 +968,7 @@ class FusionRouter:
|
|||
try:
|
||||
response = await self._completion(model=model, messages=current_messages, stream=False, **call_kwargs)
|
||||
except asyncio.CancelledError:
|
||||
complete_fusion_budget_call(call_metadata, cost_known=False)
|
||||
cancel_fusion_budget_call(call_metadata)
|
||||
raise
|
||||
except Exception:
|
||||
complete_fusion_budget_call(call_metadata, cost_known=True)
|
||||
|
|
@ -1049,7 +1050,7 @@ class FusionRouter:
|
|||
**kwargs,
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
complete_fusion_budget_call(metadata, cost_known=False)
|
||||
cancel_fusion_budget_call(metadata)
|
||||
raise
|
||||
except Exception:
|
||||
complete_fusion_budget_call(metadata, cost_known=True)
|
||||
|
|
|
|||
|
|
@ -71,6 +71,18 @@ def complete_fusion_budget_call(
|
|||
unpriced.append(token)
|
||||
|
||||
|
||||
def cancel_fusion_budget_call(metadata: Mapping[str, object]) -> None:
|
||||
"""Finish a call deliberately cancelled by Fusion without marking the whole request unpriced.
|
||||
|
||||
Panel and analyst timeouts actively cancel their in-flight child call. They
|
||||
are different from a completed provider call whose cost callback went
|
||||
missing: the latter must retain the conservative full-reservation fallback,
|
||||
while the former must not turn one timed-out advisory member into a charge
|
||||
for every possible Fusion call.
|
||||
"""
|
||||
complete_fusion_budget_call(metadata, cost_known=True)
|
||||
|
||||
|
||||
async def wait_for_fusion_budget_calls(
|
||||
metadata: Mapping[str, object],
|
||||
*,
|
||||
|
|
|
|||
|
|
@ -296,6 +296,34 @@ def _raise_on_strategy_router_write_violation(
|
|||
)
|
||||
|
||||
|
||||
def _raise_if_non_admin_configures_fusion(
|
||||
*,
|
||||
incoming_params: GenericLiteLLMParams | None,
|
||||
existing_params: GenericLiteLLMParams | None,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> None:
|
||||
"""Keep the virtual model as the runtime ACL without allowing tenant privilege escalation."""
|
||||
if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN:
|
||||
return
|
||||
incoming_model: Final = getattr(incoming_params, "model", None)
|
||||
existing_model: Final = getattr(existing_params, "model", None)
|
||||
config_supplied: Final = any(
|
||||
params is not None and params.fusion_router_config is not None for params in (incoming_params, existing_params)
|
||||
)
|
||||
if config_supplied or any(
|
||||
isinstance(model, str) and is_fusion_router_model(model) for model in (incoming_model, existing_model)
|
||||
):
|
||||
raise ProxyException(
|
||||
message=(
|
||||
"Only proxy admins can create or edit Fusion models because their configured "
|
||||
"models and Search Tool execute with the Fusion model's authorization."
|
||||
),
|
||||
type=ProxyErrorTypes.auth_error.value,
|
||||
code=status.HTTP_403_FORBIDDEN,
|
||||
param="litellm_params.fusion_router_config",
|
||||
)
|
||||
|
||||
|
||||
HEURISTIC_V2_SLOT_LOCK_KEY: Final = 5_872_301
|
||||
_HEURISTIC_V2_LOCK_SQL: Final = "SELECT 1 AS locked FROM pg_advisory_xact_lock($1)"
|
||||
_HEURISTIC_V2_DB_ROWS_SQL: Final = """
|
||||
|
|
@ -784,6 +812,12 @@ async def patch_model(
|
|||
premium_user=premium_user,
|
||||
)
|
||||
|
||||
_raise_if_non_admin_configures_fusion(
|
||||
incoming_params=patch_data.litellm_params,
|
||||
existing_params=db_model.litellm_params,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
# Pause/resume (`blocked`) is a proxy-admin-only privilege. Team admins
|
||||
# passed the auth check above for team-scoped models, but they must not
|
||||
# be able to unblock (or block) a model their proxy admin has paused.
|
||||
|
|
@ -1935,6 +1969,12 @@ async def add_new_model(
|
|||
premium_user=premium_user,
|
||||
)
|
||||
|
||||
_raise_if_non_admin_configures_fusion(
|
||||
incoming_params=model_params.litellm_params,
|
||||
existing_params=None,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
ModelManagementAuthChecks.can_user_attach_credential(
|
||||
litellm_params=model_params.litellm_params,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
|
|
@ -2115,6 +2155,12 @@ async def update_model(
|
|||
premium_user=premium_user,
|
||||
)
|
||||
|
||||
_raise_if_non_admin_configures_fusion(
|
||||
incoming_params=model_params.litellm_params,
|
||||
existing_params=deployment.litellm_params,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
ModelManagementAuthChecks.can_user_attach_credential(
|
||||
litellm_params=model_params.litellm_params,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
|
|
|
|||
|
|
@ -2550,6 +2550,10 @@ class Router:
|
|||
model=model,
|
||||
llm_provider="",
|
||||
)
|
||||
# The public Fusion name is intentionally the runtime ACL. Only
|
||||
# proxy admins may define its hidden dependency graph, so callers
|
||||
# cannot use a team-scoped model write to escalate into models
|
||||
# they could not otherwise reach.
|
||||
self._validate_fusion_proxy_context(request_kwargs=kwargs)
|
||||
response = ( # rebind-ok: one mutually exclusive dispatch branch assigns it
|
||||
await fusion_router.acompletion(
|
||||
|
|
@ -9544,6 +9548,10 @@ class Router:
|
|||
param=None,
|
||||
code=403,
|
||||
) from exc
|
||||
# Search is part of the administrator-defined Fusion capability, not a
|
||||
# caller-selected Search API request. Rechecking the caller's direct
|
||||
# Search Tool allowlist here would make an allowed Fusion model fail on
|
||||
# its private dependency and expose its internal configuration.
|
||||
return await self.asearch(
|
||||
model=model,
|
||||
query=query,
|
||||
|
|
|
|||
|
|
@ -4029,6 +4029,44 @@ class TestStrategyRouterWriteValidation:
|
|||
assert violation is not None
|
||||
assert "panel_models" in violation
|
||||
|
||||
def test_only_proxy_admin_can_define_or_edit_fusion_dependencies(self):
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
_raise_if_non_admin_configures_fusion,
|
||||
)
|
||||
from litellm.proxy.proxy_server import ProxyException
|
||||
from litellm.types.router import updateLiteLLMParams
|
||||
|
||||
admin = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
team_admin = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER)
|
||||
fusion = LiteLLM_Params(
|
||||
model="fusion_router",
|
||||
fusion_router_config={"outer_model": "outer", "panel_models": ["panel-a"]},
|
||||
)
|
||||
|
||||
_raise_if_non_admin_configures_fusion(
|
||||
incoming_params=fusion,
|
||||
existing_params=None,
|
||||
user_api_key_dict=admin,
|
||||
)
|
||||
with pytest.raises(ProxyException, match="Only proxy admins"):
|
||||
_raise_if_non_admin_configures_fusion(
|
||||
incoming_params=fusion,
|
||||
existing_params=None,
|
||||
user_api_key_dict=team_admin,
|
||||
)
|
||||
with pytest.raises(ProxyException, match="Only proxy admins"):
|
||||
_raise_if_non_admin_configures_fusion(
|
||||
incoming_params=updateLiteLLMParams(),
|
||||
existing_params=fusion,
|
||||
user_api_key_dict=team_admin,
|
||||
)
|
||||
|
||||
_raise_if_non_admin_configures_fusion(
|
||||
incoming_params=LiteLLM_Params(model="openai/gpt-4o"),
|
||||
existing_params=None,
|
||||
user_api_key_dict=team_admin,
|
||||
)
|
||||
|
||||
def test_fusion_config_only_patch_is_validated_against_stored_marker(self):
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
_strategy_router_write_violation,
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ from unittest.mock import AsyncMock
|
|||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.constants import FUSION_BUDGET_ACTIVE_KEY, FUSION_BUDGET_UNPRICED_CALL_IDS_KEY
|
||||
from litellm.fusion_router import (
|
||||
FUSION_TOOL_NAME,
|
||||
FusionCompletionCaller,
|
||||
|
|
@ -18,7 +19,10 @@ from litellm.fusion_router import (
|
|||
fusion_router_dependencies,
|
||||
validate_fusion_router_write,
|
||||
)
|
||||
from litellm.litellm_core_utils.fusion_budget import complete_fusion_budget_call
|
||||
from litellm.litellm_core_utils.fusion_budget import (
|
||||
complete_fusion_budget_call,
|
||||
fusion_budget_reconciliation_cost,
|
||||
)
|
||||
from litellm.router import Router
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.utils import ModelResponseStream
|
||||
|
|
@ -631,11 +635,19 @@ async def test_analyst_timeout_degrades_to_raw_panel_responses() -> None:
|
|||
raise AssertionError("unreachable")
|
||||
|
||||
completion = HangingAnalystCompletion()
|
||||
reservation: dict[str, object] = {
|
||||
"reserved_cost": 10.0,
|
||||
FUSION_BUDGET_ACTIVE_KEY: True,
|
||||
}
|
||||
response = await asyncio.wait_for(
|
||||
_router(completion, panel_timeout_seconds=0.2).acompletion(
|
||||
messages=[{"role": "user", "content": "Hard question"}],
|
||||
stream=False,
|
||||
request_kwargs={},
|
||||
request_kwargs={
|
||||
"litellm_metadata": {
|
||||
"user_api_key_budget_reservation": reservation,
|
||||
}
|
||||
},
|
||||
),
|
||||
timeout=2,
|
||||
)
|
||||
|
|
@ -645,6 +657,8 @@ async def test_analyst_timeout_degrades_to_raw_panel_responses() -> None:
|
|||
assert completion.analyst_started.is_set()
|
||||
assert completion.analyst_cancelled.is_set()
|
||||
assert response._hidden_params["fusion"]["analysis_available"] is False
|
||||
assert reservation.get(FUSION_BUDGET_UNPRICED_CALL_IDS_KEY) in (None, [])
|
||||
assert fusion_budget_reconciliation_cost(reservation, known_cost=0.4) == pytest.approx(0.4)
|
||||
payload = json.loads(completion.calls[-1]["messages"][-1]["content"])
|
||||
assert [item["content"] for item in payload["responses"]] == ["Panel A", "Panel B"]
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue