mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
Merge 2ad30c04c7 into d392e7faae
This commit is contained in:
ryan-crabbe-berri
GitHub
commit
9061c1fe48
2 changed files with 201 additions and 146 deletions
|
|
@ -13,6 +13,10 @@ a mint-time read races the boundary; the reset job zeroes the counter in the sam
|
|||
pass), the key must still be refused with "over 1d budget". That check polls because
|
||||
enforcement's cached auth view lags the DB write; any 200 or non-budget error fails
|
||||
immediately.
|
||||
|
||||
The blocks-then-resets check also sweeps the personal / team / team-member mint
|
||||
shapes: the same budget_limits pair rides a key minted to roomy (100.0)
|
||||
surroundings, so only the key's own windows can block regardless of who holds it.
|
||||
"""
|
||||
|
||||
import time
|
||||
|
|
@ -56,20 +60,20 @@ def _drive_to_block(client: BudgetClient, key: str) -> StreamingResponse:
|
|||
pytest.fail("budget never enforced before block")
|
||||
|
||||
|
||||
@pytest.mark.covers("quota_management.budget.key_multi_window.blocks_then_resets")
|
||||
def test_short_window_blocks_then_resets(client: BudgetClient, resources: ResourceManager) -> None:
|
||||
key = client.generate_key(
|
||||
models=[MODEL],
|
||||
budget_limits=[
|
||||
BudgetWindow(budget_duration=SHORT_WINDOW, max_budget=TINY_CAP),
|
||||
BudgetWindow(budget_duration="1m", max_budget=1.0), # roomy: never blocks
|
||||
],
|
||||
)
|
||||
resources.defer(lambda: client.delete_key(key))
|
||||
def _short_roomy_limits() -> list[BudgetWindow]:
|
||||
return [
|
||||
BudgetWindow(budget_duration=SHORT_WINDOW, max_budget=TINY_CAP),
|
||||
BudgetWindow(budget_duration="1m", max_budget=1.0), # roomy: never blocks
|
||||
]
|
||||
|
||||
|
||||
def _assert_short_window_blocks_then_resets(client: BudgetClient, key: str) -> None:
|
||||
# 1. exhaust the tight window -> litellm returns budget_exceeded
|
||||
start = time.monotonic()
|
||||
_drive_to_block(client, key)
|
||||
blocked = _drive_to_block(client, key)
|
||||
assert f"over {SHORT_WINDOW} budget" in blocked.body, (
|
||||
f"block must be attributed to the {SHORT_WINDOW} window, got: {blocked.body[:200]}"
|
||||
)
|
||||
|
||||
# 2. the window resets at the next wall-clock-aligned boundary (up to a window
|
||||
# after start), then the reset job (~15-20s rescheduler) zeroes the spend.
|
||||
|
|
@ -89,58 +93,106 @@ def test_short_window_blocks_then_resets(client: BudgetClient, resources: Resour
|
|||
pytest.fail(f"{WINDOW_SECONDS}s window never reset within 150s")
|
||||
|
||||
|
||||
@pytest.mark.covers("quota_management.budget.key_multi_window.blocks_then_resets")
|
||||
def test_long_window_blocks_after_short_window_resets(client: BudgetClient, resources: ResourceManager) -> None:
|
||||
key = client.generate_key(
|
||||
models=[MODEL],
|
||||
budget_limits=[
|
||||
BudgetWindow(budget_duration=SHORT_WINDOW, max_budget=TINY_CAP),
|
||||
BudgetWindow(budget_duration=LONG_WINDOW, max_budget=LONG_CAP),
|
||||
],
|
||||
)
|
||||
resources.defer(lambda: client.delete_key(key))
|
||||
|
||||
# 1. drive the key to get blocked by SHORT_WINDOW, assert it's budget error
|
||||
blocked = _drive_to_block(client, key)
|
||||
assert blocked.status_code == 429, f"budget block was not a 429: {blocked.status_code} {blocked.body[:200]}"
|
||||
|
||||
# 2. check the reset times of both budget windows after we drove to being blocked
|
||||
blocked_reset_at = window_reset_at(client.key_budget_windows(key), SHORT_WINDOW)
|
||||
assert blocked_reset_at is not None, "short window missing from /key/info budget_limits"
|
||||
blocked_long_reset_at = window_reset_at(client.key_budget_windows(key), LONG_WINDOW)
|
||||
assert blocked_long_reset_at is not None, "long window missing from /key/info budget_limits"
|
||||
|
||||
# 3. poll every 5s for the SHORT_WINDOW reset time until it is past it, fails if it doesnt reset
|
||||
deadline = time.monotonic() + RESET_DEADLINE_SECONDS
|
||||
while time.monotonic() < deadline:
|
||||
time.sleep(5)
|
||||
current = window_reset_at(client.key_budget_windows(key), SHORT_WINDOW)
|
||||
if current is not None and current > blocked_reset_at:
|
||||
break
|
||||
else:
|
||||
pytest.fail(
|
||||
f"{SHORT_WINDOW} window's reset_at never advanced past {blocked_reset_at} within {RESET_DEADLINE_SECONDS}s"
|
||||
)
|
||||
|
||||
# 4. short window just reset in 3, so now make a call, check that its blocked (should be blocked by LONG_WINDOW because short window reset), also make sure its budget error
|
||||
deadline = time.monotonic() + RESET_DEADLINE_SECONDS
|
||||
last_body = ""
|
||||
while time.monotonic() < deadline:
|
||||
result = _call(client, key)
|
||||
if result.ok:
|
||||
rolled = window_reset_at(client.key_budget_windows(key), LONG_WINDOW) != blocked_long_reset_at
|
||||
pytest.fail(
|
||||
f"{LONG_WINDOW} window failed to block after the {SHORT_WINDOW} window reset"
|
||||
+ (f" (the {LONG_WINDOW} window itself rolled mid-test - boundary crossed; rerun)" if rolled else "")
|
||||
def _mint_key_of_kind(client: BudgetClient, resources: ResourceManager, kind: str) -> str:
|
||||
"""A key carrying the tight+roomy budget_limits pair, minted to roomy (100.0)
|
||||
surroundings so only the key's own windows can block."""
|
||||
match kind:
|
||||
case "personal":
|
||||
user_id = client.create_user(max_budget=100.0)
|
||||
resources.defer(lambda: client.delete_user(user_id))
|
||||
key = client.generate_key(models=[MODEL], user_id=user_id, budget_limits=_short_roomy_limits())
|
||||
case "team":
|
||||
team_id = client.create_team(alias=f"e2e-mw-team-{unique_marker()}", max_budget=100.0)
|
||||
resources.defer(lambda: client.delete_team(team_id))
|
||||
key = client.generate_key(models=[MODEL], team_id=team_id, budget_limits=_short_roomy_limits())
|
||||
case "team_member":
|
||||
team_id = client.create_team(alias=f"e2e-mw-team-{unique_marker()}", max_budget=100.0)
|
||||
resources.defer(lambda: client.delete_team(team_id))
|
||||
member_id = client.create_user(max_budget=100.0)
|
||||
resources.defer(lambda: client.delete_user(member_id))
|
||||
client.add_team_member(team_id, member_id, max_budget_in_team=100.0)
|
||||
key = client.generate_key(
|
||||
models=[MODEL], team_id=team_id, user_id=member_id, budget_limits=_short_roomy_limits()
|
||||
)
|
||||
assert is_budget_block(result), (
|
||||
f"non-budget error while waiting for {LONG_WINDOW} attribution: "
|
||||
f"status={result.status_code} body={result.body[:200]}"
|
||||
case _:
|
||||
pytest.fail(f"unknown key kind: {kind}")
|
||||
resources.defer(lambda: client.delete_key(key))
|
||||
return key
|
||||
|
||||
|
||||
class TestKeyMultiWindowBudget:
|
||||
@pytest.mark.covers("quota_management.budget.key_multi_window.blocks_then_resets")
|
||||
def test_short_window_blocks_then_resets(self, client: BudgetClient, resources: ResourceManager) -> None:
|
||||
key = client.generate_key(models=[MODEL], budget_limits=_short_roomy_limits())
|
||||
resources.defer(lambda: client.delete_key(key))
|
||||
|
||||
_assert_short_window_blocks_then_resets(client, key)
|
||||
|
||||
|
||||
@pytest.mark.covers("quota_management.budget.key_multi_window.blocks_then_resets")
|
||||
@pytest.mark.parametrize("kind", ["personal", "team", "team_member"])
|
||||
def test_short_window_blocks_then_resets_across_key_kinds(
|
||||
self, client: BudgetClient, resources: ResourceManager, kind: str
|
||||
) -> None:
|
||||
key = _mint_key_of_kind(client, resources, kind)
|
||||
|
||||
_assert_short_window_blocks_then_resets(client, key)
|
||||
|
||||
|
||||
@pytest.mark.covers("quota_management.budget.key_multi_window.blocks_then_resets")
|
||||
def test_long_window_blocks_after_short_window_resets(
|
||||
self, client: BudgetClient, resources: ResourceManager
|
||||
) -> None:
|
||||
key = client.generate_key(
|
||||
models=[MODEL],
|
||||
budget_limits=[
|
||||
BudgetWindow(budget_duration=SHORT_WINDOW, max_budget=TINY_CAP),
|
||||
BudgetWindow(budget_duration=LONG_WINDOW, max_budget=LONG_CAP),
|
||||
],
|
||||
)
|
||||
resources.defer(lambda: client.delete_key(key))
|
||||
|
||||
# 1. drive the key to get blocked by SHORT_WINDOW, assert it's budget error
|
||||
blocked = _drive_to_block(client, key)
|
||||
assert blocked.status_code == 429, f"budget block was not a 429: {blocked.status_code} {blocked.body[:200]}"
|
||||
|
||||
# 2. check the reset times of both budget windows after we drove to being blocked
|
||||
blocked_reset_at = window_reset_at(client.key_budget_windows(key), SHORT_WINDOW)
|
||||
assert blocked_reset_at is not None, "short window missing from /key/info budget_limits"
|
||||
blocked_long_reset_at = window_reset_at(client.key_budget_windows(key), LONG_WINDOW)
|
||||
assert blocked_long_reset_at is not None, "long window missing from /key/info budget_limits"
|
||||
|
||||
# 3. poll every 5s for the SHORT_WINDOW reset time until it is past it, fails if it doesnt reset
|
||||
deadline = time.monotonic() + RESET_DEADLINE_SECONDS
|
||||
while time.monotonic() < deadline:
|
||||
time.sleep(5)
|
||||
current = window_reset_at(client.key_budget_windows(key), SHORT_WINDOW)
|
||||
if current is not None and current > blocked_reset_at:
|
||||
break
|
||||
else:
|
||||
pytest.fail(
|
||||
f"{SHORT_WINDOW} window's reset_at never advanced past {blocked_reset_at} within {RESET_DEADLINE_SECONDS}s"
|
||||
)
|
||||
|
||||
# 4. short window just reset in 3, so now make a call, check that its blocked (should be blocked by LONG_WINDOW because short window reset), also make sure its budget error
|
||||
deadline = time.monotonic() + RESET_DEADLINE_SECONDS
|
||||
last_body = ""
|
||||
while time.monotonic() < deadline:
|
||||
result = _call(client, key)
|
||||
if result.ok:
|
||||
rolled = window_reset_at(client.key_budget_windows(key), LONG_WINDOW) != blocked_long_reset_at
|
||||
pytest.fail(
|
||||
f"{LONG_WINDOW} window failed to block after the {SHORT_WINDOW} window reset"
|
||||
+ (f" (the {LONG_WINDOW} window itself rolled mid-test - boundary crossed; rerun)" if rolled else "")
|
||||
)
|
||||
assert is_budget_block(result), (
|
||||
f"non-budget error while waiting for {LONG_WINDOW} attribution: "
|
||||
f"status={result.status_code} body={result.body[:200]}"
|
||||
)
|
||||
if f"over {LONG_WINDOW} budget" in result.body:
|
||||
return
|
||||
last_body = result.body
|
||||
time.sleep(5)
|
||||
pytest.fail(
|
||||
f"block never attributed to the {LONG_WINDOW} window within {RESET_DEADLINE_SECONDS}s: {last_body[:200]}"
|
||||
)
|
||||
if f"over {LONG_WINDOW} budget" in result.body:
|
||||
return
|
||||
last_body = result.body
|
||||
time.sleep(5)
|
||||
pytest.fail(
|
||||
f"block never attributed to the {LONG_WINDOW} window within {RESET_DEADLINE_SECONDS}s: {last_body[:200]}"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -51,96 +51,99 @@ def _drive_to_block(client: BudgetClient, key: str) -> StreamingResponse:
|
|||
pytest.fail("team budget never enforced before block")
|
||||
|
||||
|
||||
@pytest.mark.covers("quota_management.budget.team_multi_window.blocks_then_resets")
|
||||
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=SHORT_WINDOW, 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, models=["claude-haiku-4-5"])
|
||||
resources.defer(lambda: client.delete_key(key))
|
||||
class TestTeamMultiWindowBudget:
|
||||
@pytest.mark.covers("quota_management.budget.team_multi_window.blocks_then_resets")
|
||||
def test_team_short_window_blocks_then_resets(self, client: BudgetClient, resources: ResourceManager) -> None:
|
||||
team_id = client.create_team(
|
||||
alias=f"e2e-team-window-{unique_marker()}",
|
||||
budget_limits=[
|
||||
BudgetWindow(budget_duration=SHORT_WINDOW, 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, models=["claude-haiku-4-5"])
|
||||
resources.defer(lambda: client.delete_key(key))
|
||||
|
||||
# 1. exhaust the tight window -> litellm returns budget_exceeded
|
||||
start = time.monotonic()
|
||||
_drive_to_block(client, key)
|
||||
# 1. exhaust the tight window -> litellm returns budget_exceeded
|
||||
start = time.monotonic()
|
||||
_drive_to_block(client, key)
|
||||
|
||||
# 2. the window resets at the next wall-clock-aligned boundary (up to a window
|
||||
# after start), then the reset job (~15-20s rescheduler) zeroes the spend.
|
||||
# Allow generous headroom for that alignment + rescheduler latency; a stuck
|
||||
# rescheduler is caught by the wait-loop timeout, not this elapsed bound.
|
||||
deadline = time.monotonic() + 150
|
||||
while time.monotonic() < deadline:
|
||||
time.sleep(5)
|
||||
result = _call(client, key)
|
||||
if result.ok:
|
||||
elapsed = time.monotonic() - start
|
||||
assert elapsed < WINDOW_SECONDS + 90, 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 150s")
|
||||
# 2. the window resets at the next wall-clock-aligned boundary (up to a window
|
||||
# after start), then the reset job (~15-20s rescheduler) zeroes the spend.
|
||||
# Allow generous headroom for that alignment + rescheduler latency; a stuck
|
||||
# rescheduler is caught by the wait-loop timeout, not this elapsed bound.
|
||||
deadline = time.monotonic() + 150
|
||||
while time.monotonic() < deadline:
|
||||
time.sleep(5)
|
||||
result = _call(client, key)
|
||||
if result.ok:
|
||||
elapsed = time.monotonic() - start
|
||||
assert elapsed < WINDOW_SECONDS + 90, 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 150s")
|
||||
|
||||
|
||||
@pytest.mark.covers("quota_management.budget.team_multi_window.blocks_then_resets")
|
||||
def test_team_long_window_blocks_after_short_window_resets(client: BudgetClient, resources: ResourceManager) -> None:
|
||||
@pytest.mark.covers("quota_management.budget.team_multi_window.blocks_then_resets")
|
||||
def test_team_long_window_blocks_after_short_window_resets(
|
||||
self, client: BudgetClient, resources: ResourceManager
|
||||
) -> None:
|
||||
|
||||
# 0. key with a short budget window and a long budget window
|
||||
team_id = client.create_team(
|
||||
alias=f"e2e-team-long-window-{unique_marker()}",
|
||||
budget_limits=[
|
||||
BudgetWindow(budget_duration=SHORT_WINDOW, max_budget=TINY_CAP),
|
||||
BudgetWindow(budget_duration=LONG_WINDOW, max_budget=LONG_CAP),
|
||||
],
|
||||
)
|
||||
resources.defer(lambda: client.delete_team(team_id))
|
||||
key = client.generate_key(team_id=team_id, models=["claude-haiku-4-5"])
|
||||
resources.defer(lambda: client.delete_key(key))
|
||||
# 0. key with a short budget window and a long budget window
|
||||
team_id = client.create_team(
|
||||
alias=f"e2e-team-long-window-{unique_marker()}",
|
||||
budget_limits=[
|
||||
BudgetWindow(budget_duration=SHORT_WINDOW, max_budget=TINY_CAP),
|
||||
BudgetWindow(budget_duration=LONG_WINDOW, max_budget=LONG_CAP),
|
||||
],
|
||||
)
|
||||
resources.defer(lambda: client.delete_team(team_id))
|
||||
key = client.generate_key(team_id=team_id, models=["claude-haiku-4-5"])
|
||||
resources.defer(lambda: client.delete_key(key))
|
||||
|
||||
# 1. drive the key to being blocked, assert its blocked by budget budget_exceeded
|
||||
blocked = _drive_to_block(client, key)
|
||||
assert blocked.status_code == 429, f"budget block was not a 429: {blocked.status_code} {blocked.body[:200]}"
|
||||
# 1. drive the key to being blocked, assert its blocked by budget budget_exceeded
|
||||
blocked = _drive_to_block(client, key)
|
||||
assert blocked.status_code == 429, f"budget block was not a 429: {blocked.status_code} {blocked.body[:200]}"
|
||||
|
||||
# 2. check the the teams budget windows
|
||||
blocked_reset_at = window_reset_at(client.team_budget_windows(team_id), SHORT_WINDOW)
|
||||
assert blocked_reset_at is not None, "short window missing from /team/info budget_limits"
|
||||
blocked_long_reset_at = window_reset_at(client.team_budget_windows(team_id), LONG_WINDOW)
|
||||
assert blocked_long_reset_at is not None, "long window missing from /team/info budget_limits"
|
||||
# 2. check the the teams budget windows
|
||||
blocked_reset_at = window_reset_at(client.team_budget_windows(team_id), SHORT_WINDOW)
|
||||
assert blocked_reset_at is not None, "short window missing from /team/info budget_limits"
|
||||
blocked_long_reset_at = window_reset_at(client.team_budget_windows(team_id), LONG_WINDOW)
|
||||
assert blocked_long_reset_at is not None, "long window missing from /team/info budget_limits"
|
||||
|
||||
# 3. keep checking that the short budget window reset, if it doesnt within the deadline then fail
|
||||
deadline = time.monotonic() + RESET_DEADLINE_SECONDS
|
||||
while time.monotonic() < deadline:
|
||||
time.sleep(5)
|
||||
current = window_reset_at(client.team_budget_windows(team_id), SHORT_WINDOW)
|
||||
if current is not None and current > blocked_reset_at:
|
||||
break
|
||||
else:
|
||||
pytest.fail(
|
||||
f"team {SHORT_WINDOW} window's reset_at never advanced past "
|
||||
f"{blocked_reset_at} within {RESET_DEADLINE_SECONDS}s"
|
||||
)
|
||||
|
||||
# 4. short window just reset in 3, so now make another call, assert that the long window blocks the next call with budget_exceeded
|
||||
deadline = time.monotonic() + RESET_DEADLINE_SECONDS
|
||||
last_body = ""
|
||||
while time.monotonic() < deadline:
|
||||
result = _call(client, key)
|
||||
if result.ok:
|
||||
rolled = window_reset_at(client.team_budget_windows(team_id), LONG_WINDOW) != blocked_long_reset_at
|
||||
# 3. keep checking that the short budget window reset, if it doesnt within the deadline then fail
|
||||
deadline = time.monotonic() + RESET_DEADLINE_SECONDS
|
||||
while time.monotonic() < deadline:
|
||||
time.sleep(5)
|
||||
current = window_reset_at(client.team_budget_windows(team_id), SHORT_WINDOW)
|
||||
if current is not None and current > blocked_reset_at:
|
||||
break
|
||||
else:
|
||||
pytest.fail(
|
||||
f"team {LONG_WINDOW} window failed to block after the {SHORT_WINDOW} window reset"
|
||||
+ (f" (the {LONG_WINDOW} window itself rolled mid-test - boundary crossed; rerun)" if rolled else "")
|
||||
f"team {SHORT_WINDOW} window's reset_at never advanced past "
|
||||
f"{blocked_reset_at} within {RESET_DEADLINE_SECONDS}s"
|
||||
)
|
||||
assert is_budget_block(result), (
|
||||
f"non-budget error while waiting for {LONG_WINDOW} attribution: "
|
||||
f"status={result.status_code} body={result.body[:200]}"
|
||||
|
||||
# 4. short window just reset in 3, so now make another call, assert that the long window blocks the next call with budget_exceeded
|
||||
deadline = time.monotonic() + RESET_DEADLINE_SECONDS
|
||||
last_body = ""
|
||||
while time.monotonic() < deadline:
|
||||
result = _call(client, key)
|
||||
if result.ok:
|
||||
rolled = window_reset_at(client.team_budget_windows(team_id), LONG_WINDOW) != blocked_long_reset_at
|
||||
pytest.fail(
|
||||
f"team {LONG_WINDOW} window failed to block after the {SHORT_WINDOW} window reset"
|
||||
+ (f" (the {LONG_WINDOW} window itself rolled mid-test - boundary crossed; rerun)" if rolled else "")
|
||||
)
|
||||
assert is_budget_block(result), (
|
||||
f"non-budget error while waiting for {LONG_WINDOW} attribution: "
|
||||
f"status={result.status_code} body={result.body[:200]}"
|
||||
)
|
||||
if f"over {LONG_WINDOW} budget" in result.body:
|
||||
return
|
||||
last_body = result.body
|
||||
time.sleep(5)
|
||||
pytest.fail(
|
||||
f"team block never attributed to the {LONG_WINDOW} window within {RESET_DEADLINE_SECONDS}s: {last_body[:200]}"
|
||||
)
|
||||
if f"over {LONG_WINDOW} budget" in result.body:
|
||||
return
|
||||
last_body = result.body
|
||||
time.sleep(5)
|
||||
pytest.fail(
|
||||
f"team block never attributed to the {LONG_WINDOW} window within {RESET_DEADLINE_SECONDS}s: {last_body[:200]}"
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue