mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
fix: add new tests, split gateway
This commit is contained in:
mubashir1osmani
mateo-berri
parent
95aec209da
commit
74d5d1c3a7
10 changed files with 558 additions and 29 deletions
|
|
@ -12,7 +12,7 @@ from __future__ import annotations
|
|||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from pydantic import BaseModel, RootModel
|
||||
from pydantic import AliasPath, BaseModel, Field, RootModel
|
||||
|
||||
from e2e_gateway import Gateway, build_gateway
|
||||
from e2e_http import NoBody, StreamingResponse, Success, unwrap
|
||||
|
|
@ -65,6 +65,7 @@ class TeamNewBody(BaseModel):
|
|||
team_alias: str
|
||||
max_budget: float | None = None
|
||||
organization_id: str | None = None
|
||||
budget_limits: list[BudgetWindow] | None = None
|
||||
|
||||
|
||||
class TeamNewResponse(BaseModel):
|
||||
|
|
@ -81,6 +82,29 @@ class TeamMemberAddBody(BaseModel):
|
|||
max_budget_in_team: float | None = None
|
||||
|
||||
|
||||
class TeamMemberUpdateBody(BaseModel):
|
||||
team_id: str
|
||||
user_id: str
|
||||
max_budget_in_team: float | None = None
|
||||
budget_duration: str | None = None
|
||||
|
||||
|
||||
class TeamMembershipRow(BaseModel):
|
||||
user_id: str | None = None
|
||||
budget_reset_at: str | None = Field(
|
||||
default=None,
|
||||
validation_alias=AliasPath("litellm_budget_table", "budget_reset_at"),
|
||||
)
|
||||
|
||||
|
||||
class TeamInfoParams(BaseModel):
|
||||
team_id: str
|
||||
|
||||
|
||||
class TeamInfoResponse(BaseModel):
|
||||
team_memberships: list[TeamMembershipRow] = []
|
||||
|
||||
|
||||
class TagNewBody(BaseModel):
|
||||
name: str
|
||||
max_budget: float
|
||||
|
|
@ -125,9 +149,7 @@ def is_budget_block(result: StreamingResponse) -> bool:
|
|||
return not result.ok and "budget_exceeded" in result.body
|
||||
|
||||
|
||||
def model_budget(
|
||||
model: str, limit: float, period: str = "30d"
|
||||
) -> dict[str, ModelBudgetEntry]:
|
||||
def model_budget(model: str, limit: float, period: str = "30d") -> dict[str, ModelBudgetEntry]:
|
||||
"""A model_max_budget entry: per-model cap with a reset window."""
|
||||
return {model: ModelBudgetEntry(budget_limit=limit, time_period=period)}
|
||||
|
||||
|
|
@ -254,6 +276,7 @@ class BudgetClient:
|
|||
alias: str,
|
||||
max_budget: float | None = None,
|
||||
organization_id: str | None = None,
|
||||
budget_limits: list[BudgetWindow] | None = None,
|
||||
) -> str:
|
||||
return unwrap(
|
||||
self.gateway.transport.post(
|
||||
|
|
@ -263,6 +286,7 @@ class BudgetClient:
|
|||
team_alias=alias,
|
||||
max_budget=max_budget,
|
||||
organization_id=organization_id,
|
||||
budget_limits=budget_limits,
|
||||
),
|
||||
response_type=TeamNewResponse,
|
||||
)
|
||||
|
|
@ -276,9 +300,7 @@ class BudgetClient:
|
|||
response_type=NoBody,
|
||||
)
|
||||
|
||||
def add_team_member(
|
||||
self, team_id: str, user_id: str, *, max_budget_in_team: float | None = None
|
||||
) -> None:
|
||||
def add_team_member(self, team_id: str, user_id: str, *, max_budget_in_team: float | None = None) -> None:
|
||||
resp = self.gateway.transport.send(
|
||||
"/team/member_add",
|
||||
headers=self.gateway.transport.master,
|
||||
|
|
@ -290,6 +312,45 @@ class BudgetClient:
|
|||
)
|
||||
assert resp.ok, resp.body
|
||||
|
||||
def update_team_member(
|
||||
self,
|
||||
team_id: str,
|
||||
user_id: str,
|
||||
*,
|
||||
max_budget_in_team: float | None = None,
|
||||
budget_duration: str | None = None,
|
||||
) -> None:
|
||||
resp = self.gateway.transport.send(
|
||||
"/team/member_update",
|
||||
headers=self.gateway.transport.master,
|
||||
json=TeamMemberUpdateBody(
|
||||
team_id=team_id,
|
||||
user_id=user_id,
|
||||
max_budget_in_team=max_budget_in_team,
|
||||
budget_duration=budget_duration,
|
||||
),
|
||||
)
|
||||
assert resp.ok, resp.body
|
||||
|
||||
def member_budget_reset_at(self, team_id: str, user_id: str) -> str | None:
|
||||
"""The member's per-team budget_reset_at as /team/info reports it, or None if
|
||||
no reset is scheduled. The reset job advances this each time the window
|
||||
elapses; a job that skips the row leaves it pinned forever."""
|
||||
result = self.gateway.transport.get(
|
||||
"/team/info",
|
||||
headers=self.gateway.transport.master,
|
||||
params=TeamInfoParams(team_id=team_id),
|
||||
response_type=TeamInfoResponse,
|
||||
)
|
||||
match result:
|
||||
case Success(data=data):
|
||||
return next(
|
||||
(row.budget_reset_at for row in data.team_memberships if row.user_id == user_id),
|
||||
None,
|
||||
)
|
||||
case _:
|
||||
return None
|
||||
|
||||
# ---- tag ------------------------------------------------------------
|
||||
|
||||
def create_tag(self, name: str, *, max_budget: float) -> str:
|
||||
|
|
|
|||
118
tests/e2e/budgets/test_spend_counter_reseed_e2e.py
Normal file
118
tests/e2e/budgets/test_spend_counter_reseed_e2e.py
Normal file
|
|
@ -0,0 +1,118 @@
|
|||
"""Live e2e: concurrent cold-counter reseeds keep the spend counter equal to DB spend (#26829).
|
||||
|
||||
Regression for the cross-pod spend-counter multiplication. Real requests build a key's
|
||||
DB spend through the spend writer; the Redis spend counter then expires (the e2e proxy
|
||||
sets a short default_redis_ttl) and goes cold. The proxy runs several workers sharing one
|
||||
Redis, so a concurrent burst makes more than one worker reseed the same cold counter at
|
||||
once. The fix seeds with SET NX - one worker initializes the counter at the DB spend and
|
||||
the rest read it back - so the counter still equals the DB spend (plus the burst's own
|
||||
small cost). The pre-#26829 additive reseed stacked the DB spend once per worker, leaving
|
||||
the counter at ~N x the real spend.
|
||||
|
||||
The test reads the shared counter straight from Redis and asserts it equals the DB spend,
|
||||
not a multiple. It also asserts the counter actually went cold before the burst, so a proxy
|
||||
that never expires the counter (no short TTL) fails loudly instead of passing vacuously.
|
||||
Skipped when the e2e Redis is not reachable.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import os
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from threading import Barrier
|
||||
|
||||
import pytest
|
||||
|
||||
from budget_client import BudgetClient
|
||||
from e2e_config import unique_marker
|
||||
from e2e_http import StreamingResponse
|
||||
from lifecycle import ResourceManager
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
MODEL = "claude-haiku-4-5"
|
||||
ACCUMULATE_CALLS = 24
|
||||
BURST = 6
|
||||
# proxy_batch_write_at (60s) flushes the spend to the DB and default_redis_ttl (20s)
|
||||
# expires the counter; this waits out both.
|
||||
COLD_WAIT_SECONDS = 80
|
||||
|
||||
|
||||
def _redis():
|
||||
import redis
|
||||
|
||||
return redis.Redis(
|
||||
host=os.getenv("E2E_REDIS_HOST", "localhost"),
|
||||
port=int(os.getenv("E2E_REDIS_PORT", "6380")),
|
||||
password=os.getenv("REDIS_PASSWORD") or None,
|
||||
decode_responses=True,
|
||||
socket_connect_timeout=2,
|
||||
)
|
||||
|
||||
|
||||
def _spend_counter(rds, key: str) -> float | None:
|
||||
"""The shared spend counter for `key`, or None if it is cold. The counter key is
|
||||
``{cache namespace}:spend:key:{sha256(key)}``; matched by suffix so the configured
|
||||
namespace need not be hard-coded."""
|
||||
digest = hashlib.sha256(key.encode()).hexdigest()
|
||||
matches = list(rds.scan_iter(match=f"*spend:key:{digest}"))
|
||||
if not matches:
|
||||
return None
|
||||
raw = rds.get(matches[0])
|
||||
return float(raw) if raw is not None else None
|
||||
|
||||
|
||||
def _chat(client: BudgetClient, key: str) -> StreamingResponse:
|
||||
return client.chat(key, MODEL, f"reseed {unique_marker()}", max_tokens=16)
|
||||
|
||||
|
||||
def _accumulate(client: BudgetClient, key: str, count: int) -> None:
|
||||
with ThreadPoolExecutor(max_workers=8) as pool:
|
||||
list(pool.map(lambda _: _chat(client, key), range(count)))
|
||||
|
||||
|
||||
def _burst(client: BudgetClient, key: str, count: int) -> None:
|
||||
"""Fire `count` requests that start together, so multiple workers reseed the cold
|
||||
counter concurrently rather than one warming it before the others arrive."""
|
||||
barrier = Barrier(count)
|
||||
|
||||
def one(_: int) -> StreamingResponse:
|
||||
barrier.wait()
|
||||
return _chat(client, key)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=count) as pool:
|
||||
list(pool.map(one, range(count)))
|
||||
|
||||
|
||||
def test_cold_counter_reseed_keeps_counter_equal_to_db_spend(
|
||||
client: BudgetClient, resources: ResourceManager
|
||||
) -> None:
|
||||
try:
|
||||
rds = _redis()
|
||||
rds.ping()
|
||||
except Exception as exc: # noqa: BLE001 - any connect failure means skip
|
||||
pytest.skip(f"e2e redis not reachable (set E2E_REDIS_HOST/E2E_REDIS_PORT): {exc}")
|
||||
|
||||
key = client.generate_key(max_budget=1.0, models=[MODEL])
|
||||
resources.defer(lambda: client.delete_key(key))
|
||||
|
||||
_accumulate(client, key, ACCUMULATE_CALLS)
|
||||
time.sleep(COLD_WAIT_SECONDS)
|
||||
|
||||
assert _spend_counter(rds, key) is None, (
|
||||
"the spend counter never went cold; default_redis_ttl must be short enough for it "
|
||||
"to expire, otherwise the burst reads a warm counter and the reseed is never exercised"
|
||||
)
|
||||
db_spend = client.gateway.key_info(key).spend or 0.0
|
||||
assert db_spend > 0, f"no DB spend accumulated from real calls: {db_spend}"
|
||||
|
||||
_burst(client, key, BURST)
|
||||
time.sleep(3)
|
||||
|
||||
counter = _spend_counter(rds, key)
|
||||
assert counter is not None, "the burst did not reseed the cold counter"
|
||||
assert db_spend * 0.95 <= counter < db_spend * 1.7, (
|
||||
f"redis spend counter {counter} does not equal DB spend {db_spend} (expected ~equal "
|
||||
f"plus the burst's small cost); a near-multiple means the cold-counter reseed stacked "
|
||||
f"the DB spend once per worker instead of seeding it once (#26829)"
|
||||
)
|
||||
99
tests/e2e/budgets/test_team_member_budget_e2e.py
Normal file
99
tests/e2e/budgets/test_team_member_budget_e2e.py
Normal file
|
|
@ -0,0 +1,99 @@
|
|||
"""Live e2e: a team member's per-team budget attributes spend and enforces a cap.
|
||||
|
||||
The team carries a large budget while the one enrolled member is capped at a tiny
|
||||
per-team budget, so any block is member-level, not team-level. Two scenarios share
|
||||
that single member:
|
||||
- attribution: the member's calls land in the spend logs tagged with both the team_id
|
||||
and the member's user_id, so per-member spend can be billed back
|
||||
- enforcement: once the member's spend passes the per-team budget, calls are blocked
|
||||
with budget_exceeded while the team's own budget is nowhere near exhausted
|
||||
|
||||
Per-member budgets enforce off batch-written spend (~60s), so a quick burst all goes
|
||||
through; the block only lands once that spend flushes.
|
||||
"""
|
||||
|
||||
import time
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
|
||||
from budget_client import BudgetClient, is_budget_block
|
||||
from e2e_config import unique_marker
|
||||
from e2e_http import Success, require_successful_call
|
||||
from models import ChatBody, ChatMessage
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
MODEL = "claude-haiku-4-5"
|
||||
TEAM_BUDGET = 100.0
|
||||
MEMBER_BUDGET = 3e-6
|
||||
BURST = 6
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Member:
|
||||
team_id: str
|
||||
user_id: str
|
||||
key: str
|
||||
|
||||
|
||||
@pytest.fixture(scope="class")
|
||||
def member(client: BudgetClient) -> Iterator[_Member]:
|
||||
"""A team with a large budget plus one member capped at a tiny per-team budget,
|
||||
and that member's key. Shared across the class; torn down when it finishes."""
|
||||
marker = unique_marker()
|
||||
team_id = client.create_team(alias=f"e2e-team-member-{marker}", max_budget=TEAM_BUDGET)
|
||||
user_id = client.create_user(max_budget=TEAM_BUDGET)
|
||||
client.add_team_member(team_id, user_id, max_budget_in_team=MEMBER_BUDGET)
|
||||
key = client.generate_key(team_id=team_id, user_id=user_id)
|
||||
yield _Member(team_id=team_id, user_id=user_id, key=key)
|
||||
client.delete_key(key)
|
||||
client.delete_user(user_id)
|
||||
client.delete_team(team_id)
|
||||
|
||||
|
||||
def _send(client: BudgetClient, key: str) -> str | None:
|
||||
"""One member call; its response id (== the spend-log request_id) if it went
|
||||
through, else None."""
|
||||
match client.gateway.chat(
|
||||
key,
|
||||
ChatBody(
|
||||
model=MODEL,
|
||||
messages=[ChatMessage(role="user", content=f"hi {unique_marker()}")],
|
||||
max_tokens=16,
|
||||
),
|
||||
):
|
||||
case Success(data=response):
|
||||
return response.id
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
class TestTeamMemberBudget:
|
||||
def test_member_spend_attributed_to_team_and_user(self, client: BudgetClient, member: _Member) -> None:
|
||||
sent = frozenset(rid for rid in (_send(client, member.key) for _ in range(BURST)) if rid)
|
||||
assert sent, "no member call went through; cannot check attribution"
|
||||
|
||||
rows = client.gateway.poll_logs_for_key(
|
||||
member.key, predicate=lambda rs: bool(sent & {r.request_id for r in rs})
|
||||
)
|
||||
logged = [row for row in rows if row.request_id in sent]
|
||||
assert logged, f"none of the member's {len(sent)} calls reached the spend logs"
|
||||
|
||||
for row in logged:
|
||||
assert row.team_id == member.team_id, (
|
||||
f"call {row.request_id} logged under team {row.team_id}, not the member's team {member.team_id}"
|
||||
)
|
||||
assert row.user == member.user_id, (
|
||||
f"call {row.request_id} logged under user {row.user}, not member {member.user_id}"
|
||||
)
|
||||
|
||||
def test_member_spend_over_budget_is_blocked(self, client: BudgetClient, member: _Member) -> None:
|
||||
for _ in range(40):
|
||||
result = client.chat(member.key, MODEL, f"spend {unique_marker()}", max_tokens=16)
|
||||
if is_budget_block(result):
|
||||
return
|
||||
require_successful_call(result)
|
||||
time.sleep(2)
|
||||
pytest.fail("per-member budget never enforced within the call budget")
|
||||
47
tests/e2e/budgets/test_team_member_budget_reset_e2e.py
Normal file
47
tests/e2e/budgets/test_team_member_budget_reset_e2e.py
Normal file
|
|
@ -0,0 +1,47 @@
|
|||
import time
|
||||
from datetime import datetime
|
||||
|
||||
import pytest
|
||||
|
||||
from budget_client import BudgetClient
|
||||
from e2e_config import unique_marker
|
||||
from e2e_http import require_successful_call
|
||||
from lifecycle import ResourceManager
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
MEMBER_BUDGET = 1.0 # default member budget is $50, we're testing with a smaller value
|
||||
|
||||
def _as_datetime(value: str) -> datetime:
|
||||
return datetime.fromisoformat(value.replace("Z", "+00:00"))
|
||||
|
||||
|
||||
def test_team_member_budget_reset_keeps_advancing(client: BudgetClient, resources: ResourceManager) -> None:
|
||||
team_id = client.create_team(alias=f"e2e-member-reset-{unique_marker()}", max_budget=100.0)
|
||||
resources.defer(lambda: client.delete_team(team_id))
|
||||
user_id = client.create_user(max_budget=100.0)
|
||||
resources.defer(lambda: client.delete_user(user_id))
|
||||
|
||||
# add the member, then update them onto a short per-team budget window
|
||||
client.add_team_member(team_id, user_id, max_budget_in_team=MEMBER_BUDGET)
|
||||
client.update_team_member(team_id, user_id, max_budget_in_team=MEMBER_BUDGET, budget_duration="30s")
|
||||
|
||||
scheduled = client.member_budget_reset_at(team_id, user_id)
|
||||
assert scheduled, "updating the member with a budget_duration set no budget_reset_at"
|
||||
first_reset = _as_datetime(scheduled)
|
||||
|
||||
# the member can spend within the team while the window is live
|
||||
key = client.generate_key(team_id=team_id, user_id=user_id)
|
||||
resources.defer(lambda: client.delete_key(key))
|
||||
require_successful_call(client.chat(key, "claude-haiku-4-5", f"reset {unique_marker()}", max_tokens=16))
|
||||
|
||||
# once the window elapses the reset job must move budget_reset_at forward; a job
|
||||
# that skips the member's budget row (the #25109 regression) leaves it pinned at
|
||||
# first_reset forever
|
||||
deadline = time.monotonic() + 90
|
||||
while time.monotonic() < deadline:
|
||||
time.sleep(5)
|
||||
current = client.member_budget_reset_at(team_id, user_id)
|
||||
if current and _as_datetime(current) > first_reset:
|
||||
return
|
||||
pytest.fail(f"member budget_reset_at never advanced past {first_reset.isoformat()} in 90s")
|
||||
71
tests/e2e/budgets/test_team_multi_window_budget_e2e.py
Normal file
71
tests/e2e/budgets/test_team_multi_window_budget_e2e.py
Normal file
|
|
@ -0,0 +1,71 @@
|
|||
"""Live e2e: a team's multi-window budgets (budget_limits) enforce AND reset per window.
|
||||
|
||||
The team analog of test_multi_window_budget_e2e.py (which covers keys). A team is
|
||||
created with a tight 30s window and a roomy 1m window; a key on that team blocks once
|
||||
the tight window's cap is exceeded, then - once the 30s elapses and the reset job runs
|
||||
(rescheduled fast via PROXY_BUDGET_RESCHEDULER_* in docker-compose) - the window resets
|
||||
and calls flow again. This exercises the reset_budget_windows TEAM branch (raw SQL over
|
||||
LiteLLM_TeamTable.budget_limits, the literal #25109 path), which had no live coverage.
|
||||
|
||||
Currently fails at team creation: /team/new writes the raw window list straight to the
|
||||
Json? column, where Prisma rejects it (500), unlike the key path and /team/update which
|
||||
json.dumps it first. Left failing rather than weakened - it passes once that write is
|
||||
fixed.
|
||||
"""
|
||||
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from budget_client import BudgetClient, is_budget_block
|
||||
from e2e_config import unique_marker
|
||||
from e2e_http import require_successful_call
|
||||
from lifecycle import ResourceManager
|
||||
from models import BudgetWindow
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
WINDOW_SECONDS = 30
|
||||
|
||||
|
||||
def _call(client: BudgetClient, key: str):
|
||||
return client.chat(key, "claude-haiku-4-5", f"team-window {unique_marker()}", max_tokens=16)
|
||||
|
||||
|
||||
def test_team_short_window_blocks_then_resets(client: BudgetClient, resources: ResourceManager) -> None:
|
||||
team_id = client.create_team(
|
||||
alias=f"e2e-team-window-{unique_marker()}",
|
||||
budget_limits=[
|
||||
BudgetWindow(budget_duration=f"{WINDOW_SECONDS}s", max_budget=3e-6),
|
||||
BudgetWindow(budget_duration="1m", max_budget=1.0), # roomy: never blocks
|
||||
],
|
||||
)
|
||||
resources.defer(lambda: client.delete_team(team_id))
|
||||
key = client.generate_key(team_id=team_id)
|
||||
resources.defer(lambda: client.delete_key(key))
|
||||
|
||||
# 1. exhaust the tight window -> litellm returns budget_exceeded
|
||||
start = time.monotonic()
|
||||
blocked = False
|
||||
for _ in range(20):
|
||||
result = _call(client, key)
|
||||
if is_budget_block(result):
|
||||
blocked = True
|
||||
break
|
||||
require_successful_call(result)
|
||||
time.sleep(2)
|
||||
assert blocked, f"team {WINDOW_SECONDS}s window never enforced"
|
||||
|
||||
# 2. the window resets at the next wall-clock-aligned boundary + the reset job. When
|
||||
# a call flows again the window has reset; the elapsed clock must be short enough
|
||||
# that this is the 30s window resetting, not the roomy 1m one.
|
||||
deadline = time.monotonic() + 90
|
||||
while time.monotonic() < deadline:
|
||||
time.sleep(5)
|
||||
result = _call(client, key)
|
||||
if result.ok:
|
||||
elapsed = time.monotonic() - start
|
||||
assert elapsed < WINDOW_SECONDS + 45, f"reset took {elapsed:.0f}s - too long for a {WINDOW_SECONDS}s window"
|
||||
return
|
||||
assert is_budget_block(result), f"non-budget error during reset wait: {result.body[:200]}"
|
||||
pytest.fail(f"team {WINDOW_SECONDS}s window never reset within 90s")
|
||||
|
|
@ -21,7 +21,7 @@ from typing import Iterator
|
|||
import pytest
|
||||
import requests
|
||||
|
||||
from e2e_config import PROXY_BASE_URL
|
||||
from e2e_config import CONTROL_PLANE_BASE_URL, PROXY_BASE_URL
|
||||
from lifecycle import GatewayProvider, ResourceManager
|
||||
|
||||
|
||||
|
|
@ -35,15 +35,27 @@ def pytest_configure(config: pytest.Config) -> None:
|
|||
)
|
||||
|
||||
|
||||
def _liveness_reason(label: str, base_url: str) -> str | None:
|
||||
"""None if `base_url` answers its liveness probe, else a skip reason."""
|
||||
try:
|
||||
resp = requests.get(f"{base_url}/health/liveliness", timeout=5)
|
||||
except requests.RequestException as exc:
|
||||
return f"No live {label} at {base_url}: {exc}"
|
||||
if resp.status_code >= 500:
|
||||
return f"{label} at {base_url} returned {resp.status_code}"
|
||||
return None
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=1)
|
||||
def _proxy_skip_reason() -> str | None:
|
||||
"""Probe the proxy once per session. None if it answers, else a skip reason."""
|
||||
try:
|
||||
resp = requests.get(f"{PROXY_BASE_URL}/health/liveliness", timeout=5)
|
||||
except requests.RequestException as exc:
|
||||
return f"No live proxy at {PROXY_BASE_URL}: {exc}"
|
||||
if resp.status_code >= 500:
|
||||
return f"Proxy at {PROXY_BASE_URL} returned {resp.status_code}"
|
||||
"""Probe the proxy once per session. None if it answers, else a skip reason. In
|
||||
a split deployment the management/admin control plane is a separate service, so
|
||||
require it too (when it differs) - else its tests would fail rather than skip."""
|
||||
reason = _liveness_reason("proxy", PROXY_BASE_URL)
|
||||
if reason is not None:
|
||||
return reason
|
||||
if CONTROL_PLANE_BASE_URL != PROXY_BASE_URL:
|
||||
return _liveness_reason("control plane", CONTROL_PLANE_BASE_URL)
|
||||
return None
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -10,6 +10,17 @@ import uuid
|
|||
PROXY_BASE_URL = os.environ.get("LITELLM_PROXY_URL", "http://localhost:4000").rstrip("/")
|
||||
MASTER_KEY = os.environ.get("LITELLM_MASTER_KEY", "sk-1234")
|
||||
|
||||
# Control-plane (management/admin) base URL. In a split control-plane/data-plane
|
||||
# deployment the LLM data plane (PROXY_BASE_URL: /chat, /embeddings, native
|
||||
# passthrough) and the management API (keys, users, teams, orgs, budgets, spend,
|
||||
# model info, /openapi.json) are served by *different* services. The suite drives
|
||||
# both through one Transport that routes by path (see transport.SplitTransport).
|
||||
# Defaults to PROXY_BASE_URL so a monolithic proxy serving everything on one URL
|
||||
# behaves exactly as before.
|
||||
CONTROL_PLANE_BASE_URL = os.environ.get(
|
||||
"LITELLM_CONTROL_PLANE_URL", PROXY_BASE_URL
|
||||
).rstrip("/")
|
||||
|
||||
# Writes on the proxy are eventually consistent (e.g. spend rows flush on
|
||||
# proxy_batch_write_at, ~60s). Read-backs poll to this deadline, never sleep-once.
|
||||
POLL_TIMEOUT = float(os.environ.get("E2E_POLL_TIMEOUT", "120"))
|
||||
|
|
|
|||
|
|
@ -39,13 +39,14 @@ from models import (
|
|||
SpendLogsParams,
|
||||
)
|
||||
from e2e_config import (
|
||||
CONTROL_PLANE_BASE_URL,
|
||||
MASTER_KEY,
|
||||
POLL_INTERVAL,
|
||||
POLL_TIMEOUT,
|
||||
PROXY_BASE_URL,
|
||||
REQUEST_TIMEOUT,
|
||||
)
|
||||
from transport import HttpTransport, Transport
|
||||
from transport import HttpTransport, SplitTransport, Transport
|
||||
|
||||
RowsPredicate = Callable[[list[SpendLogRow]], bool]
|
||||
|
||||
|
|
@ -188,13 +189,22 @@ class Gateway:
|
|||
|
||||
|
||||
def build_gateway() -> Gateway:
|
||||
"""The Gateway every suite's client is built from: an HttpTransport pointed at
|
||||
the configured proxy, with the shared poll budget."""
|
||||
"""The Gateway every suite's client is built from: a SplitTransport that routes
|
||||
LLM calls to the data plane (PROXY_BASE_URL) and management/admin calls to the
|
||||
control plane (CONTROL_PLANE_BASE_URL), with the shared poll budget. The two
|
||||
base URLs are the same for a monolithic proxy, so routing is then a no-op."""
|
||||
return Gateway(
|
||||
transport=HttpTransport(
|
||||
base_url=PROXY_BASE_URL,
|
||||
master_key=MASTER_KEY,
|
||||
request_timeout=REQUEST_TIMEOUT,
|
||||
transport=SplitTransport(
|
||||
data=HttpTransport(
|
||||
base_url=PROXY_BASE_URL,
|
||||
master_key=MASTER_KEY,
|
||||
request_timeout=REQUEST_TIMEOUT,
|
||||
),
|
||||
control=HttpTransport(
|
||||
base_url=CONTROL_PLANE_BASE_URL,
|
||||
master_key=MASTER_KEY,
|
||||
request_timeout=REQUEST_TIMEOUT,
|
||||
),
|
||||
),
|
||||
poll_timeout=POLL_TIMEOUT,
|
||||
poll_interval=POLL_INTERVAL,
|
||||
|
|
|
|||
|
|
@ -32,6 +32,8 @@ class KeyGenerateBody(BaseModel):
|
|||
budget_id: str | None = None
|
||||
model_max_budget: dict[str, ModelBudgetEntry] | None = None
|
||||
budget_limits: list[BudgetWindow] | None = None
|
||||
tpm_limit: int | None = None
|
||||
rpm_limit: int | None = None
|
||||
|
||||
|
||||
class KeyGenerateResponse(BaseModel):
|
||||
|
|
@ -134,6 +136,8 @@ class SpendLogRow(BaseModel):
|
|||
cache_hit: str | None = None
|
||||
call_type: str | None = None
|
||||
custom_llm_provider: str | None = None
|
||||
team_id: str | None = None
|
||||
user: str | None = None
|
||||
end_user: str | None = None
|
||||
prompt_tokens: int | None = None
|
||||
completion_tokens: int | None = None
|
||||
|
|
@ -228,14 +232,10 @@ class CustomPricing(BaseModel):
|
|||
def token_cost(self, prompt_tokens: int, completion_tokens: int) -> float:
|
||||
"""Spend for a fresh (uncached) call under these rates: the proxy's
|
||||
custom-pricing formula (prompt * input + completion * output)."""
|
||||
assert (
|
||||
self.input_cost_per_token is not None
|
||||
and self.output_cost_per_token is not None
|
||||
), "custom pricing has no per-token rates"
|
||||
return (
|
||||
prompt_tokens * self.input_cost_per_token
|
||||
+ completion_tokens * self.output_cost_per_token
|
||||
assert self.input_cost_per_token is not None and self.output_cost_per_token is not None, (
|
||||
"custom pricing has no per-token rates"
|
||||
)
|
||||
return prompt_tokens * self.input_cost_per_token + completion_tokens * self.output_cost_per_token
|
||||
|
||||
|
||||
class ModelInfoEntry(BaseModel):
|
||||
|
|
|
|||
|
|
@ -142,3 +142,103 @@ class HttpTransport:
|
|||
params=params,
|
||||
timeout=self.request_timeout,
|
||||
)
|
||||
|
||||
|
||||
# Top-level management/admin route groups. In a split deployment these are served
|
||||
# by the control plane (a different service from the LLM data plane). LLM routes
|
||||
# (/chat, /embeddings, and native passthrough like /gemini, /anthropic) are NOT
|
||||
# here and fall through to the data plane. Matched as path prefixes.
|
||||
CONTROL_PLANE_PREFIXES: tuple[str, ...] = (
|
||||
"/key",
|
||||
"/user",
|
||||
"/team",
|
||||
"/organization",
|
||||
"/customer",
|
||||
"/tag",
|
||||
"/budget",
|
||||
"/model/info",
|
||||
"/spend",
|
||||
"/global",
|
||||
"/openapi.json",
|
||||
)
|
||||
|
||||
|
||||
def is_control_plane_path(path: str) -> bool:
|
||||
"""True if `path` is a management/admin route (served by the control plane in a
|
||||
split deployment), false for LLM data-plane routes."""
|
||||
return path.startswith(CONTROL_PLANE_PREFIXES)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SplitTransport:
|
||||
"""A Transport that dispatches each call by path to one of two backends: the
|
||||
management/admin control plane or the LLM data plane.
|
||||
|
||||
Litellm can run as a split control-plane/data-plane deployment where the two
|
||||
surfaces live on different services. Clients here stay plane-agnostic — they
|
||||
keep calling ``transport.post("/budget/new", ...)`` or
|
||||
``transport.send("/chat/completions", ...)`` — and routing happens in one place
|
||||
by path (see ``CONTROL_PLANE_PREFIXES``). When ``control`` and ``data`` share a
|
||||
base URL (the monolithic default), routing is a no-op. ``bearer``/``master``
|
||||
are plane-agnostic (same master key both planes), so they come from ``data``.
|
||||
"""
|
||||
|
||||
data: HttpTransport
|
||||
control: HttpTransport
|
||||
|
||||
def _route(self, path: str) -> HttpTransport:
|
||||
return self.control if is_control_plane_path(path) else self.data
|
||||
|
||||
def bearer(self, key: str) -> AuthHeaders:
|
||||
return self.data.bearer(key)
|
||||
|
||||
@property
|
||||
def master(self) -> AuthHeaders:
|
||||
return self.data.master
|
||||
|
||||
def post[R: BaseModel](
|
||||
self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R]
|
||||
) -> Result[R]:
|
||||
return self._route(path).post(
|
||||
path, headers=headers, json=json, response_type=response_type
|
||||
)
|
||||
|
||||
def get[R: BaseModel](
|
||||
self,
|
||||
path: str,
|
||||
*,
|
||||
headers: BaseModel,
|
||||
params: BaseModel,
|
||||
response_type: type[R],
|
||||
) -> Result[R]:
|
||||
return self._route(path).get(
|
||||
path, headers=headers, params=params, response_type=response_type
|
||||
)
|
||||
|
||||
def delete[R: BaseModel](
|
||||
self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R]
|
||||
) -> Result[R]:
|
||||
return self._route(path).delete(
|
||||
path, headers=headers, json=json, response_type=response_type
|
||||
)
|
||||
|
||||
def stream(
|
||||
self, path: str, *, headers: BaseModel, json: BaseModel
|
||||
) -> StreamingResponse:
|
||||
return self._route(path).stream(path, headers=headers, json=json)
|
||||
|
||||
def send(
|
||||
self,
|
||||
path: str,
|
||||
*,
|
||||
headers: BaseModel,
|
||||
json: BaseModel,
|
||||
params: BaseModel | None = None,
|
||||
stream: bool = False,
|
||||
) -> StreamingResponse:
|
||||
return self._route(path).send(
|
||||
path, headers=headers, json=json, params=params, stream=stream
|
||||
)
|
||||
|
||||
def probe(self, path: str, *, params: BaseModel) -> ProbeResult:
|
||||
return self._route(path).probe(path, params=params)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue