mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
8f31c7171d
* fix(router): keep a lone deployment warm under router allowed_fails unless the policy names the exception Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(router): give cooldown fixtures a sibling deployment so allowed_fails=0 still forces a cooldown * fix(router): keep benching a lone deployment that has a fallback group or an empty deployment policy * fix(router): accept dictionary fallback targets in the lone deployment check Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
626 lines
28 KiB
Python
626 lines
28 KiB
Python
from __future__ import annotations
|
|
|
|
import itertools
|
|
import json
|
|
import re
|
|
import uuid
|
|
from collections.abc import Callable, Iterator, Mapping, Sequence
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from contextlib import ExitStack
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Final
|
|
|
|
import anthropic
|
|
import httpx
|
|
import openai
|
|
import pytest
|
|
import yaml
|
|
from integration._support.client import GATEWAY_LIMITS, Gateway, Scenario, eventually, gateway_from_environment
|
|
from integration._support.openai_wire import chat_reply, responses_reply
|
|
from integration._support.process import graceful_stop_seconds, owned_proxy_process
|
|
from integration._support.prometheus_series import PROXY_FAILURES, Sample, scrape
|
|
from integration._support.redis_process import OwnedRedis, owned_redis
|
|
from integration._support.wire import Reply, Request, Wire, wire_server
|
|
from integration.routing.test_payment_required_mapping import (
|
|
OPENAI_MODEL,
|
|
PAYMENT_REQUIRED,
|
|
anthropic_client,
|
|
anthropic_error,
|
|
anthropic_params,
|
|
assert_payment_required_row,
|
|
chat_body,
|
|
deployment,
|
|
is_model_info_probe,
|
|
model_list_reply,
|
|
new_group,
|
|
new_marker,
|
|
openai_client,
|
|
openai_params,
|
|
openai_refusal,
|
|
post,
|
|
provider_calls,
|
|
)
|
|
from pydantic import JsonValue
|
|
from redis import Redis
|
|
|
|
pytestmark: Final = [pytest.mark.timeout(300), pytest.mark.shares_worker_state]
|
|
|
|
COOLDOWN_SECONDS: Final = 300
|
|
RELOAD_SECONDS: Final = 2
|
|
NO_DEPLOYMENTS: Final = "No deployments available"
|
|
COOLED_DOWN: Final = "litellm_deployment_cooled_down_total"
|
|
PRIMARY_GROUP: Final = new_group()
|
|
FALLBACK_GROUP: Final = new_group()
|
|
MARKER: Final = re.compile(rb"pr402-[0-9a-f]{32}(?:-\d+)?")
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class Outcome:
|
|
marker: str
|
|
status: int
|
|
detail: str
|
|
call_id: str
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class Call:
|
|
marker: str
|
|
path: str
|
|
body: Mapping[str, JsonValue]
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class CooldownRig:
|
|
gateway: Gateway
|
|
cache: OwnedRedis
|
|
log: Path
|
|
|
|
def cooled(self, deployment_id: str) -> bool:
|
|
return cooled(self.cache, deployment_id)
|
|
|
|
|
|
def cooled(cache: OwnedRedis, deployment_id: str) -> bool:
|
|
with Redis(host=cache.host, port=cache.port) as client:
|
|
return int(client.exists(f"deployment:{deployment_id}:cooldown")) == 1
|
|
|
|
|
|
def marker_in(request: Request) -> str:
|
|
found: Final = MARKER.search(request.body)
|
|
assert found is not None, (request.method, request.target, request.headers, request.body[:300])
|
|
return found.group(0).decode()
|
|
|
|
|
|
def markers_received(wire: Wire) -> tuple[str, ...]:
|
|
return tuple(marker_in(request) for request in provider_calls(wire))
|
|
|
|
|
|
def refusing_reply(request: Request) -> Reply:
|
|
return anthropic_error(402, f"scripted 402 {marker_in(request)}")
|
|
|
|
|
|
def healthy_reply(request: Request) -> Reply:
|
|
if is_model_info_probe(request):
|
|
return model_list_reply()
|
|
marker: Final = marker_in(request)
|
|
body: Final = json.loads(request.body)
|
|
stream: Final = isinstance(body, dict) and body.get("stream") is True
|
|
if request.target.endswith("/responses"):
|
|
return responses_reply(f"resp_{marker}", OPENAI_MODEL, f"served {marker}", stream=stream)
|
|
return chat_reply(f"chatcmpl-{marker}", OPENAI_MODEL, f"served {marker}", stream=stream)
|
|
|
|
|
|
def router_settings(cache: OwnedRedis, **extra: JsonValue) -> dict[str, JsonValue]:
|
|
return {
|
|
"num_retries": 0,
|
|
"cooldown_time": COOLDOWN_SECONDS,
|
|
"redis_host": cache.host,
|
|
"redis_port": cache.port,
|
|
**extra,
|
|
}
|
|
|
|
|
|
def cooldown_config(
|
|
directory: Path,
|
|
settings: Mapping[str, JsonValue],
|
|
*,
|
|
model_list: Sequence[Mapping[str, JsonValue]] = (),
|
|
prometheus: bool = True,
|
|
) -> Path:
|
|
base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
|
assert isinstance(base, dict)
|
|
litellm_settings: Final = (
|
|
{**base["litellm_settings"], "callbacks": ["prometheus"]} if prometheus else base["litellm_settings"]
|
|
)
|
|
config: Final = {
|
|
**base,
|
|
"litellm_settings": litellm_settings,
|
|
"router_settings": dict(settings),
|
|
**({"model_list": [dict(entry) for entry in model_list]} if model_list else {}),
|
|
}
|
|
path: Final = directory / f"cooldown-{uuid.uuid4().hex[:8]}.yaml"
|
|
path.write_text(yaml.safe_dump(config))
|
|
return path
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[CooldownRig]:
|
|
directory: Final = tmp_path_factory.mktemp("payment-required-cooldown")
|
|
with ExitStack() as stack:
|
|
upstream: Final = stack.enter_context(gateway_from_environment())
|
|
cache: Final = stack.enter_context(owned_redis(directory))
|
|
settings: Final = router_settings(cache, fallbacks=[{PRIMARY_GROUP: [FALLBACK_GROUP]}])
|
|
owned: Final = stack.enter_context(
|
|
owned_proxy_process(
|
|
upstream,
|
|
directory,
|
|
{"PROXY_CONFIG_RELOAD_INTERVAL_SECONDS": str(RELOAD_SECONDS)},
|
|
config=cooldown_config(directory, settings),
|
|
remove_environment=("PROMETHEUS_MULTIPROC_DIR",),
|
|
)
|
|
)
|
|
yield CooldownRig(owned.gateway, cache, owned.log)
|
|
|
|
|
|
def two_deployments(scenario: Scenario, group: str, refusing: Wire, healthy: Wire) -> str:
|
|
refused_id: Final = deployment(scenario, group, anthropic_params(refusing))
|
|
deployment(scenario, group, openai_params(healthy))
|
|
return refused_id
|
|
|
|
|
|
def raw_chat(gateway: Gateway, group: str, marker: str, **extra: JsonValue) -> Outcome:
|
|
response: Final = post(gateway, "/v1/chat/completions", chat_body(group, marker, **extra))
|
|
return Outcome(marker, response.status_code, response.text, response.headers["x-litellm-call-id"])
|
|
|
|
|
|
def raw_call(gateway: Gateway, call: Call) -> Outcome:
|
|
response: Final = post(gateway, call.path, call.body)
|
|
return Outcome(call.marker, response.status_code, response.text, response.headers["x-litellm-call-id"])
|
|
|
|
|
|
def openai_chat(client: openai.OpenAI, group: str, marker: str) -> Outcome:
|
|
try:
|
|
raw: Final = client.chat.completions.with_raw_response.create(
|
|
model=group, messages=[{"role": "user", "content": marker}]
|
|
)
|
|
except openai.APIStatusError as error:
|
|
return Outcome(marker, error.status_code, error.message, error.response.headers["x-litellm-call-id"])
|
|
return Outcome(marker, raw.status_code, raw.parse().id, raw.headers["x-litellm-call-id"])
|
|
|
|
|
|
def anthropic_message(client: anthropic.Anthropic, group: str, marker: str) -> Outcome:
|
|
try:
|
|
raw: Final = client.messages.with_raw_response.create(
|
|
model=group, max_tokens=32, messages=[{"role": "user", "content": marker}]
|
|
)
|
|
except anthropic.APIStatusError as error:
|
|
return Outcome(marker, error.status_code, error.message, error.response.headers["x-litellm-call-id"])
|
|
block: Final = raw.parse().content[0]
|
|
assert isinstance(block, anthropic.types.TextBlock), block
|
|
return Outcome(marker, raw.status_code, block.text, raw.headers["x-litellm-call-id"])
|
|
|
|
|
|
def until_refused(attempt: Callable[[str], Outcome], tries: int = 40) -> tuple[Outcome, ...]:
|
|
def walk() -> Iterator[Outcome]:
|
|
for _ in range(tries):
|
|
outcome: Final = attempt(new_marker())
|
|
yield outcome
|
|
if outcome.status != 200:
|
|
return
|
|
|
|
history: Final = tuple(walk())
|
|
assert history[-1].status != 200, "the refusing deployment was never picked"
|
|
return history
|
|
|
|
|
|
def assert_refused(outcome: Outcome) -> None:
|
|
assert outcome.status == 402, (outcome.status, outcome.detail)
|
|
assert PAYMENT_REQUIRED in outcome.detail, outcome.detail
|
|
|
|
|
|
def assert_served(outcomes: Sequence[Outcome], identity: Callable[[str], str]) -> None:
|
|
assert all(outcome.status == 200 for outcome in outcomes), [(o.status, o.detail[:200]) for o in outcomes]
|
|
assert all(identity(outcome.marker) in outcome.detail for outcome in outcomes), [o.detail[:200] for o in outcomes]
|
|
|
|
|
|
def chat_identity(marker: str) -> str:
|
|
return f"chatcmpl-{marker}"
|
|
|
|
|
|
def served_text(marker: str) -> str:
|
|
return f"served {marker}"
|
|
|
|
|
|
def response_identity(marker: str) -> str:
|
|
return f"resp_{marker}"
|
|
|
|
|
|
def markers(outcomes: Sequence[Outcome]) -> tuple[str, ...]:
|
|
return tuple(outcome.marker for outcome in outcomes)
|
|
|
|
|
|
def failure_samples(gateway: Gateway, group: str) -> tuple[Sample, ...]:
|
|
return tuple(
|
|
sample
|
|
for sample in scrape(gateway)
|
|
if sample.name == PROXY_FAILURES and sample.labels.get("requested_model") == group
|
|
)
|
|
|
|
|
|
def cooled_down_samples(gateway: Gateway, deployment_id: str) -> tuple[Sample, ...]:
|
|
return tuple(
|
|
sample
|
|
for sample in scrape(gateway)
|
|
if sample.name == COOLED_DOWN and sample.labels.get("model_id") == deployment_id
|
|
)
|
|
|
|
|
|
def assert_payment_required_metrics(gateway: Gateway, group: str, deployment_id: str) -> None:
|
|
failures: Final = eventually(lambda: failure_samples(gateway, group), lambda found: len(found) >= 1, seconds=10)
|
|
assert [(sample.labels["exception_class"], sample.labels["exception_status"]) for sample in failures] == [
|
|
("Anthropic.PaymentRequiredError", "402")
|
|
], failures
|
|
cooldowns: Final = eventually(
|
|
lambda: cooled_down_samples(gateway, deployment_id), lambda found: len(found) >= 1, seconds=10
|
|
)
|
|
assert [sample.labels["exception_status"] for sample in cooldowns] == ["402"], cooldowns
|
|
|
|
|
|
def test_chat_openai_sdk_402_cools_the_deployment_and_the_healthy_one_serves_after(rig: CooldownRig) -> None:
|
|
group: Final = new_group()
|
|
with (
|
|
wire_server(refusing_reply) as refusing,
|
|
wire_server(healthy_reply) as healthy,
|
|
rig.gateway.scenario() as scenario,
|
|
):
|
|
refused_id: Final = two_deployments(scenario, group, refusing, healthy)
|
|
client: Final = openai_client(rig.gateway)
|
|
history: Final = until_refused(lambda marker: openai_chat(client, group, marker))
|
|
assert_refused(history[-1])
|
|
assert_served(history[:-1], chat_identity)
|
|
eventually(lambda: rig.cooled(refused_id), lambda seen: seen, seconds=10)
|
|
assert markers_received(refusing) == (history[-1].marker,)
|
|
assert markers_received(healthy) == markers(history[:-1])
|
|
served: Final = tuple(openai_chat(client, group, new_marker()) for _ in range(20))
|
|
assert_served(served, chat_identity)
|
|
assert markers_received(healthy) == markers(served)
|
|
assert markers_received(refusing) == ()
|
|
assert_payment_required_row(history[-1].call_id, refused_id)
|
|
assert_payment_required_metrics(rig.gateway, group, refused_id)
|
|
|
|
|
|
def test_messages_anthropic_sdk_402_cools_the_deployment_and_the_healthy_one_serves_after(rig: CooldownRig) -> None:
|
|
group: Final = new_group()
|
|
with (
|
|
wire_server(refusing_reply) as refusing,
|
|
wire_server(healthy_reply) as healthy,
|
|
rig.gateway.scenario() as scenario,
|
|
):
|
|
refused_id: Final = two_deployments(scenario, group, refusing, healthy)
|
|
client: Final = anthropic_client(rig.gateway)
|
|
history: Final = until_refused(lambda marker: anthropic_message(client, group, marker))
|
|
assert_refused(history[-1])
|
|
assert_served(history[:-1], served_text)
|
|
eventually(lambda: rig.cooled(refused_id), lambda seen: seen, seconds=10)
|
|
assert markers_received(refusing) == (history[-1].marker,)
|
|
assert markers_received(healthy) == markers(history[:-1])
|
|
served: Final = tuple(anthropic_message(client, group, new_marker()) for _ in range(20))
|
|
assert_served(served, served_text)
|
|
assert markers_received(healthy) == markers(served)
|
|
assert markers_received(refusing) == ()
|
|
assert_payment_required_row(history[-1].call_id, refused_id)
|
|
|
|
|
|
@pytest.mark.parametrize("stream", (False, True), ids=("json", "stream"))
|
|
def test_responses_httpx_402_cools_the_deployment_and_the_healthy_one_serves_after(
|
|
rig: CooldownRig, stream: bool
|
|
) -> None:
|
|
group: Final = new_group()
|
|
with (
|
|
wire_server(refusing_reply) as refusing,
|
|
wire_server(healthy_reply) as healthy,
|
|
rig.gateway.scenario() as scenario,
|
|
):
|
|
refused_id: Final = two_deployments(scenario, group, refusing, healthy)
|
|
|
|
def attempt(marker: str) -> Outcome:
|
|
return raw_call(
|
|
rig.gateway, Call(marker, "/v1/responses", {"model": group, "input": marker, "stream": stream})
|
|
)
|
|
|
|
history: Final = until_refused(attempt)
|
|
assert_refused(history[-1])
|
|
assert_served(history[:-1], response_identity)
|
|
eventually(lambda: rig.cooled(refused_id), lambda seen: seen, seconds=10)
|
|
assert markers_received(refusing) == (history[-1].marker,)
|
|
assert markers_received(healthy) == markers(history[:-1])
|
|
served: Final = tuple(attempt(new_marker()) for _ in range(20))
|
|
assert_served(served, response_identity)
|
|
assert markers_received(healthy) == markers(served)
|
|
assert markers_received(refusing) == ()
|
|
assert_payment_required_row(history[-1].call_id, refused_id)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"model_info",
|
|
({}, {"allowed_fails_policy": None}, {"allowed_fails_policy": {}}),
|
|
ids=("missing", "null", "empty"),
|
|
)
|
|
def test_a_lone_402_deployment_without_a_policy_is_never_cooled(
|
|
rig: CooldownRig, model_info: dict[str, JsonValue]
|
|
) -> None:
|
|
group: Final = new_group()
|
|
with wire_server(refusing_reply) as refusing, rig.gateway.scenario() as scenario:
|
|
refused_id: Final = deployment(scenario, group, anthropic_params(refusing), model_info=model_info)
|
|
refused: Final = tuple(raw_chat(rig.gateway, group, new_marker()) for _ in range(5))
|
|
for outcome in refused:
|
|
assert_refused(outcome)
|
|
assert_payment_required_row(outcome.call_id, refused_id)
|
|
assert not rig.cooled(refused_id)
|
|
sixth: Final = raw_chat(rig.gateway, group, new_marker())
|
|
assert_refused(sixth)
|
|
assert markers_received(refusing) == (*markers(refused), sixth.marker)
|
|
|
|
|
|
@pytest.mark.timeout(900)
|
|
def test_a_router_allowed_fails_policy_cools_a_lone_402_deployment_past_its_count(tmp_path: Path) -> None:
|
|
group: Final = new_group()
|
|
with ExitStack() as stack:
|
|
upstream: Final = stack.enter_context(gateway_from_environment())
|
|
cache: Final = stack.enter_context(owned_redis(tmp_path))
|
|
settings: Final = router_settings(cache, allowed_fails_policy={"BadRequestErrorAllowedFails": 1})
|
|
owned: Final = stack.enter_context(
|
|
owned_proxy_process(
|
|
upstream,
|
|
tmp_path,
|
|
{},
|
|
config=cooldown_config(tmp_path, settings),
|
|
remove_environment=("PROMETHEUS_MULTIPROC_DIR",),
|
|
)
|
|
)
|
|
refusing: Final = stack.enter_context(wire_server(refusing_reply))
|
|
scenario: Final = stack.enter_context(owned.gateway.scenario())
|
|
refused_id: Final = deployment(scenario, group, anthropic_params(refusing))
|
|
first: Final = raw_chat(owned.gateway, group, new_marker())
|
|
second: Final = raw_chat(owned.gateway, group, new_marker())
|
|
assert_refused(first)
|
|
assert_refused(second)
|
|
eventually(lambda: cooled(cache, refused_id), lambda seen: seen, seconds=10)
|
|
third: Final = raw_chat(owned.gateway, group, new_marker())
|
|
assert third.status == 429 and NO_DEPLOYMENTS in third.detail, (third.status, third.detail)
|
|
assert markers_received(refusing) == (first.marker, second.marker)
|
|
|
|
|
|
def failing_reply(request: Request) -> Reply:
|
|
return model_list_reply() if is_model_info_probe(request) else anthropic_error(500, "upstream down")
|
|
|
|
|
|
def assert_failed(outcome: Outcome) -> None:
|
|
assert outcome.status == 500, (outcome.status, outcome.detail)
|
|
assert "upstream down" in outcome.detail, outcome.detail
|
|
|
|
|
|
@pytest.mark.timeout(900)
|
|
def test_a_router_allowed_fails_count_leaves_a_lone_deployment_warm_but_benches_one_with_a_sibling_or_fallback(
|
|
tmp_path: Path,
|
|
) -> None:
|
|
lone_group: Final = new_group()
|
|
pair_group: Final = new_group()
|
|
primary_group: Final = new_group()
|
|
backup_group: Final = new_group()
|
|
with ExitStack() as stack:
|
|
upstream: Final = stack.enter_context(gateway_from_environment())
|
|
cache: Final = stack.enter_context(owned_redis(tmp_path))
|
|
settings: Final = router_settings(cache, allowed_fails=1, fallbacks=[{primary_group: [backup_group]}])
|
|
owned: Final = stack.enter_context(
|
|
owned_proxy_process(
|
|
upstream,
|
|
tmp_path,
|
|
{},
|
|
config=cooldown_config(tmp_path, settings),
|
|
remove_environment=("PROMETHEUS_MULTIPROC_DIR",),
|
|
)
|
|
)
|
|
failing: Final = stack.enter_context(wire_server(failing_reply))
|
|
healthy: Final = stack.enter_context(wire_server(healthy_reply))
|
|
scenario: Final = stack.enter_context(owned.gateway.scenario())
|
|
lone_id: Final = deployment(scenario, lone_group, anthropic_params(failing))
|
|
failed: Final = tuple(raw_chat(owned.gateway, lone_group, new_marker()) for _ in range(4))
|
|
for outcome in failed:
|
|
assert_failed(outcome)
|
|
assert not cooled(cache, lone_id)
|
|
assert markers_received(failing) == markers(failed)
|
|
|
|
paired_id: Final = two_deployments(scenario, pair_group, failing, healthy)
|
|
history: Final = tuple(
|
|
itertools.takewhile(
|
|
lambda outcome: not cooled(cache, paired_id),
|
|
(raw_chat(owned.gateway, pair_group, new_marker()) for _ in range(40)),
|
|
)
|
|
)
|
|
assert cooled(cache, paired_id), [(outcome.status, outcome.detail[:80]) for outcome in history]
|
|
assert sum(outcome.status == 500 for outcome in history) == 1
|
|
served: Final = tuple(raw_chat(owned.gateway, pair_group, new_marker()) for _ in range(10))
|
|
assert_served(served, chat_identity)
|
|
assert markers_received(healthy) == (*markers(o for o in history if o.status == 200), *markers(served))
|
|
|
|
primary_id: Final = deployment(scenario, primary_group, anthropic_params(failing))
|
|
deployment(scenario, backup_group, openai_params(healthy))
|
|
failing.drain()
|
|
healthy.drain()
|
|
fell_back: Final = tuple(
|
|
itertools.takewhile(
|
|
lambda outcome: not cooled(cache, primary_id),
|
|
(raw_chat(owned.gateway, primary_group, new_marker()) for _ in range(40)),
|
|
)
|
|
)
|
|
assert cooled(cache, primary_id) and fell_back, [(o.status, o.detail[:80]) for o in fell_back]
|
|
assert_served(fell_back, chat_identity)
|
|
backed: Final = tuple(raw_chat(owned.gateway, primary_group, new_marker()) for _ in range(4))
|
|
assert_served(backed, chat_identity)
|
|
received: Final = markers_received(failing)
|
|
assert received[: len(fell_back)] == markers(fell_back)
|
|
assert len(received) <= len(fell_back) + 1, received
|
|
assert {*markers(fell_back), *markers(backed)} <= set(markers_received(healthy))
|
|
|
|
|
|
def test_a_deployment_allowed_fails_policy_cools_a_lone_402_deployment(rig: CooldownRig) -> None:
|
|
group: Final = new_group()
|
|
policy: Final = {"allowed_fails_policy": {"BadRequestErrorAllowedFails": 0}}
|
|
with wire_server(refusing_reply) as refusing, rig.gateway.scenario() as scenario:
|
|
refused_id: Final = deployment(scenario, group, anthropic_params(refusing), model_info=policy)
|
|
first: Final = raw_chat(rig.gateway, group, new_marker())
|
|
assert_refused(first)
|
|
eventually(lambda: rig.cooled(refused_id), lambda seen: seen, seconds=10)
|
|
second: Final = raw_chat(rig.gateway, group, new_marker())
|
|
assert second.status == 429 and NO_DEPLOYMENTS in second.detail, (second.status, second.detail)
|
|
assert markers_received(refusing) == (first.marker,)
|
|
|
|
|
|
def test_a_plain_allowed_fails_count_leaves_a_lone_402_deployment_warm(rig: CooldownRig) -> None:
|
|
group: Final = new_group()
|
|
with wire_server(refusing_reply) as refusing, rig.gateway.scenario() as scenario:
|
|
refused_id: Final = deployment(scenario, group, anthropic_params(refusing), model_info={"allowed_fails": 0})
|
|
refused: Final = tuple(raw_chat(rig.gateway, group, new_marker()) for _ in range(5))
|
|
for outcome in refused:
|
|
assert_refused(outcome)
|
|
assert not rig.cooled(refused_id)
|
|
sixth: Final = raw_chat(rig.gateway, group, new_marker())
|
|
assert_refused(sixth)
|
|
assert markers_received(refusing) == (*markers(refused), sixth.marker)
|
|
|
|
|
|
def test_an_openai_compatible_402_cools_the_deployment_too(rig: CooldownRig) -> None:
|
|
group: Final = new_group()
|
|
with (
|
|
wire_server(openai_refusal) as refusing,
|
|
wire_server(healthy_reply) as healthy,
|
|
rig.gateway.scenario() as scenario,
|
|
):
|
|
refused_id: Final = deployment(scenario, group, openai_params(refusing))
|
|
deployment(scenario, group, openai_params(healthy))
|
|
history: Final = until_refused(lambda marker: raw_chat(rig.gateway, group, marker))
|
|
assert history[-1].status == 402 and "scripted 402" in history[-1].detail, history[-1]
|
|
assert_served(history[:-1], chat_identity)
|
|
eventually(lambda: rig.cooled(refused_id), lambda seen: seen, seconds=10)
|
|
assert markers_received(refusing) == (history[-1].marker,)
|
|
served: Final = tuple(raw_chat(rig.gateway, group, new_marker()) for _ in range(20))
|
|
assert_served(served, chat_identity)
|
|
assert markers_received(refusing) == ()
|
|
assert markers_received(healthy) == (*markers(history[:-1]), *markers(served))
|
|
|
|
|
|
def test_a_lone_402_primary_with_a_fallback_group_stays_warm_and_the_fallback_answers(rig: CooldownRig) -> None:
|
|
with (
|
|
wire_server(refusing_reply) as refusing,
|
|
wire_server(healthy_reply) as healthy,
|
|
rig.gateway.scenario() as scenario,
|
|
):
|
|
primary_id: Final = deployment(scenario, PRIMARY_GROUP, anthropic_params(refusing))
|
|
deployment(scenario, FALLBACK_GROUP, openai_params(healthy))
|
|
served: Final = tuple(raw_chat(rig.gateway, PRIMARY_GROUP, new_marker()) for _ in range(3))
|
|
assert_served(served, chat_identity)
|
|
assert markers_received(refusing) == markers(served)
|
|
assert markers_received(healthy) == markers(served)
|
|
assert not rig.cooled(primary_id)
|
|
|
|
|
|
def test_a_deployment_policy_added_while_402_traffic_flows_starts_the_cooldown(rig: CooldownRig) -> None:
|
|
group: Final = new_group()
|
|
with wire_server(refusing_reply) as refusing, rig.gateway.scenario() as scenario:
|
|
refused_id: Final = deployment(scenario, group, anthropic_params(refusing))
|
|
before: Final = tuple(raw_chat(rig.gateway, group, new_marker()) for _ in range(3))
|
|
for outcome in before:
|
|
assert_refused(outcome)
|
|
assert not rig.cooled(refused_id)
|
|
policy: Final = {"id": refused_id, "allowed_fails_policy": {"BadRequestErrorAllowedFails": 0}}
|
|
patched: Final = rig.gateway.request("PATCH", f"/model/{refused_id}/update", {"model_info": policy})
|
|
assert patched.status_code == 200, patched.text
|
|
|
|
def probe() -> tuple[int, bool]:
|
|
return raw_chat(rig.gateway, group, new_marker()).status, rig.cooled(refused_id)
|
|
|
|
eventually(probe, lambda seen: seen[1], seconds=RELOAD_SECONDS * 10)
|
|
after: Final = raw_chat(rig.gateway, group, new_marker())
|
|
assert after.status == 429 and NO_DEPLOYMENTS in after.detail, (after.status, after.detail)
|
|
assert len(markers_received(refusing)) >= 4
|
|
|
|
|
|
def test_a_cooled_402_deployment_comes_back_after_its_own_cooldown_time(rig: CooldownRig) -> None:
|
|
group: Final = new_group()
|
|
with (
|
|
wire_server(refusing_reply) as refusing,
|
|
wire_server(healthy_reply) as healthy,
|
|
rig.gateway.scenario() as scenario,
|
|
):
|
|
refused_id: Final = deployment(scenario, group, {**anthropic_params(refusing), "cooldown_time": 2})
|
|
deployment(scenario, group, openai_params(healthy))
|
|
history: Final = until_refused(lambda marker: raw_chat(rig.gateway, group, marker))
|
|
assert_refused(history[-1])
|
|
eventually(lambda: rig.cooled(refused_id), lambda seen: seen, seconds=10)
|
|
eventually(lambda: rig.cooled(refused_id), lambda seen: not seen, seconds=10)
|
|
again: Final = until_refused(lambda marker: raw_chat(rig.gateway, group, marker))
|
|
assert_refused(again[-1])
|
|
assert markers_received(refusing) == (history[-1].marker, again[-1].marker)
|
|
|
|
|
|
def burst_call(group: str, index: int) -> Call:
|
|
marker: Final = f"{new_marker()}-{index}"
|
|
stream: Final = index % 2 == 1
|
|
match index % 3:
|
|
case 0:
|
|
return Call(marker, "/v1/chat/completions", chat_body(group, marker, stream=stream))
|
|
case 1:
|
|
body: Final = {"model": group, "max_tokens": 32, "messages": [{"role": "user", "content": marker}]}
|
|
return Call(marker, "/v1/messages", {**body, "stream": stream})
|
|
case _:
|
|
return Call(marker, "/v1/responses", {"model": group, "input": marker, "stream": stream})
|
|
|
|
|
|
def ready_workers(log: Path) -> int:
|
|
return log.read_text().count("Application startup complete.")
|
|
|
|
|
|
@pytest.mark.timeout(900)
|
|
def test_chaos_burst_cooldown_holds_across_both_workers_and_a_worker_kill(tmp_path: Path) -> None:
|
|
group: Final = new_group()
|
|
refused_id: Final = f"refusing-{uuid.uuid4().hex[:12]}"
|
|
healthy_id: Final = f"healthy-{uuid.uuid4().hex[:12]}"
|
|
with ExitStack() as stack:
|
|
upstream: Final = stack.enter_context(gateway_from_environment())
|
|
cache: Final = stack.enter_context(owned_redis(tmp_path))
|
|
refusing: Final = stack.enter_context(wire_server(refusing_reply))
|
|
healthy: Final = stack.enter_context(wire_server(healthy_reply))
|
|
model_list: Final = (
|
|
{"model_name": group, "litellm_params": anthropic_params(refusing), "model_info": {"id": refused_id}},
|
|
{"model_name": group, "litellm_params": openai_params(healthy), "model_info": {"id": healthy_id}},
|
|
)
|
|
config: Final = cooldown_config(tmp_path, router_settings(cache), model_list=model_list, prometheus=False)
|
|
owned: Final = stack.enter_context(owned_proxy_process(upstream, tmp_path, {}, config=config, workers=2))
|
|
eventually(lambda: ready_workers(owned.log), lambda ready: ready == 2, seconds=graceful_stop_seconds())
|
|
client: Final = stack.enter_context(
|
|
httpx.Client(base_url=owned.gateway.client.base_url, timeout=60, trust_env=False, limits=GATEWAY_LIMITS)
|
|
)
|
|
patient: Final = Gateway(client, owned.gateway.key, owned.gateway.upstream_url)
|
|
calls: Final = tuple(burst_call(group, index) for index in range(30))
|
|
with ThreadPoolExecutor(max_workers=30) as pool:
|
|
outcomes: Final = tuple(pool.map(lambda call: raw_call(patient, call), calls))
|
|
assert {outcome.status for outcome in outcomes} <= {200, 402}, [(o.status, o.detail[:200]) for o in outcomes]
|
|
refused: Final = tuple(outcome for outcome in outcomes if outcome.status == 402)
|
|
served: Final = tuple(outcome for outcome in outcomes if outcome.status == 200)
|
|
assert refused, "the burst never reached the refusing deployment"
|
|
for outcome in refused:
|
|
assert_refused(outcome)
|
|
assert all(outcome.marker in outcome.detail for outcome in served), [o.detail[:200] for o in served]
|
|
assert sorted(markers_received(refusing)) == sorted(markers(refused))
|
|
assert sorted(markers_received(healthy)) == sorted(markers(served))
|
|
eventually(lambda: cooled(cache, refused_id), lambda seen: seen, seconds=10)
|
|
warm: Final = tuple(raw_chat(patient, group, new_marker()) for _ in range(20))
|
|
assert_served(warm, chat_identity)
|
|
assert markers_received(refusing) == ()
|
|
assert markers_received(healthy) == markers(warm)
|
|
owned.kill_worker(min(owned.worker_pids()))
|
|
after_kill: Final = tuple(raw_chat(patient, group, new_marker()) for _ in range(20))
|
|
assert_served(after_kill, chat_identity)
|
|
assert markers_received(refusing) == ()
|
|
assert markers_received(healthy) == markers(after_kill)
|