mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
fix(proxy): judge the free-model budget waiver by the group an alias routes to (#44638)
Some checks are pending
CI Coverage / assert-ci-coverage (push) Waiting to run
CodeQL / Analyze (actions) (push) Waiting to run
CodeQL / Analyze (javascript-typescript) (push) Waiting to run
CodeQL / Analyze (python) (push) Waiting to run
CodSpeed Benchmarks / benchmarks (push) Waiting to run
Helm unit test / unit-test (push) Waiting to run
Lens Worker Image / lens-worker-image (push) Waiting to run
Publish basedpyright base counts / publish (push) Waiting to run
Scorecard supply-chain security / Scorecard analysis (push) Waiting to run
Code Quality Checks / code-quality (push) Waiting to run
Code Quality Checks / python-310-import-smoke (push) Waiting to run
UI Unit Tests / ui-unit-tests (push) Waiting to run
LiteLLM Rust / rust-lint (push) Waiting to run
LiteLLM Rust / rust-test (push) Waiting to run
LiteLLM Rust / rust-wheel (push) Waiting to run
Unit Tests: Documentation Validation / documentation (push) Waiting to run
Unit Tests: Proxy DB Operations / assert-shard-coverage (push) Waiting to run
Unit Tests: Proxy DB Operations / auth-checks (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / budgets (push) Blocked by required conditions
Unit Tests / Build the Rust bridge (push) Waiting to run
Unit Tests / caching-local (push) Blocked by required conditions
VS Code Extension / vscode-extension (push) Waiting to run
Unit Tests: Proxy DB Operations / custom-logging (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / db-and-spend (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / endpoints-and-responses (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / guardrails-hooks (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / jwt-and-keys (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / key-generation (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / logging-misc (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / proxy-runtime (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / proxy-server-core (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / proxy-utils (push) Blocked by required conditions
Unit Tests / core-utils (push) Blocked by required conditions
Unit Tests / enterprise-managed-files (push) Blocked by required conditions
Unit Tests / enterprise-package (push) Blocked by required conditions
Unit Tests / enterprise-routing (push) Blocked by required conditions
Unit Tests / integrations (push) Blocked by required conditions
Unit Tests / OpenAI and Meta Providers (push) Blocked by required conditions
Unit Tests / All Other Providers (push) Blocked by required conditions
Unit Tests / Vertex AI (push) Blocked by required conditions
Unit Tests / misc (push) Blocked by required conditions
Unit Tests / misc-dirs (push) Blocked by required conditions
Unit Tests / proxy-auth (push) Blocked by required conditions
Unit Tests / proxy-endpoints (push) Blocked by required conditions
Unit Tests / proxy-extras (push) Blocked by required conditions
Unit Tests / proxy-feature-endpoints (push) Blocked by required conditions
Unit Tests / proxy-hooks-client (push) Blocked by required conditions
Unit Tests / proxy-infra (push) Blocked by required conditions
Unit Tests / proxy-infra-root (push) Blocked by required conditions
Unit Tests / proxy-server (push) Blocked by required conditions
Unit Tests / responses-caching-types (push) Blocked by required conditions
Unit Tests / unit (push) Blocked by required conditions
GitHub Actions Security Analysis / zizmor (push) Waiting to run
Some checks are pending
CI Coverage / assert-ci-coverage (push) Waiting to run
CodeQL / Analyze (actions) (push) Waiting to run
CodeQL / Analyze (javascript-typescript) (push) Waiting to run
CodeQL / Analyze (python) (push) Waiting to run
CodSpeed Benchmarks / benchmarks (push) Waiting to run
Helm unit test / unit-test (push) Waiting to run
Lens Worker Image / lens-worker-image (push) Waiting to run
Publish basedpyright base counts / publish (push) Waiting to run
Scorecard supply-chain security / Scorecard analysis (push) Waiting to run
Code Quality Checks / code-quality (push) Waiting to run
Code Quality Checks / python-310-import-smoke (push) Waiting to run
UI Unit Tests / ui-unit-tests (push) Waiting to run
LiteLLM Rust / rust-lint (push) Waiting to run
LiteLLM Rust / rust-test (push) Waiting to run
LiteLLM Rust / rust-wheel (push) Waiting to run
Unit Tests: Documentation Validation / documentation (push) Waiting to run
Unit Tests: Proxy DB Operations / assert-shard-coverage (push) Waiting to run
Unit Tests: Proxy DB Operations / auth-checks (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / budgets (push) Blocked by required conditions
Unit Tests / Build the Rust bridge (push) Waiting to run
Unit Tests / caching-local (push) Blocked by required conditions
VS Code Extension / vscode-extension (push) Waiting to run
Unit Tests: Proxy DB Operations / custom-logging (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / db-and-spend (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / endpoints-and-responses (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / guardrails-hooks (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / jwt-and-keys (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / key-generation (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / logging-misc (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / proxy-runtime (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / proxy-server-core (push) Blocked by required conditions
Unit Tests: Proxy DB Operations / proxy-utils (push) Blocked by required conditions
Unit Tests / core-utils (push) Blocked by required conditions
Unit Tests / enterprise-managed-files (push) Blocked by required conditions
Unit Tests / enterprise-package (push) Blocked by required conditions
Unit Tests / enterprise-routing (push) Blocked by required conditions
Unit Tests / integrations (push) Blocked by required conditions
Unit Tests / OpenAI and Meta Providers (push) Blocked by required conditions
Unit Tests / All Other Providers (push) Blocked by required conditions
Unit Tests / Vertex AI (push) Blocked by required conditions
Unit Tests / misc (push) Blocked by required conditions
Unit Tests / misc-dirs (push) Blocked by required conditions
Unit Tests / proxy-auth (push) Blocked by required conditions
Unit Tests / proxy-endpoints (push) Blocked by required conditions
Unit Tests / proxy-extras (push) Blocked by required conditions
Unit Tests / proxy-feature-endpoints (push) Blocked by required conditions
Unit Tests / proxy-hooks-client (push) Blocked by required conditions
Unit Tests / proxy-infra (push) Blocked by required conditions
Unit Tests / proxy-infra-root (push) Blocked by required conditions
Unit Tests / proxy-server (push) Blocked by required conditions
Unit Tests / responses-caching-types (push) Blocked by required conditions
Unit Tests / unit (push) Blocked by required conditions
GitHub Actions Security Analysis / zizmor (push) Waiting to run
* fix(proxy): judge the free-model budget waiver by the group an alias routes to A hidden model_group_alias can reuse the name of a real model_name. The router serves that name from the alias target, but two of the three checks behind the free-model budget waiver read the deployments of both the alias target and the real model that shares the name. So an over-budget key was served through an alias whose name belongs to a model with an explicit $0 price when the target is $0 only by the cost map, which the same key is refused on by its own name. The mirror case refused a free target because the shadowed name belongs to a PTU-priced deployment. Resolve the alias once and have the explicit-cost and PTU checks read the routed group, the same group the price check already reads. * test(proxy): cover the plain alias form of a shadowed free model name The same wrong verdict exists for a plain string alias, so the unpriced-target case now runs for both alias shapes. * fix(proxy): judge an alias chain's budget waiver by the deployments it is served from The explicit-price and PTU checks read the alias target through Router.get_model_list(), which follows a second alias hop when the target is itself an alias key. The router never takes that hop, so an alias chain was judged by a deployment the request never reaches. The checks now take the deployments named after the routed group, or the wildcard deployment serving it when none carries its name. * fix(proxy): refuse the budget waiver when an alias chain is served by a priced wildcard route * test(integration): cover the shadowing alias budget gate end to end Adds the audit cells for a hidden alias whose name shadows an explicitly free group: streamed SDK refusals on chat, responses and messages, embeddings, the free wildcard and mixed-group paths, per-model budgets, JWT and custom auth callers, cache hits, alias removal under traffic, provider failure and fallback, a concurrent outage burst, and a worker kill on an owned two-worker proxy * test(integration): cite the cost-map rows the shadowing alias cells rely on --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
parent
ff5084687c
commit
4a963a6c0b
3 changed files with 1189 additions and 25 deletions
|
|
@ -148,6 +148,7 @@ from litellm.repositories.user_repository import UserRepository
|
|||
from litellm.router import Router
|
||||
from litellm.types.proxy.auth.auth_checks import UserNotFoundError
|
||||
from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget
|
||||
from litellm.types.router import DeploymentTypedDict
|
||||
from litellm.utils import get_utc_datetime
|
||||
|
||||
from .auth_checks_organization import (
|
||||
|
|
@ -527,7 +528,8 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None
|
|||
# not from defaulted sparse auto-registration entries.
|
||||
# See: https://github.com/BerriAI/litellm/issues/24770
|
||||
safe_name = str(model_name).replace("\n", "").replace("\r", "")
|
||||
if not _is_cost_explicitly_configured(model_name, llm_router):
|
||||
served_deployments = _deployments_served_for(model_name, llm_router)
|
||||
if not _is_cost_explicitly_configured(served_deployments):
|
||||
verbose_proxy_logger.debug(
|
||||
"Model %s has zero cost but no explicit cost "
|
||||
"configuration in model_cost entry — treating as unknown "
|
||||
|
|
@ -538,7 +540,16 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None
|
|||
zero_cost_cache[model_name] = False
|
||||
return False
|
||||
|
||||
if _has_ptu_flat_cost(model_name, llm_router):
|
||||
if _has_explicit_positive_cost(served_deployments):
|
||||
verbose_proxy_logger.debug(
|
||||
"Model %s routes to a deployment with an explicit positive per-token price (enforce budget)",
|
||||
safe_name,
|
||||
)
|
||||
if zero_cost_cache is not None:
|
||||
zero_cost_cache[model_name] = False
|
||||
return False
|
||||
|
||||
if _has_ptu_flat_cost(served_deployments):
|
||||
verbose_proxy_logger.debug(
|
||||
"Model %s prices reserved PTU capacity as a flat cost, so its zero per-token "
|
||||
"rate is not a free model (enforce budget)",
|
||||
|
|
@ -570,40 +581,56 @@ _NO_MODEL_INFO: Final[Mapping[str, object]] = MappingProxyType({})
|
|||
_TEAM_GRANT_RELATIONS: Final[Mapping[str, object]] = MappingProxyType({"litellm_model_table": True})
|
||||
|
||||
|
||||
def _has_ptu_flat_cost(model: str, llm_router: "Router") -> bool:
|
||||
"""Whether any deployment in the model group bills reserved PTU capacity as a flat cost.
|
||||
def _deployments_served_for(model: str, llm_router: "Router") -> Sequence[DeploymentTypedDict]:
|
||||
"""
|
||||
The deployments a request to ``model`` routes to, for the gates that waive budget checks.
|
||||
|
||||
The router resolves a ``model_group_alias`` name to its target exactly once and serves that
|
||||
group, never a real deployment that shares the alias's name. ``Router.get_model_list()`` on the
|
||||
routed group reads those deployments unless the group is itself an alias key, where it would
|
||||
follow the alias a second hop the router never takes; then the deployments named after the
|
||||
group count, or the wildcard route serving the group when none carries its name, the same
|
||||
fallback ``get_model_list()`` takes.
|
||||
"""
|
||||
routed_group: Final = llm_router.routable_model_group(model)
|
||||
if llm_router.routable_model_group(routed_group) == routed_group:
|
||||
return llm_router.get_model_list(model_name=routed_group) or ()
|
||||
named: Final = llm_router.get_model_list_from_model_alias(model_name=model)
|
||||
if named:
|
||||
return named
|
||||
return tuple(
|
||||
DeploymentTypedDict(**deployment)
|
||||
for deployment in llm_router.pattern_router.get_deployments_by_pattern(model=routed_group)
|
||||
)
|
||||
|
||||
|
||||
def _has_ptu_flat_cost(deployments: Sequence[DeploymentTypedDict]) -> bool:
|
||||
"""Whether any of the deployments bills reserved PTU capacity as a flat cost.
|
||||
|
||||
Such a deployment carries an explicit zero per-token price so the flat cost is not charged
|
||||
twice, which otherwise reads here as a free model and waives every budget check for it.
|
||||
|
||||
Resolved through ``Router.get_model_list()``, which includes ``model_group_alias``, because
|
||||
this runs after the explicit-cost gate: resolving that gate alone would let an aliased PTU
|
||||
group through as free.
|
||||
"""
|
||||
for deployment in llm_router.get_model_list(model_name=model) or ():
|
||||
for deployment in deployments:
|
||||
model_info = deployment.get("model_info") or _NO_MODEL_INFO
|
||||
if model_info.get("ptu_count") is not None and model_info.get("cost_per_ptu_per_hour") is not None:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _is_cost_explicitly_configured(model: str, llm_router: "Router") -> bool:
|
||||
def _is_cost_explicitly_configured(deployments: Sequence[DeploymentTypedDict]) -> bool:
|
||||
"""
|
||||
Check if any deployment in the model group has cost fields explicitly
|
||||
set in its litellm.model_cost entry.
|
||||
Check if any of the deployments has cost fields explicitly set in its
|
||||
litellm.model_cost entry.
|
||||
|
||||
When Router._create_deployment() registers a model not in the global
|
||||
cost map, it creates a sparse entry like {"id": "<hash>"} with no cost
|
||||
fields. _get_model_info_helper() then defaults missing costs to 0.
|
||||
This function detects that scenario by checking the raw model_cost entry.
|
||||
|
||||
The group is resolved through ``Router.get_model_list()``, the same resolution
|
||||
``get_model_group_info()`` applies when the caller reads the cost a few lines earlier, so the
|
||||
two lookups cannot disagree, including for names defined in ``Router.model_group_alias``.
|
||||
It also reaches a deployment that prices itself through its ``model_info`` block, whose entry
|
||||
lands in the cost map under the deployment id.
|
||||
"""
|
||||
for deployment in llm_router.get_model_list(model_name=model) or ():
|
||||
for deployment in deployments:
|
||||
model_id = (deployment.get("model_info") or _EMPTY_COST_ENTRY).get("id")
|
||||
if model_id is None:
|
||||
continue
|
||||
|
|
@ -613,6 +640,26 @@ def _is_cost_explicitly_configured(model: str, llm_router: "Router") -> bool:
|
|||
return False
|
||||
|
||||
|
||||
def _has_explicit_positive_cost(deployments: Sequence[DeploymentTypedDict]) -> bool:
|
||||
"""Whether any of the deployments carries an explicit positive per-token price.
|
||||
|
||||
The group's price is read from the deployments named after the routed group and from
|
||||
that group's own alias target, while an alias chain is served from the deployments named
|
||||
after its routed group or from a wildcard route; a priced route among those serves the
|
||||
request at its own rate whatever the group's price reads.
|
||||
"""
|
||||
for deployment in deployments:
|
||||
model_id = (deployment.get("model_info") or _EMPTY_COST_ENTRY).get("id")
|
||||
if model_id is None:
|
||||
continue
|
||||
raw_entry = litellm.model_cost.get(model_id, _EMPTY_COST_ENTRY)
|
||||
if _is_positive_cost(raw_entry.get("input_cost_per_token")) or _is_positive_cost(
|
||||
raw_entry.get("output_cost_per_token")
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
_EMPTY_COST_ENTRY: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
||||
|
||||
|
|
@ -672,7 +719,7 @@ def model_has_no_cost_mapping(model: str | None, llm_router: Router | None) -> b
|
|||
if _model_group_has_pricing(model=model, llm_router=llm_router):
|
||||
return False
|
||||
|
||||
return not _is_cost_explicitly_configured(model=model, llm_router=llm_router)
|
||||
return not _is_cost_explicitly_configured(llm_router.get_model_list(model_name=model) or ())
|
||||
|
||||
|
||||
def _unpriced_models_in_request(model: str | list[str] | None, llm_router: Router | None) -> tuple[str, ...]:
|
||||
|
|
|
|||
|
|
@ -1,20 +1,35 @@
|
|||
import base64
|
||||
import json
|
||||
import math
|
||||
import signal
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Iterator, Mapping
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass
|
||||
from hashlib import sha256
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import anthropic
|
||||
import httpx
|
||||
import jwt
|
||||
import openai
|
||||
import psutil
|
||||
import pytest
|
||||
import yaml
|
||||
from cryptography.hazmat.primitives.asymmetric import rsa
|
||||
from integration._support.anthropic_thinking import JSON_OBJECT
|
||||
from integration._support.client import Gateway, object_value
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.client import Gateway, Scenario, eventually, object_value
|
||||
from integration._support.database import write_rows
|
||||
from integration._support.process import group_members, owned_proxy, owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from integration.authorization._hidden_alias_budget import (
|
||||
BUDGET,
|
||||
BUDGET_EXCEEDED,
|
||||
CHAT_REPLY,
|
||||
GATEWAY_BURST,
|
||||
PEER_BURST,
|
||||
PROXY_BUDGET_USER,
|
||||
RESPONSES_REPLY,
|
||||
AliasRig,
|
||||
|
|
@ -29,9 +44,13 @@ from integration.authorization._hidden_alias_budget import (
|
|||
error_type,
|
||||
exhausted_key,
|
||||
fresh_chat,
|
||||
fresh_message,
|
||||
fresh_post,
|
||||
fresh_response,
|
||||
hidden,
|
||||
install_aliases,
|
||||
landed,
|
||||
landed_all_once,
|
||||
landed_once,
|
||||
openai_client,
|
||||
remove_aliases,
|
||||
|
|
@ -46,6 +65,11 @@ from pydantic import JsonValue
|
|||
|
||||
pytestmark: Final = pytest.mark.timeout(240)
|
||||
|
||||
# $0 comes from litellm's own cost-map row (model_prices_and_context_window.json, added 2025-10-14 in 900a57f5b1f):
|
||||
# Lemonade is a local inference server, so there is no vendor price for the row to drift from. A row that gains a
|
||||
# price is stale, not a regression: pick another chat row the map prices at $0 and serves over the OpenAI wire
|
||||
UNPRICED_FREE_PROVIDER_MODEL: Final = "lemonade/Gemma-3-4b-it-GGUF"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig() -> Iterator[AliasRig]:
|
||||
|
|
@ -506,6 +530,222 @@ def test_hidden_alias_to_a_group_priced_by_the_cost_map_stays_budgeted(rig: Alia
|
|||
assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 0
|
||||
|
||||
|
||||
def _reported_per_token_price(rig: AliasRig, group: str) -> tuple[float, float]:
|
||||
entries: Final = rig.gateway.get("/model/info")["data"]
|
||||
assert isinstance(entries, list), entries
|
||||
infos: Final = tuple(
|
||||
object_value(object_value(entry)["model_info"])
|
||||
for entry in entries
|
||||
if object_value(entry)["model_name"] == group
|
||||
)
|
||||
assert len(infos) == 1, infos
|
||||
return float(str(infos[0]["input_cost_per_token"])), float(str(infos[0]["output_cost_per_token"]))
|
||||
|
||||
|
||||
def _serving_deployment(candidate: Gateway, model: str) -> str | None:
|
||||
response: Final = fresh_chat(candidate, model, candidate.key, "route-" + uuid.uuid4().hex)
|
||||
return response.headers.get("x-litellm-model-id") if response.status_code == 200 else None
|
||||
|
||||
|
||||
def _routes_alike(candidate: Gateway, name: str, target: str) -> bool:
|
||||
served: Final = _serving_deployment(candidate, name)
|
||||
return served is not None and served == _serving_deployment(candidate, target)
|
||||
|
||||
|
||||
def _await_alias(candidate: Gateway, name: str, target: str) -> None:
|
||||
"""Wait until ``candidate`` serves ``name`` from ``target``'s deployment.
|
||||
|
||||
``/router/settings`` answers from the stored config before a replica's next reload applies it, and a
|
||||
shadowing alias's name answers 200 either way, so the serving deployment is the only signal on the wire.
|
||||
"""
|
||||
eventually(lambda: _routes_alike(candidate, name, target), bool, seconds=90)
|
||||
|
||||
|
||||
def _shadowing_alias(
|
||||
rig: AliasRig, scenario: Scenario, target: str, candidates: tuple[Gateway, ...], **shadow: JsonValue
|
||||
) -> str:
|
||||
"""A hidden alias whose name is also a real group, live on every candidate so the name routes to ``target``."""
|
||||
name: Final = scenario.model(**shadow)
|
||||
install_aliases(rig.gateway, {name: hidden(target)})
|
||||
scenario.cleanups.callback(remove_aliases, rig.gateway, frozenset({name}))
|
||||
for candidate in candidates:
|
||||
_await_alias(candidate, name, target)
|
||||
return name
|
||||
|
||||
|
||||
def _unpriced_free_target(rig: AliasRig, scenario: Scenario, provider_model: str = UNPRICED_FREE_PROVIDER_MODEL) -> str:
|
||||
target: Final = scenario.model(model=provider_model)
|
||||
assert _reported_per_token_price(rig, target) == (0.0, 0.0), (
|
||||
f"{provider_model} must price at $0 in the cost map with no explicit price on the deployment, "
|
||||
"or the shadowed alias is refused by the price gate and never reaches the one under test"
|
||||
)
|
||||
return target
|
||||
|
||||
|
||||
def _shadow_over_unpriced_target(rig: AliasRig, scenario: Scenario) -> tuple[str, str]:
|
||||
target: Final = _unpriced_free_target(rig, scenario)
|
||||
shadow: Final = _shadowing_alias(
|
||||
rig, scenario, target, (rig.gateway, rig.peer), input_cost_per_token=0, output_cost_per_token=0
|
||||
)
|
||||
return target, shadow
|
||||
|
||||
|
||||
def _refusal(response: httpx.Response) -> tuple[int, str]:
|
||||
return response.status_code, error_type(response)
|
||||
|
||||
|
||||
def _assert_refused_like_target(rig: AliasRig, candidate: Gateway, key: str, shadow: str, target: str) -> None:
|
||||
for route, send in (("chat", fresh_chat), ("responses", fresh_response), ("messages", fresh_message)):
|
||||
shadow_marker: Final = f"shadow-{route}-" + uuid.uuid4().hex
|
||||
target_marker: Final = f"target-{route}-" + uuid.uuid4().hex
|
||||
via_shadow: Final = send(candidate, shadow, key, shadow_marker)
|
||||
by_name: Final = send(candidate, target, key, target_marker)
|
||||
assert _refusal(by_name) == (BUDGET_EXCEEDED, "budget_exceeded"), by_name.text
|
||||
assert via_shadow.status_code == by_name.status_code, via_shadow.text
|
||||
assert error_type(via_shadow) == error_type(by_name), via_shadow.text
|
||||
observed: Final = upstream_requests(rig.gateway.upstream_url)
|
||||
assert upstream_hits(observed, shadow_marker) == 0
|
||||
assert upstream_hits(observed, target_marker) == 0
|
||||
|
||||
|
||||
def test_hidden_alias_shadowing_an_explicitly_free_group_is_refused_like_its_unpriced_target(rig: AliasRig) -> None:
|
||||
with rig.gateway.scenario() as scenario:
|
||||
target, shadow = _shadow_over_unpriced_target(rig, scenario)
|
||||
key: Final = exhausted_key(rig, scenario)
|
||||
settle_chat(rig, target, key, BUDGET_EXCEEDED)
|
||||
_assert_refused_like_target(rig, rig.gateway, key, shadow, target)
|
||||
_assert_refused_like_target(rig, rig.peer, key, shadow, target)
|
||||
|
||||
|
||||
def _assert_refused_on_every_route(rig: AliasRig, candidate: Gateway, key: str, model: str) -> None:
|
||||
for route, send in (("chat", fresh_chat), ("responses", fresh_response), ("messages", fresh_message)):
|
||||
marker: Final = f"chain-{route}-" + uuid.uuid4().hex
|
||||
refused: Final = send(candidate, model, key, marker)
|
||||
assert refused.status_code == BUDGET_EXCEEDED, refused.text
|
||||
assert error_type(refused) == "budget_exceeded", refused.text
|
||||
assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 0
|
||||
|
||||
|
||||
def test_hidden_alias_shadow_chain_is_judged_by_the_deployment_it_is_served_from(rig: AliasRig) -> None:
|
||||
"""A shadowing alias routes one hop to its target's own deployment, even when that target is an alias too."""
|
||||
with rig.gateway.scenario() as scenario:
|
||||
middle, shadow = _shadow_over_unpriced_target(rig, scenario)
|
||||
served_from: Final = _serving_deployment(rig.gateway, middle)
|
||||
install_aliases(rig.gateway, {middle: hidden(rig.free)})
|
||||
scenario.cleanups.callback(remove_aliases, rig.gateway, frozenset({middle}))
|
||||
for candidate in (rig.gateway, rig.peer):
|
||||
_await_alias(candidate, middle, rig.free)
|
||||
assert _serving_deployment(candidate, shadow) == served_from
|
||||
key: Final = exhausted_key(rig, scenario)
|
||||
for candidate in (rig.gateway, rig.peer):
|
||||
marker: Final = "chain-target-" + uuid.uuid4().hex
|
||||
by_name: Final = fresh_chat(candidate, middle, key, marker)
|
||||
assert by_name.status_code == 200, by_name.text
|
||||
assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 1
|
||||
_assert_refused_on_every_route(rig, candidate, key, shadow)
|
||||
|
||||
|
||||
def _await_route(candidate: Gateway, name: str, route: str) -> None:
|
||||
eventually(lambda: _serving_deployment(candidate, name), lambda served: served == route, seconds=90)
|
||||
|
||||
|
||||
def _deployment_named(rig: AliasRig, scenario: Scenario, name: str, **parameters: JsonValue) -> str:
|
||||
created: Final = rig.gateway.post(
|
||||
"/model/new",
|
||||
{
|
||||
"model_name": name,
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"api_key": "integration-provider-key",
|
||||
"api_base": f"{rig.gateway.upstream_url}/v1",
|
||||
**parameters,
|
||||
},
|
||||
"model_info": {},
|
||||
},
|
||||
)
|
||||
identity: Final = str(object_value(created["model_info"])["id"])
|
||||
scenario.cleanups.callback(scenario.delete_model, identity)
|
||||
return identity
|
||||
|
||||
|
||||
def _priced_wildcard_route(rig: AliasRig, scenario: Scenario, prefix: str) -> str:
|
||||
"""A wildcard deployment serving every ``<prefix>/...`` name at an explicit positive per-token price."""
|
||||
return _deployment_named(
|
||||
rig, scenario, f"{prefix}/*", model="openai/*", input_cost_per_token=0.00001, output_cost_per_token=0.00002
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("entry_on_pattern", [False, True], ids=["entry_off_pattern", "entry_on_pattern"])
|
||||
def test_hidden_alias_chain_served_by_a_priced_wildcard_route_is_refused(rig: AliasRig, entry_on_pattern: bool) -> None:
|
||||
"""A chain the router serves from a priced wildcard route is budgeted at that route's price.
|
||||
|
||||
The chain's middle name is only an alias key to the free group, so the group's price reads $0
|
||||
through it while the request is served by the wildcard route matching the middle name.
|
||||
"""
|
||||
prefix: Final = "wild" + uuid.uuid4().hex[:8]
|
||||
middle: Final = f"{prefix}/gpt-4o-mini"
|
||||
entry: Final = (f"{prefix}/" if entry_on_pattern else "chain-entry-") + uuid.uuid4().hex
|
||||
with rig.gateway.scenario() as scenario:
|
||||
route: Final = _priced_wildcard_route(rig, scenario, prefix)
|
||||
install_aliases(rig.gateway, {entry: hidden(middle), middle: hidden(rig.free)})
|
||||
scenario.cleanups.callback(remove_aliases, rig.gateway, frozenset({entry, middle}))
|
||||
for candidate in (rig.gateway, rig.peer):
|
||||
_await_route(candidate, entry, route)
|
||||
_await_alias(candidate, middle, rig.free)
|
||||
key: Final = exhausted_key(rig, scenario)
|
||||
settle_chat(rig, rig.paid, key, BUDGET_EXCEEDED)
|
||||
for candidate in (rig.gateway, rig.peer):
|
||||
_assert_refused_on_every_route(rig, candidate, key, entry)
|
||||
|
||||
|
||||
def _ptu_shadow_config(rig: AliasRig, directory: Path, name: str) -> Path:
|
||||
base: Final = _base_config()["model_list"]
|
||||
assert isinstance(base, list), base
|
||||
deployment: Final[JsonValue] = {
|
||||
"model_name": name,
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"api_key": "integration-provider-key",
|
||||
"api_base": f"{rig.gateway.upstream_url}/v1",
|
||||
"input_cost_per_token": 0,
|
||||
"output_cost_per_token": 0,
|
||||
},
|
||||
"model_info": {"ptu_count": 100, "cost_per_ptu_per_hour": 2.0},
|
||||
}
|
||||
return _own_config(directory, "ptu-shadow.yaml", "model_list", [*base, deployment])
|
||||
|
||||
|
||||
@pytest.mark.timeout(480)
|
||||
def test_hidden_alias_shadowing_a_ptu_group_is_served_through_its_free_target(rig: AliasRig, tmp_path: Path) -> None:
|
||||
shadow: Final = "ptu-shadow-" + uuid.uuid4().hex
|
||||
marker: Final = "ptu-shadow-served-" + uuid.uuid4().hex
|
||||
config: Final = _ptu_shadow_config(rig, tmp_path, shadow)
|
||||
with rig.gateway.scenario() as scenario:
|
||||
install_aliases(rig.gateway, {shadow: hidden(rig.free)})
|
||||
scenario.cleanups.callback(remove_aliases, rig.gateway, frozenset({shadow}))
|
||||
key: Final = exhausted_key(rig, scenario)
|
||||
with owned_proxy(rig.gateway, tmp_path, {}, config=config) as candidate:
|
||||
_await_alias(candidate, shadow, rig.free)
|
||||
settle_candidate(candidate, rig.paid, key, BUDGET_EXCEEDED, seconds=120)
|
||||
served: Final = fresh_chat(candidate, shadow, key, marker)
|
||||
assert served.status_code == 200, served.text
|
||||
assert served.json()["choices"][0]["message"]["content"] == CHAT_REPLY, served.text
|
||||
assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 1
|
||||
assert_free_row(landed_once(key, marker), shadow)
|
||||
|
||||
|
||||
def test_hidden_alias_shadowing_an_explicitly_free_group_to_a_paid_target_stays_budgeted(rig: AliasRig) -> None:
|
||||
marker: Final = "paid-shadow-" + uuid.uuid4().hex
|
||||
with rig.gateway.scenario() as scenario:
|
||||
shadow: Final = _shadowing_alias(
|
||||
rig, scenario, rig.paid, (rig.gateway, rig.peer), input_cost_per_token=0, output_cost_per_token=0
|
||||
)
|
||||
key: Final = exhausted_key(rig, scenario)
|
||||
refused: Final = fresh_chat(rig.gateway, shadow, key, marker)
|
||||
assert _refusal(refused) == (BUDGET_EXCEEDED, "budget_exceeded"), refused.text
|
||||
assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 0
|
||||
|
||||
|
||||
def test_hidden_alias_to_a_missing_group_is_refused_and_the_proxy_stays_healthy(rig: AliasRig) -> None:
|
||||
with rig.gateway.scenario() as scenario:
|
||||
key: Final = exhausted_key(rig, scenario)
|
||||
|
|
@ -575,3 +815,578 @@ def _assert_hidden_aliases_unlisted(rig: AliasRig, candidate: Gateway) -> None:
|
|||
def test_hidden_alias_stays_out_of_model_listings(rig: AliasRig) -> None:
|
||||
_assert_hidden_aliases_unlisted(rig, rig.gateway)
|
||||
_assert_hidden_aliases_unlisted(rig, rig.peer)
|
||||
|
||||
|
||||
def _deployment_id(rig: AliasRig, name: str) -> str:
|
||||
entries: Final = rig.gateway.get("/model/info")["data"]
|
||||
assert isinstance(entries, list), entries
|
||||
identities: Final = tuple(
|
||||
str(object_value(object_value(entry)["model_info"])["id"])
|
||||
for entry in entries
|
||||
if object_value(entry)["model_name"] == name
|
||||
)
|
||||
assert len(identities) == 1, identities
|
||||
return identities[0]
|
||||
|
||||
|
||||
def _assert_budget_refusal(status_code: int, body: str) -> None:
|
||||
assert status_code == BUDGET_EXCEEDED, body
|
||||
assert "budget_exceeded" in body, body
|
||||
|
||||
|
||||
def _assert_no_served_row(key: str, marker: str) -> None:
|
||||
assert all(row["status"] != "success" for row in landed(key, marker)), landed(key, marker)
|
||||
|
||||
|
||||
def test_exhausted_key_is_refused_on_shadowing_alias_through_streamed_openai_chat(rig: AliasRig) -> None:
|
||||
marker: Final = "shadow-chat-stream-" + uuid.uuid4().hex
|
||||
with rig.gateway.scenario() as scenario:
|
||||
_, shadow = _shadow_over_unpriced_target(rig, scenario)
|
||||
key: Final = exhausted_key(rig, scenario)
|
||||
with openai_client(rig.gateway, key) as client, pytest.raises(openai.APIStatusError) as refused:
|
||||
client.chat.completions.create(
|
||||
model=shadow,
|
||||
messages=[{"role": "user", "content": marker}],
|
||||
stream=True,
|
||||
extra_headers=spend_marker(marker),
|
||||
)
|
||||
_assert_budget_refusal(refused.value.status_code, str(refused.value.body))
|
||||
assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 0
|
||||
|
||||
|
||||
async def test_exhausted_key_is_refused_on_shadowing_alias_through_streamed_openai_responses(rig: AliasRig) -> None:
|
||||
marker: Final = "shadow-responses-stream-" + uuid.uuid4().hex
|
||||
with rig.gateway.scenario() as scenario:
|
||||
_, shadow = _shadow_over_unpriced_target(rig, scenario)
|
||||
key: Final = exhausted_key(rig, scenario)
|
||||
async with async_openai_client(rig.gateway, key) as client:
|
||||
with pytest.raises(openai.APIStatusError) as refused:
|
||||
await client.responses.create(
|
||||
model=shadow, input=marker, stream=True, extra_headers=spend_marker(marker)
|
||||
)
|
||||
_assert_budget_refusal(refused.value.status_code, str(refused.value.body))
|
||||
assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 0
|
||||
|
||||
|
||||
async def test_exhausted_key_is_refused_on_shadowing_alias_through_streamed_anthropic_messages(rig: AliasRig) -> None:
|
||||
marker: Final = "shadow-messages-stream-" + uuid.uuid4().hex
|
||||
with rig.gateway.scenario() as scenario:
|
||||
_, shadow = _shadow_over_unpriced_target(rig, scenario)
|
||||
key: Final = exhausted_key(rig, scenario)
|
||||
async with async_anthropic_client(rig.gateway, key) as client:
|
||||
with pytest.raises(anthropic.APIStatusError) as refused:
|
||||
async with client.messages.stream(
|
||||
model=shadow,
|
||||
max_tokens=16,
|
||||
messages=[{"role": "user", "content": marker}],
|
||||
extra_headers=spend_marker(marker),
|
||||
) as stream:
|
||||
await stream.get_final_message()
|
||||
_assert_budget_refusal(refused.value.status_code, str(refused.value.body))
|
||||
assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 0
|
||||
|
||||
|
||||
def _await_served_within(candidate: Gateway, name: str, deployments: frozenset[str]) -> None:
|
||||
eventually(lambda: _serving_deployment(candidate, name), lambda served: served in deployments, seconds=90)
|
||||
|
||||
|
||||
def _assert_served_by(
|
||||
rig: AliasRig, candidate: Gateway, name: str, key: str, deployments: frozenset[str], prefix: str
|
||||
) -> str:
|
||||
marker: Final = f"{prefix}-" + uuid.uuid4().hex
|
||||
served: Final = fresh_chat(candidate, name, key, marker)
|
||||
assert served.status_code == 200, served.text
|
||||
assert served.headers["x-litellm-model-id"] in deployments, dict(served.headers)
|
||||
assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 1
|
||||
response_id: Final = str(JSON_OBJECT.validate_json(served.content)["id"])
|
||||
row: Final = landed_once(key, marker)
|
||||
assert row["request_id"] == response_id, row
|
||||
assert_free_row(row, name)
|
||||
return response_id
|
||||
|
||||
|
||||
def test_key_with_headroom_is_served_by_the_target_through_a_shadowing_alias(rig: AliasRig) -> None:
|
||||
with rig.gateway.scenario() as scenario:
|
||||
target, shadow = _shadow_over_unpriced_target(rig, scenario)
|
||||
key: Final = scenario.key(max_budget=BUDGET)
|
||||
for candidate in (rig.gateway, rig.peer):
|
||||
_assert_served_by(rig, candidate, shadow, key, frozenset({_deployment_id(rig, target)}), "headroom-shadow")
|
||||
|
||||
|
||||
def test_exhausted_key_reaches_hidden_alias_to_a_name_served_by_a_free_wildcard_route(rig: AliasRig) -> None:
|
||||
prefix: Final = "freewild" + uuid.uuid4().hex[:8]
|
||||
entry: Final = "free-wildcard-entry-" + uuid.uuid4().hex
|
||||
with rig.gateway.scenario() as scenario:
|
||||
route: Final = _deployment_named(
|
||||
rig, scenario, f"{prefix}/*", model="openai/*", input_cost_per_token=0, output_cost_per_token=0
|
||||
)
|
||||
install_aliases(rig.gateway, {entry: hidden(f"{prefix}/gpt-4o-mini")})
|
||||
scenario.cleanups.callback(remove_aliases, rig.gateway, frozenset({entry}))
|
||||
for candidate in (rig.gateway, rig.peer):
|
||||
_await_route(candidate, entry, route)
|
||||
key: Final = exhausted_key(rig, scenario)
|
||||
for candidate in (rig.gateway, rig.peer):
|
||||
_assert_served_by(rig, candidate, entry, key, frozenset({route}), "free-wildcard")
|
||||
|
||||
|
||||
def test_exhausted_key_reaches_a_group_with_one_explicitly_free_deployment(rig: AliasRig) -> None:
|
||||
group: Final = "mixed-free-" + uuid.uuid4().hex
|
||||
entry: Final = "mixed-free-entry-" + uuid.uuid4().hex
|
||||
with rig.gateway.scenario() as scenario:
|
||||
members: Final = frozenset(
|
||||
{
|
||||
_deployment_named(rig, scenario, group, input_cost_per_token=0, output_cost_per_token=0),
|
||||
_deployment_named(rig, scenario, group, model=UNPRICED_FREE_PROVIDER_MODEL),
|
||||
}
|
||||
)
|
||||
install_aliases(rig.gateway, {entry: hidden(group)})
|
||||
scenario.cleanups.callback(remove_aliases, rig.gateway, frozenset({entry}))
|
||||
for candidate in (rig.gateway, rig.peer):
|
||||
_await_served_within(candidate, entry, members)
|
||||
key: Final = exhausted_key(rig, scenario)
|
||||
for candidate in (rig.gateway, rig.peer):
|
||||
for name in (group, entry):
|
||||
_assert_served_by(rig, candidate, name, key, members, "mixed")
|
||||
|
||||
|
||||
def _script_target(rig: AliasRig, statuses: tuple[int, ...]) -> None:
|
||||
provider_model: Final = UNPRICED_FREE_PROVIDER_MODEL.split("/", 1)[1]
|
||||
scripted: Final = httpx.post(
|
||||
f"{rig.gateway.upstream_url}/__scripts/{provider_model}",
|
||||
json={"statuses": list(statuses)},
|
||||
timeout=15,
|
||||
trust_env=False,
|
||||
)
|
||||
assert scripted.status_code == 200, scripted.text
|
||||
|
||||
|
||||
def _clear_target_script(rig: AliasRig) -> None:
|
||||
provider_model: Final = UNPRICED_FREE_PROVIDER_MODEL.split("/", 1)[1]
|
||||
cleared: Final = httpx.delete(f"{rig.gateway.upstream_url}/__scripts/{provider_model}", timeout=15, trust_env=False)
|
||||
assert cleared.status_code in (200, 404), cleared.text
|
||||
|
||||
|
||||
def _assert_provider_failure(response: httpx.Response) -> None:
|
||||
assert response.status_code == 500, response.text
|
||||
assert "Controlled provider failure" in response.text, response.text
|
||||
assert "budget" not in response.text.lower(), response.text
|
||||
|
||||
|
||||
def test_provider_failure_behind_a_shadowing_alias_reaches_the_caller(rig: AliasRig) -> None:
|
||||
failed_marker: Final = "shadow-provider-failure-" + uuid.uuid4().hex
|
||||
with rig.gateway.scenario() as scenario:
|
||||
target, shadow = _shadow_over_unpriced_target(rig, scenario)
|
||||
healthy: Final = scenario.key(max_budget=BUDGET)
|
||||
exhausted: Final = exhausted_key(rig, scenario)
|
||||
settle_chat(rig, target, exhausted, BUDGET_EXCEEDED)
|
||||
_assert_refused_like_target(rig, rig.gateway, exhausted, shadow, target)
|
||||
scenario.cleanups.callback(_clear_target_script, rig)
|
||||
_script_target(rig, (500,))
|
||||
_assert_provider_failure(fresh_chat(rig.gateway, shadow, healthy, failed_marker))
|
||||
assert upstream_hits(upstream_requests(rig.gateway.upstream_url), failed_marker) == 1
|
||||
_clear_target_script(rig)
|
||||
_assert_served_by(
|
||||
rig, rig.gateway, shadow, healthy, frozenset({_deployment_id(rig, target)}), "shadow-provider-recovered"
|
||||
)
|
||||
_assert_no_served_row(healthy, failed_marker)
|
||||
_assert_refused_like_target(rig, rig.gateway, exhausted, shadow, target)
|
||||
|
||||
|
||||
def test_unauthenticated_request_to_a_shadowing_alias_is_rejected(rig: AliasRig) -> None:
|
||||
marker: Final = "shadow-unauthenticated-" + uuid.uuid4().hex
|
||||
with rig.gateway.scenario() as scenario:
|
||||
_, shadow = _shadow_over_unpriced_target(rig, scenario)
|
||||
response: Final = httpx.post(
|
||||
f"{base_url(rig.gateway)}/v1/chat/completions",
|
||||
json={"model": shadow, "messages": [{"role": "user", "content": marker}]},
|
||||
timeout=60,
|
||||
trust_env=False,
|
||||
)
|
||||
assert response.status_code == 401, response.text
|
||||
assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 0
|
||||
|
||||
|
||||
# $0 comes from litellm's own cost-map row (model_prices_and_context_window.json, added 2025-12-08 in #17692). The cell
|
||||
# needs an embedding row priced $0 only by the map and served over the OpenAI wire; a row that gains a price is stale,
|
||||
# not a regression: pick another such row
|
||||
UNPRICED_FREE_EMBEDDING_MODEL: Final = "fireworks_ai/accounts/fireworks/models/qwen3-embedding-0p6b"
|
||||
|
||||
|
||||
def _embed(candidate: Gateway, model: str, key: str, marker: str) -> httpx.Response:
|
||||
return fresh_post(candidate, "/v1/embeddings", {"model": model, "input": marker}, key, marker)
|
||||
|
||||
|
||||
def _assert_embeddings_judged_like_target(
|
||||
candidate: Gateway, shadow: str, target: str, healthy: str, exhausted: str
|
||||
) -> None:
|
||||
marker: Final = "embed-shadow-" + uuid.uuid4().hex
|
||||
served: Final = _embed(candidate, shadow, healthy, marker)
|
||||
assert served.status_code == 200, served.text
|
||||
vectors: Final = JSON_OBJECT.validate_json(served.content)["data"]
|
||||
assert isinstance(vectors, list) and len(vectors) == 1, served.text
|
||||
assert_free_row(landed_once(healthy, marker), shadow)
|
||||
via_shadow: Final = _embed(candidate, shadow, exhausted, "embed-refused-shadow-" + uuid.uuid4().hex)
|
||||
by_name: Final = _embed(candidate, target, exhausted, "embed-refused-target-" + uuid.uuid4().hex)
|
||||
assert _refusal(by_name) == (BUDGET_EXCEEDED, "budget_exceeded"), by_name.text
|
||||
assert _refusal(via_shadow) == _refusal(by_name), via_shadow.text
|
||||
|
||||
|
||||
def test_exhausted_key_is_refused_on_a_shadowing_alias_through_embeddings(rig: AliasRig) -> None:
|
||||
with rig.gateway.scenario() as scenario:
|
||||
target: Final = _unpriced_free_target(rig, scenario, UNPRICED_FREE_EMBEDDING_MODEL)
|
||||
shadow: Final = _shadowing_alias(
|
||||
rig, scenario, target, (rig.gateway, rig.peer), input_cost_per_token=0, output_cost_per_token=0
|
||||
)
|
||||
healthy: Final = scenario.key(max_budget=BUDGET)
|
||||
exhausted: Final = exhausted_key(rig, scenario)
|
||||
for candidate in (rig.gateway, rig.peer):
|
||||
_assert_embeddings_judged_like_target(candidate, shadow, target, healthy, exhausted)
|
||||
|
||||
|
||||
def test_failed_primary_does_not_fall_back_to_a_shadowing_alias_for_exhausted_key(rig: AliasRig) -> None:
|
||||
marker: Final = "fallback-shadow-" + uuid.uuid4().hex
|
||||
with rig.gateway.scenario() as scenario:
|
||||
_, shadow = _shadow_over_unpriced_target(rig, scenario)
|
||||
key: Final = exhausted_key(rig, scenario)
|
||||
scenario.cleanups.callback(clear_provider_script, rig)
|
||||
script_provider(rig, 1)
|
||||
_assert_provider_failure(fresh_chat(rig.gateway, rig.failing_free, key, marker, {"fallbacks": [shadow]}))
|
||||
assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 1
|
||||
_assert_no_served_row(key, marker)
|
||||
|
||||
|
||||
def test_removing_a_shadowing_alias_restores_the_explicitly_free_deployment(rig: AliasRig) -> None:
|
||||
with rig.gateway.scenario() as scenario:
|
||||
target, shadow = _shadow_over_unpriced_target(rig, scenario)
|
||||
real: Final = _deployment_id(rig, shadow)
|
||||
key: Final = exhausted_key(rig, scenario)
|
||||
settle_chat(rig, target, key, BUDGET_EXCEEDED)
|
||||
for candidate in (rig.gateway, rig.peer):
|
||||
_assert_refused_like_target(rig, candidate, key, shadow, target)
|
||||
remove_aliases(rig.gateway, frozenset({shadow}))
|
||||
for candidate in (rig.gateway, rig.peer):
|
||||
_await_route(candidate, shadow, real)
|
||||
_assert_served_by(rig, candidate, shadow, key, frozenset({real}), "restored")
|
||||
|
||||
|
||||
ZERO_MODEL_BUDGET: Final[JsonValue] = {"budget_limit": 0, "time_period": "1d"}
|
||||
MODEL_BUDGET_SQL: Final = 'UPDATE "LiteLLM_VerificationToken" SET model_max_budget = %s::jsonb WHERE token = %s'
|
||||
|
||||
|
||||
def _key_with_zero_model_budgets(scenario: Scenario, names: tuple[str, ...]) -> str:
|
||||
key: Final = scenario.key(max_budget=BUDGET)
|
||||
budgets: Final = json.dumps({name: ZERO_MODEL_BUDGET for name in names})
|
||||
write_rows(MODEL_BUDGET_SQL, (budgets, sha256(key.encode()).hexdigest()))
|
||||
return key
|
||||
|
||||
|
||||
def _assert_model_budget_refused_alike(rig: AliasRig, candidate: Gateway, key: str, shadow: str, target: str) -> None:
|
||||
shadow_marker: Final = "model-budget-shadow-" + uuid.uuid4().hex
|
||||
target_marker: Final = "model-budget-target-" + uuid.uuid4().hex
|
||||
via_shadow: Final = fresh_chat(candidate, shadow, key, shadow_marker)
|
||||
by_name: Final = fresh_chat(candidate, target, key, target_marker)
|
||||
assert by_name.status_code == BUDGET_EXCEEDED, by_name.text
|
||||
assert _refusal(via_shadow) == _refusal(by_name), via_shadow.text
|
||||
observed: Final = upstream_requests(rig.gateway.upstream_url)
|
||||
assert upstream_hits(observed, shadow_marker) == 0
|
||||
assert upstream_hits(observed, target_marker) == 0
|
||||
|
||||
|
||||
def test_per_model_budget_on_a_shadowing_alias_is_judged_like_its_target(rig: AliasRig) -> None:
|
||||
with rig.gateway.scenario() as scenario:
|
||||
target, shadow = _shadow_over_unpriced_target(rig, scenario)
|
||||
key: Final = _key_with_zero_model_budgets(scenario, (shadow, target))
|
||||
for candidate in (rig.gateway, rig.peer):
|
||||
_assert_model_budget_refused_alike(rig, candidate, key, shadow, target)
|
||||
|
||||
|
||||
def test_exhausted_key_is_not_served_a_cached_reply_through_a_shadowing_alias(rig: AliasRig) -> None:
|
||||
marker: Final = "shadow-cache-twin-" + uuid.uuid4().hex
|
||||
with rig.gateway.scenario() as scenario:
|
||||
_, shadow = _shadow_over_unpriced_target(rig, scenario)
|
||||
healthy: Final = scenario.key(max_budget=BUDGET)
|
||||
exhausted: Final = exhausted_key(rig, scenario)
|
||||
first: Final = fresh_chat(rig.gateway, shadow, healthy, marker)
|
||||
assert first.status_code == 200, first.text
|
||||
warmed: Final = fresh_chat(rig.gateway, shadow, healthy, marker)
|
||||
assert warmed.status_code == 200, warmed.text
|
||||
assert warmed.json()["id"] == first.json()["id"], warmed.text
|
||||
assert "x-litellm-cache-key" in warmed.headers, dict(warmed.headers)
|
||||
refused: Final = fresh_chat(rig.gateway, shadow, exhausted, marker)
|
||||
_assert_budget_refusal(refused.status_code, refused.text)
|
||||
assert "x-litellm-cache-key" not in refused.headers, dict(refused.headers)
|
||||
assert upstream_hits(upstream_requests(rig.gateway.upstream_url), marker) == 1
|
||||
_assert_no_served_row(exhausted, marker)
|
||||
|
||||
|
||||
BUDGET_VERDICTS: Final = frozenset({200, BUDGET_EXCEEDED})
|
||||
|
||||
|
||||
def _only_budget_verdicts(seen: frozenset[int]) -> bool:
|
||||
assert seen <= BUDGET_VERDICTS, seen
|
||||
return seen == frozenset({200})
|
||||
|
||||
|
||||
def test_shadowing_alias_removed_under_traffic_keeps_every_answer_a_budget_verdict(rig: AliasRig) -> None:
|
||||
with rig.gateway.scenario() as scenario:
|
||||
target, shadow = _shadow_over_unpriced_target(rig, scenario)
|
||||
real: Final = _deployment_id(rig, shadow)
|
||||
key: Final = exhausted_key(rig, scenario)
|
||||
settle_chat(rig, shadow, key, BUDGET_EXCEEDED)
|
||||
remove_aliases(rig.gateway, frozenset({shadow}))
|
||||
eventually(lambda: chat_statuses(rig.gateway, shadow, key, GATEWAY_BURST), _only_budget_verdicts, seconds=90)
|
||||
eventually(lambda: chat_statuses(rig.peer, shadow, key, PEER_BURST), _only_budget_verdicts, seconds=90)
|
||||
for candidate in (rig.gateway, rig.peer):
|
||||
_assert_served_by(rig, candidate, shadow, key, frozenset({real}), "after-removal")
|
||||
assert _serving_deployment(rig.gateway, target) == _deployment_id(rig, target)
|
||||
|
||||
|
||||
JWT_KEY_ID: Final = "integration-shadow-jwt-key"
|
||||
|
||||
|
||||
def _jwt_config(directory: Path) -> Path:
|
||||
general: Final = object_value(_base_config()["general_settings"])
|
||||
settings: Final[JsonValue] = {
|
||||
**general,
|
||||
"enable_jwt_auth": True,
|
||||
"litellm_jwtauth": {"team_id_jwt_field": "team_id"},
|
||||
}
|
||||
return _own_config(directory, "shadow-jwt.yaml", "general_settings", settings)
|
||||
|
||||
|
||||
def _jwk_integer(value: int) -> str:
|
||||
return base64.urlsafe_b64encode(value.to_bytes((value.bit_length() + 7) // 8, "big")).rstrip(b"=").decode()
|
||||
|
||||
|
||||
def _jwks(private_key: rsa.RSAPrivateKey) -> bytes:
|
||||
numbers: Final = private_key.public_key().public_numbers()
|
||||
key: Final = {"kty": "RSA", "use": "sig", "alg": "RS256", "kid": JWT_KEY_ID}
|
||||
return json.dumps({"keys": [{**key, "n": _jwk_integer(numbers.n), "e": _jwk_integer(numbers.e)}]}).encode()
|
||||
|
||||
|
||||
@pytest.mark.timeout(480)
|
||||
def test_jwt_caller_over_team_budget_is_refused_on_a_shadowing_alias(rig: AliasRig, tmp_path: Path) -> None:
|
||||
private_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048)
|
||||
jwks_body: Final = _jwks(private_key)
|
||||
|
||||
def respond(request: Request) -> Reply:
|
||||
assert request.method == "GET", request
|
||||
return Reply(body=jwks_body)
|
||||
|
||||
with wire_server(respond) as jwks, rig.gateway.scenario() as scenario:
|
||||
target, shadow = _shadow_over_unpriced_target(rig, scenario)
|
||||
team: Final = scenario.team(max_budget=BUDGET)
|
||||
issued: Final = int(time.time())
|
||||
token: Final = jwt.encode(
|
||||
{"sub": "integration-shadow-jwt-caller", "team_id": team, "iat": issued, "exp": issued + 900},
|
||||
private_key,
|
||||
algorithm="RS256",
|
||||
headers={"kid": JWT_KEY_ID},
|
||||
)
|
||||
with owned_proxy(
|
||||
rig.gateway, tmp_path, {"JWT_PUBLIC_KEY_URL": jwks.url}, config=_jwt_config(tmp_path)
|
||||
) as candidate:
|
||||
_await_alias(candidate, shadow, target)
|
||||
first: Final = fresh_chat(candidate, rig.paid, token, "jwt-exhaust-" + uuid.uuid4().hex)
|
||||
assert first.status_code == 200, first.text
|
||||
settle_candidate(candidate, rig.paid, token, BUDGET_EXCEEDED, seconds=120)
|
||||
settle_candidate(candidate, target, token, BUDGET_EXCEEDED, seconds=120)
|
||||
_assert_refused_like_target(rig, candidate, token, shadow, target)
|
||||
free_marker: Final = "jwt-free-" + uuid.uuid4().hex
|
||||
served: Final = fresh_chat(candidate, rig.hidden_free, token, free_marker)
|
||||
assert served.status_code == 200, served.text
|
||||
assert upstream_hits(upstream_requests(rig.gateway.upstream_url), free_marker) == 1
|
||||
|
||||
|
||||
CUSTOM_AUTH_DOUBLE: Final = """
|
||||
import os
|
||||
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
||||
|
||||
|
||||
async def user_api_key_auth(request: Request, api_key: str) -> UserAPIKeyAuth:
|
||||
if api_key == os.environ["LITELLM_MASTER_KEY"]:
|
||||
return UserAPIKeyAuth(api_key=api_key, user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
if api_key != os.environ["INTEGRATION_CUSTOM_AUTH_KEY"]:
|
||||
raise Exception("unknown caller")
|
||||
budgets = {
|
||||
name: {"budget_limit": 0, "time_period": "1d"}
|
||||
for name in os.environ["INTEGRATION_ZERO_BUDGET_MODELS"].split(",")
|
||||
}
|
||||
return UserAPIKeyAuth(api_key=api_key, user_id="integration-custom-caller", model_max_budget=budgets)
|
||||
"""
|
||||
|
||||
|
||||
def _custom_auth_config(directory: Path) -> Path:
|
||||
(directory / "custom_auth_double.py").write_text(CUSTOM_AUTH_DOUBLE)
|
||||
base: Final = _base_config()
|
||||
general: Final = object_value(base["general_settings"])
|
||||
litellm_settings: Final = object_value(base["litellm_settings"])
|
||||
path: Final = directory / "shadow-custom-auth.yaml"
|
||||
path.write_text(
|
||||
yaml.safe_dump(
|
||||
{
|
||||
**base,
|
||||
"general_settings": {**general, "custom_auth": "custom_auth_double.user_api_key_auth"},
|
||||
"litellm_settings": {**litellm_settings, "enable_post_custom_auth_checks": True},
|
||||
}
|
||||
)
|
||||
)
|
||||
return path
|
||||
|
||||
|
||||
@pytest.mark.timeout(480)
|
||||
def test_custom_auth_caller_with_zero_model_budgets_is_refused_on_a_shadowing_alias(
|
||||
rig: AliasRig, tmp_path: Path
|
||||
) -> None:
|
||||
caller: Final = "sk-custom-" + uuid.uuid4().hex
|
||||
with rig.gateway.scenario() as scenario:
|
||||
target, shadow = _shadow_over_unpriced_target(rig, scenario)
|
||||
overrides: Final = {
|
||||
"INTEGRATION_CUSTOM_AUTH_KEY": caller,
|
||||
"INTEGRATION_ZERO_BUDGET_MODELS": ",".join((shadow, target, rig.free)),
|
||||
}
|
||||
with owned_proxy(rig.gateway, tmp_path, overrides, config=_custom_auth_config(tmp_path)) as candidate:
|
||||
_await_alias(candidate, shadow, target)
|
||||
settle_candidate(candidate, target, caller, BUDGET_EXCEEDED, seconds=120)
|
||||
_assert_refused_like_target(rig, candidate, caller, shadow, target)
|
||||
free_marker: Final = "custom-free-" + uuid.uuid4().hex
|
||||
served: Final = fresh_chat(candidate, rig.free, caller, free_marker)
|
||||
assert served.status_code == 200, served.text
|
||||
assert upstream_hits(upstream_requests(rig.gateway.upstream_url), free_marker) == 1
|
||||
unknown: Final = fresh_chat(candidate, shadow, "sk-unknown-" + uuid.uuid4().hex, "custom-unknown")
|
||||
assert unknown.status_code == 401, unknown.text
|
||||
|
||||
|
||||
BURST_SERVED: Final = 40
|
||||
BURST_FAILURES: Final = 8
|
||||
BURST_PROBES: Final = 8
|
||||
BURST_KINDS: Final = ("chat", "chat-stream", "responses", "responses-stream", "messages", "messages-stream")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _BurstCall:
|
||||
kind: str
|
||||
marker: str
|
||||
|
||||
|
||||
def _burst_calls() -> tuple[_BurstCall, ...]:
|
||||
served: Final = tuple(
|
||||
_BurstCall(BURST_KINDS[index % len(BURST_KINDS)], "burst-" + uuid.uuid4().hex) for index in range(BURST_SERVED)
|
||||
)
|
||||
probes: Final = tuple(_BurstCall("probe", "probe-" + uuid.uuid4().hex) for _ in range(BURST_PROBES))
|
||||
liveliness: Final = tuple(_BurstCall("liveliness", "live-" + uuid.uuid4().hex) for _ in range(BURST_PROBES))
|
||||
return served + probes + liveliness
|
||||
|
||||
|
||||
def _send_burst_call(rig: AliasRig, shadow: str, healthy: str, exhausted: str, call: _BurstCall) -> httpx.Response:
|
||||
stream: Final[Mapping[str, JsonValue]] = {"stream": True}
|
||||
match call.kind:
|
||||
case "chat":
|
||||
return fresh_chat(rig.gateway, shadow, healthy, call.marker)
|
||||
case "chat-stream":
|
||||
return fresh_chat(rig.gateway, shadow, healthy, call.marker, stream)
|
||||
case "responses":
|
||||
return fresh_response(rig.gateway, shadow, healthy, call.marker)
|
||||
case "responses-stream":
|
||||
return fresh_response(rig.gateway, shadow, healthy, call.marker, stream)
|
||||
case "messages":
|
||||
return fresh_message(rig.gateway, shadow, healthy, call.marker)
|
||||
case "messages-stream":
|
||||
return fresh_message(rig.gateway, shadow, healthy, call.marker, stream)
|
||||
case "probe":
|
||||
return fresh_chat(rig.gateway, shadow, exhausted, call.marker)
|
||||
case "liveliness":
|
||||
return httpx.get(f"{base_url(rig.gateway)}/health/liveliness", timeout=60, trust_env=False)
|
||||
case _:
|
||||
raise AssertionError(call.kind)
|
||||
|
||||
|
||||
def _served_calls(responses: Mapping[_BurstCall, httpx.Response]) -> Mapping[_BurstCall, httpx.Response]:
|
||||
return {call: response for call, response in responses.items() if call.kind in BURST_KINDS}
|
||||
|
||||
|
||||
def _calls_of_kind(responses: Mapping[_BurstCall, httpx.Response], kind: str) -> tuple[_BurstCall, ...]:
|
||||
return tuple(call for call in responses if call.kind == kind)
|
||||
|
||||
|
||||
def test_shadowing_alias_burst_with_provider_outage_lands_every_served_id_once(rig: AliasRig) -> None:
|
||||
with rig.gateway.scenario() as scenario:
|
||||
_, shadow = _shadow_over_unpriced_target(rig, scenario)
|
||||
healthy: Final = scenario.key(max_budget=BUDGET)
|
||||
exhausted: Final = exhausted_key(rig, scenario)
|
||||
scenario.cleanups.callback(_clear_target_script, rig)
|
||||
_script_target(rig, (500,) * BURST_FAILURES + (200,) * BURST_SERVED)
|
||||
calls: Final = _burst_calls()
|
||||
|
||||
def send(call: _BurstCall) -> httpx.Response:
|
||||
return _send_burst_call(rig, shadow, healthy, exhausted, call)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=len(calls)) as pool:
|
||||
responses: Final = dict(zip(calls, pool.map(send, calls), strict=True))
|
||||
served: Final = _served_calls(responses)
|
||||
failed: Final = {call: response for call, response in served.items() if response.status_code != 200}
|
||||
assert len(failed) == BURST_FAILURES, {call.kind: response.status_code for call, response in failed.items()}
|
||||
for response in failed.values():
|
||||
_assert_provider_failure(response)
|
||||
for call in _calls_of_kind(responses, "probe"):
|
||||
_assert_budget_refusal(responses[call].status_code, responses[call].text)
|
||||
for call in _calls_of_kind(responses, "liveliness"):
|
||||
assert responses[call].status_code == 200, responses[call].text
|
||||
served_markers: Final = frozenset(call.marker for call in served if call not in failed)
|
||||
assert len(served_markers) == BURST_SERVED - BURST_FAILURES
|
||||
observed: Final = upstream_requests(rig.gateway.upstream_url)
|
||||
assert all(upstream_hits(observed, call.marker) == 1 for call in served)
|
||||
assert all(upstream_hits(observed, call.marker) == 0 for call in _calls_of_kind(responses, "probe"))
|
||||
for row in landed_all_once(healthy, served_markers):
|
||||
assert_free_row(row, shadow)
|
||||
for call in failed:
|
||||
_assert_no_served_row(healthy, call.marker)
|
||||
for call in _calls_of_kind(responses, "probe"):
|
||||
_assert_no_served_row(exhausted, call.marker)
|
||||
|
||||
|
||||
def _is_worker(process: psutil.Process, root: int) -> bool:
|
||||
try:
|
||||
return (
|
||||
process.ppid() == root
|
||||
and process.status() != psutil.STATUS_ZOMBIE
|
||||
and "spawn_main" in " ".join(process.cmdline())
|
||||
)
|
||||
except psutil.Error:
|
||||
return False
|
||||
|
||||
|
||||
def _workers(root: int) -> tuple[psutil.Process, ...]:
|
||||
return tuple(process for process in group_members(root) if _is_worker(process, root))
|
||||
|
||||
|
||||
def _replaced(victim: psutil.Process) -> Callable[[tuple[psutil.Process, ...]], bool]:
|
||||
return lambda workers: len(workers) == 2 and victim.pid not in {worker.pid for worker in workers}
|
||||
|
||||
|
||||
@pytest.mark.timeout(600)
|
||||
def test_killing_one_worker_keeps_the_shadowing_alias_refused(rig: AliasRig, tmp_path: Path) -> None:
|
||||
config: Final = _own_config(
|
||||
tmp_path, "shadow-two-workers.yaml", "general_settings", _base_config()["general_settings"]
|
||||
)
|
||||
with rig.gateway.scenario() as scenario:
|
||||
target, shadow = _shadow_over_unpriced_target(rig, scenario)
|
||||
healthy: Final = scenario.key(max_budget=BUDGET)
|
||||
exhausted: Final = exhausted_key(rig, scenario)
|
||||
with owned_proxy_process(rig.gateway, tmp_path, {}, config=config, workers=2) as owned:
|
||||
candidate: Final = owned.gateway
|
||||
root: Final = owned.process.pid
|
||||
before: Final = eventually(lambda: _workers(root), lambda workers: len(workers) == 2, seconds=60)
|
||||
_await_alias(candidate, shadow, target)
|
||||
settle_candidate(candidate, target, exhausted, BUDGET_EXCEEDED, seconds=120)
|
||||
settle_candidate(candidate, shadow, exhausted, BUDGET_EXCEEDED, seconds=120)
|
||||
victim: Final = before[0]
|
||||
victim.send_signal(signal.SIGKILL)
|
||||
eventually(lambda: _workers(root), _replaced(victim), seconds=120)
|
||||
settle_candidate(candidate, shadow, exhausted, BUDGET_EXCEEDED, seconds=120)
|
||||
assert chat_statuses(candidate, shadow, exhausted, GATEWAY_BURST) == {BUDGET_EXCEEDED}
|
||||
assert chat_statuses(candidate, shadow, healthy, GATEWAY_BURST) == {200}
|
||||
assert chat_statuses(candidate, target, exhausted, GATEWAY_BURST) == {BUDGET_EXCEEDED}
|
||||
|
|
|
|||
|
|
@ -16,6 +16,54 @@ from litellm.proxy.auth.auth_checks import _is_model_cost_zero
|
|||
from litellm.router import Router
|
||||
|
||||
|
||||
UNPRICED_ZERO_COST_MODEL = "ollama/unpriced-zero-cost-target"
|
||||
|
||||
|
||||
def _explicitly_free(name: str, model: str = "gpt-3.5-turbo", **model_info: float) -> dict:
|
||||
return {
|
||||
"model_name": name,
|
||||
"litellm_params": {
|
||||
"model": model,
|
||||
"api_key": "sk-fake",
|
||||
"input_cost_per_token": 0.0,
|
||||
"output_cost_per_token": 0.0,
|
||||
},
|
||||
"model_info": {"id": f"{name}-id", **model_info},
|
||||
}
|
||||
|
||||
|
||||
def _free_by_cost_map_only(name: str) -> dict:
|
||||
litellm.model_cost[UNPRICED_ZERO_COST_MODEL] = {
|
||||
"input_cost_per_token": 0.0,
|
||||
"output_cost_per_token": 0.0,
|
||||
"litellm_provider": "ollama",
|
||||
"mode": "chat",
|
||||
}
|
||||
return {
|
||||
"model_name": name,
|
||||
"litellm_params": {"model": UNPRICED_ZERO_COST_MODEL, "api_base": "http://localhost:11434"},
|
||||
"model_info": {"id": f"{name}-id"},
|
||||
}
|
||||
|
||||
|
||||
def _explicitly_free_ollama_wildcard() -> dict:
|
||||
return {
|
||||
"model_name": "ollama/*",
|
||||
"litellm_params": {
|
||||
"model": "ollama/*",
|
||||
"api_base": "http://localhost:11434",
|
||||
"input_cost_per_token": 0.0,
|
||||
"output_cost_per_token": 0.0,
|
||||
},
|
||||
"model_info": {"id": "free-wildcard-id"},
|
||||
}
|
||||
|
||||
|
||||
def _served_model(router: Router, name: str) -> str:
|
||||
deployment = router.get_available_deployment(model=name, messages=[{"role": "user", "content": "hi"}])
|
||||
return deployment["litellm_params"]["model"]
|
||||
|
||||
|
||||
class TestUnmappedModelBudgetEnforcement:
|
||||
"""Unmapped models must NOT bypass budget checks."""
|
||||
|
||||
|
|
@ -372,11 +420,12 @@ class TestUnmappedModelBudgetEnforcement:
|
|||
assert _is_model_cost_zero(model="hidden-alias", llm_router=router) is False
|
||||
|
||||
@pytest.mark.parametrize("alias_name_first", [True, False])
|
||||
def test_alias_shadowing_a_real_group_gives_each_name_its_own_verdict(self, alias_name_first: bool):
|
||||
"""An alias whose name is also a real PTU-priced group never shares a verdict with its target.
|
||||
def test_alias_shadowing_a_real_group_answers_for_its_target_in_either_order(self, alias_name_first: bool):
|
||||
"""An alias whose name is also a real PTU-priced group is judged by the free target it routes to.
|
||||
|
||||
The verdict is cached per requested name, so whichever name is asked first, the free target
|
||||
stays free and the shadowed PTU name stays enforced.
|
||||
The router resolves the alias before it looks at deployments, so the PTU deployment sharing
|
||||
the alias's name is never served under it. The verdict is cached per requested name, so
|
||||
whichever name is asked first, both names read as free.
|
||||
"""
|
||||
router = Router(
|
||||
model_list=[
|
||||
|
|
@ -405,7 +454,7 @@ class TestUnmappedModelBudgetEnforcement:
|
|||
model_group_alias={"ptu-model": "free-model"},
|
||||
)
|
||||
order = ("ptu-model", "free-model") if alias_name_first else ("free-model", "ptu-model")
|
||||
expected = {"ptu-model": False, "free-model": True}
|
||||
expected = {"ptu-model": True, "free-model": True}
|
||||
|
||||
assert [_is_model_cost_zero(model=name, llm_router=router) for name in order] == [
|
||||
expected[name] for name in order
|
||||
|
|
@ -414,6 +463,124 @@ class TestUnmappedModelBudgetEnforcement:
|
|||
expected[name] for name in order
|
||||
], "the cached verdicts must match the first evaluation"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"alias_entry",
|
||||
["unpriced-target", {"model": "unpriced-target", "hidden": True}],
|
||||
ids=["plain_alias", "hidden_alias"],
|
||||
)
|
||||
def test_alias_shadowing_an_explicitly_free_group_answers_for_its_unpriced_target(
|
||||
self, alias_entry: str | dict[str, str | bool]
|
||||
):
|
||||
"""An alias keyed like an explicitly free real group is judged by its target alone, hidden or not.
|
||||
|
||||
The router serves the alias name from its target, a group whose cost-map price is zero with
|
||||
no explicit price on the deployment, so the budget stays enforced exactly as it is for the
|
||||
target by name. The shadowed explicitly free deployment is never served under that name and
|
||||
must not lend it the bypass.
|
||||
"""
|
||||
litellm.model_cost["ollama/unpriced-zero-cost-target"] = {
|
||||
"input_cost_per_token": 0.0,
|
||||
"output_cost_per_token": 0.0,
|
||||
"litellm_provider": "ollama",
|
||||
"mode": "chat",
|
||||
}
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "free-model",
|
||||
"litellm_params": {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"api_key": "sk-fake",
|
||||
"input_cost_per_token": 0.0,
|
||||
"output_cost_per_token": 0.0,
|
||||
},
|
||||
"model_info": {"id": "free-model-id"},
|
||||
},
|
||||
{
|
||||
"model_name": "unpriced-target",
|
||||
"litellm_params": {
|
||||
"model": "ollama/unpriced-zero-cost-target",
|
||||
"api_base": "http://localhost:11434",
|
||||
},
|
||||
"model_info": {"id": "unpriced-target-id"},
|
||||
},
|
||||
],
|
||||
model_group_alias={"free-model": alias_entry},
|
||||
)
|
||||
|
||||
assert _is_model_cost_zero(model="unpriced-target", llm_router=router) is False
|
||||
assert _is_model_cost_zero(model="free-model", llm_router=router) is False, (
|
||||
"the alias routes to the unpriced target, so it must be refused like the target by name"
|
||||
)
|
||||
|
||||
def test_hidden_alias_shadowing_a_ptu_group_answers_for_its_free_target(self):
|
||||
"""A hidden alias keyed like a PTU-priced real group reads as free when its target is free.
|
||||
|
||||
The PTU deployment sharing the alias's name is never served under it, so its flat cost must
|
||||
not keep the alias enforced while the router serves every call from the free target.
|
||||
"""
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "ptu-model",
|
||||
"litellm_params": {
|
||||
"model": "azure/ptu-deployment",
|
||||
"api_base": "https://fake.openai.azure.com",
|
||||
"api_key": "sk-fake",
|
||||
"input_cost_per_token": 0.0,
|
||||
"output_cost_per_token": 0.0,
|
||||
},
|
||||
"model_info": {"id": "ptu-model-id", "ptu_count": 100, "cost_per_ptu_per_hour": 2.0},
|
||||
},
|
||||
{
|
||||
"model_name": "free-model",
|
||||
"litellm_params": {
|
||||
"model": "ollama/llama2",
|
||||
"api_base": "http://localhost:11434",
|
||||
"input_cost_per_token": 0.0,
|
||||
"output_cost_per_token": 0.0,
|
||||
},
|
||||
"model_info": {"id": "free-model-id"},
|
||||
},
|
||||
],
|
||||
model_group_alias={"ptu-model": {"model": "free-model", "hidden": True}},
|
||||
)
|
||||
|
||||
assert _is_model_cost_zero(model="free-model", llm_router=router) is True
|
||||
assert _is_model_cost_zero(model="ptu-model", llm_router=router) is True, (
|
||||
"the alias routes to the free target, so it must bypass budget like the target by name"
|
||||
)
|
||||
|
||||
def test_hidden_alias_shadowing_an_explicitly_free_group_to_a_priced_target_enforces_budget(self):
|
||||
"""A hidden alias keyed like an explicitly free real group stays enforced when its target is priced."""
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "free-model",
|
||||
"litellm_params": {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"api_key": "sk-fake",
|
||||
"input_cost_per_token": 0.0,
|
||||
"output_cost_per_token": 0.0,
|
||||
},
|
||||
"model_info": {"id": "free-model-id"},
|
||||
},
|
||||
{
|
||||
"model_name": "paid-model",
|
||||
"litellm_params": {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"api_key": "sk-fake",
|
||||
"input_cost_per_token": 0.0000002,
|
||||
"output_cost_per_token": 0.0000012,
|
||||
},
|
||||
"model_info": {"id": "paid-model-id"},
|
||||
},
|
||||
],
|
||||
model_group_alias={"free-model": {"model": "paid-model", "hidden": True}},
|
||||
)
|
||||
|
||||
assert _is_model_cost_zero(model="free-model", llm_router=router) is False
|
||||
|
||||
def test_alias_chain_through_a_priced_group_enforces_budget(self):
|
||||
"""An alias to a group that is itself an alias key resolves one hop, like the router does.
|
||||
|
||||
|
|
@ -448,6 +615,141 @@ class TestUnmappedModelBudgetEnforcement:
|
|||
|
||||
assert _is_model_cost_zero(model="chain-smart", llm_router=router) is False
|
||||
|
||||
@pytest.mark.parametrize("hidden", [False, True], ids=["plain_alias", "hidden_alias"])
|
||||
def test_alias_to_an_unpriced_group_that_is_also_an_alias_enforces_budget(self, hidden: bool):
|
||||
"""An alias is judged by the group it routes to, never by where that group's own alias points.
|
||||
|
||||
The router serves ``chain-entry`` from the real ``chain-middle`` deployment, whose zero price
|
||||
comes from the cost map alone. ``chain-middle`` is also an alias key to an explicitly free
|
||||
group, a second hop the router never takes for ``chain-entry``, so that group must not lend
|
||||
it the bypass.
|
||||
"""
|
||||
|
||||
def alias(target: str) -> str | dict[str, str | bool]:
|
||||
return {"model": target, "hidden": True} if hidden else target
|
||||
|
||||
router = Router(
|
||||
model_list=[_free_by_cost_map_only("chain-middle"), _explicitly_free("free-model")],
|
||||
model_group_alias={"chain-entry": alias("chain-middle"), "chain-middle": alias("free-model")},
|
||||
)
|
||||
|
||||
assert _served_model(router, "chain-entry") == UNPRICED_ZERO_COST_MODEL
|
||||
assert _is_model_cost_zero(model="chain-entry", llm_router=router) is False
|
||||
assert _is_model_cost_zero(model="chain-middle", llm_router=router) is True, (
|
||||
"asked by its own name, chain-middle routes to the explicitly free group"
|
||||
)
|
||||
|
||||
def test_alias_to_a_free_group_that_is_also_an_alias_to_a_ptu_group_bypasses_budget(self):
|
||||
"""A PTU group one alias hop past the group that serves the request does not enforce the budget.
|
||||
|
||||
The router serves ``chain-entry`` from the real, explicitly free ``chain-middle`` deployment.
|
||||
``chain-middle`` is also an alias key to a PTU-priced group, which only a request for
|
||||
``chain-middle`` itself routes to.
|
||||
"""
|
||||
router = Router(
|
||||
model_list=[
|
||||
_explicitly_free("chain-middle"),
|
||||
_explicitly_free("ptu-model", model="azure/ptu-deployment", ptu_count=100, cost_per_ptu_per_hour=2.0),
|
||||
],
|
||||
model_group_alias={"chain-entry": "chain-middle", "chain-middle": "ptu-model"},
|
||||
)
|
||||
|
||||
assert _served_model(router, "chain-entry") == "gpt-3.5-turbo"
|
||||
assert _is_model_cost_zero(model="chain-entry", llm_router=router) is True
|
||||
assert _is_model_cost_zero(model="chain-middle", llm_router=router) is False, (
|
||||
"asked by its own name, chain-middle routes to the PTU group"
|
||||
)
|
||||
|
||||
def test_alias_chain_served_by_a_paid_wildcard_route_enforces_budget(self):
|
||||
"""An alias whose target is only an alias key is never judged by that second alias's free group.
|
||||
|
||||
``chain-entry`` resolves one hop to ``gpt-4o-mini``, which is no deployment's name, so the
|
||||
router serves it from the paid wildcard route. The free group ``gpt-4o-mini`` is aliased
|
||||
to is only reached by a request for ``gpt-4o-mini`` itself.
|
||||
"""
|
||||
router = Router(
|
||||
model_list=[
|
||||
_explicitly_free("free-model", model="ollama/llama2"),
|
||||
{"model_name": "*", "litellm_params": {"model": "openai/*", "api_key": "sk-fake"}},
|
||||
],
|
||||
model_group_alias={"chain-entry": "gpt-4o-mini", "gpt-4o-mini": "free-model"},
|
||||
)
|
||||
|
||||
assert _served_model(router, "chain-entry") == "openai/gpt-4o-mini"
|
||||
assert _is_model_cost_zero(model="chain-entry", llm_router=router) is False
|
||||
|
||||
@pytest.mark.parametrize("alias_name", ["openai/smart", "smart"], ids=["alias_on_pattern", "alias_off_pattern"])
|
||||
def test_alias_chain_served_by_an_explicitly_priced_wildcard_route_enforces_budget(self, alias_name: str):
|
||||
"""A chain the router serves from a priced wildcard route is budgeted at that route's price.
|
||||
|
||||
The group's price reads $0 through the second alias to the free group, and the wildcard
|
||||
route's cost-map entry carries explicit prices, so only their sign tells that the served
|
||||
deployment is not free.
|
||||
"""
|
||||
router = Router(
|
||||
model_list=[
|
||||
_explicitly_free("free-model", model="ollama/llama2"),
|
||||
{
|
||||
"model_name": "openai/*",
|
||||
"litellm_params": {
|
||||
"model": "openai/*",
|
||||
"api_key": "sk-fake",
|
||||
"input_cost_per_token": 0.00001,
|
||||
"output_cost_per_token": 0.00002,
|
||||
},
|
||||
"model_info": {"id": "priced-wildcard-id"},
|
||||
},
|
||||
],
|
||||
model_group_alias={alias_name: "openai/gpt-4o-mini", "openai/gpt-4o-mini": "free-model"},
|
||||
)
|
||||
|
||||
assert _served_model(router, alias_name) == "openai/gpt-4o-mini"
|
||||
assert _is_model_cost_zero(model=alias_name, llm_router=router) is False
|
||||
|
||||
def test_alias_shadowing_a_free_group_is_judged_by_its_unpriced_target_through_an_alias_chain(self):
|
||||
"""A shadowing alias stays enforced when its unpriced target is itself an alias key to a free group."""
|
||||
router = Router(
|
||||
model_list=[
|
||||
_explicitly_free("shadowed-free"),
|
||||
_free_by_cost_map_only("chain-middle"),
|
||||
_explicitly_free("free-model"),
|
||||
],
|
||||
model_group_alias={"shadowed-free": "chain-middle", "chain-middle": "free-model"},
|
||||
)
|
||||
|
||||
assert _served_model(router, "shadowed-free") == UNPRICED_ZERO_COST_MODEL
|
||||
assert _is_model_cost_zero(model="shadowed-free", llm_router=router) is False
|
||||
|
||||
@pytest.mark.parametrize("alias_name", ["ollama/fast", "fast"], ids=["alias_on_pattern", "alias_off_pattern"])
|
||||
def test_alias_to_a_name_served_by_an_explicitly_free_wildcard_route_bypasses_budget(self, alias_name: str):
|
||||
"""An alias to a name only a wildcard route serves reads that route's deployment, like the name itself."""
|
||||
router = Router(
|
||||
model_list=[_explicitly_free_ollama_wildcard()],
|
||||
model_group_alias={alias_name: "ollama/llama3"},
|
||||
)
|
||||
|
||||
assert _served_model(router, alias_name) == "ollama/llama3"
|
||||
assert _is_model_cost_zero(model="ollama/llama3", llm_router=router) is True
|
||||
assert _is_model_cost_zero(model=alias_name, llm_router=router) is True
|
||||
|
||||
def test_alias_chain_to_a_name_served_by_an_explicitly_free_wildcard_route_bypasses_budget(self):
|
||||
"""An alias to an alias key no deployment is named after reads the wildcard route serving it.
|
||||
|
||||
``ollama/fast`` resolves one hop to ``ollama/llama3``, which is only an alias key, so the
|
||||
router serves it from the explicitly free wildcard route. The group ``ollama/llama3`` is
|
||||
aliased to is only reached by a request for ``ollama/llama3`` itself.
|
||||
"""
|
||||
router = Router(
|
||||
model_list=[_explicitly_free_ollama_wildcard(), _free_by_cost_map_only("unpriced-model")],
|
||||
model_group_alias={"ollama/fast": "ollama/llama3", "ollama/llama3": "unpriced-model"},
|
||||
)
|
||||
|
||||
assert _served_model(router, "ollama/fast") == "ollama/llama3"
|
||||
assert _is_model_cost_zero(model="ollama/fast", llm_router=router) is True
|
||||
assert _is_model_cost_zero(model="ollama/llama3", llm_router=router) is False, (
|
||||
"asked by its own name, ollama/llama3 routes to the unpriced group"
|
||||
)
|
||||
|
||||
def test_handles_router_without_zero_cost_cache_attribute(self):
|
||||
"""Tolerate router-like objects (e.g. ``MagicMock`` stand-ins) that
|
||||
do not expose ``_zero_cost_cache`` — the auth check must still
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue