mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
test: fix stale and flaky tests across CircleCI, GHA and Buildkite (#45530)
* test(integration): poll the partition lock witness off the event loop The witness poll ran a blocking psycopg connect and query on the test's event loop, so the partition DDL task only progressed during the 20ms sleeps and missed the 3s deadline on loaded runners. Polling through asyncio.to_thread lets the DDL reach the held lock concurrently, and the deadline is 10s since it only bounds how long the DDL takes to start waiting * test(e2e): retry the reasoning turn until the model emits a reasoning item Whether gpt-5.4-mini emits a reasoning item next to a forced function call is up to the model, and it skipped it on four scheduled runs while a same-SHA rerun passed. The first turn now retries up to three times before the unchanged reasoning-item assertion * test(e2e): retry a Bedrock Converse model once when it reports no cache tokens Opus on Bedrock Converse reported zero cache tokens on half of the scheduled runs while a same-SHA rerun passed. A model that comes back uncached is asked once more inside the 5 minute window, and the non-zero cache token assertion is unchanged * test(e2e): poll until a deleted stored response stops being retrievable Azure kept serving a just-deleted streamed response for a moment, so a single retrieve did not raise. Both delete tests now poll retrieve until it returns an error and still require a 4xx * test(e2e): give the Together prompt cache five primed attempts Together documents its prefix cache as best-effort, and three primed attempts came back with zero cached tokens once while a same-SHA rerun passed * test(e2e): allow 30s for the idle RSS reading at session start A loaded router replica took longer than 10s to answer the session-start memory read, which pytest reruns cannot recover. The RSS budget assertion is unchanged * test(e2e): scope the MCP submission check to its own card The register call timed out at 15s on a loaded stack and leaked a submitted server, so the retries matched two "3 passing, 1 failing" cards. Registration gets 60s and the check reads only the submitted server's card * test(router): wait for the primary success record before counting it The shadow fan-out test waited only for the two shadow success events and then asserted exactly one primary success, which the logging worker sometimes had not delivered yet. It needed a rerun on 5 of 59 main runs * test(proxy): give the spend-log cleanup run a 1s budget A 0.25s budget could expire before the first batch on a loaded xdist worker, leaving rows_deleted at 0. One second still stops the 50-batch, 5-second loop on the deadline, which is what the test proves * test(autorouter): release the slow token count at 2.05s instead of 2.3s The count only has to finish after a quarter of the 8s worker budget, and the planner gives it 3s, so releasing at 2.3s left 0.7s of slack that a loaded runner used up. 2.05s is still past the quarter mark with almost 1s of slack * test(xai): run the reasoning-effort tests on grok-4.6 with a prompt that needs reasoning grok-4.7 now returns reasoning tokens without any reasoning text, so reasoning_content was never set and both tests failed on every scheduled run. Called directly, grok-4.6 returned reasoning text 6 of 6 times on a step-by-step arithmetic prompt but only some of the time on a bare greeting * test(integration): give the two worker-kill chaos tests the owned-proxy time budget Each test spends about 50s on the burst and then up to graceful_stop_seconds() stopping its owned proxy, which overran the 90s default pytest timeout on loaded runners. They now use the same 2 * graceful_stop_seconds() + 120 budget as other owned-proxy tests * test(integration): compare regional image responses without their created timestamp Two identical generations that straddle a second boundary differ only in created, which failed the byte-for-byte comparison. Every other field and both spend rows are still compared * test(integration): count only the S3 logger's flush task The set of asyncio tasks created while the logger starts also picked up client close finalizers left by earlier tests, so the count ranged from 1 to 8. The test now counts the periodic_flush task it owns and cancels * test(integration): send guardrail timeout probes eight at a time All 44 probes went at once to a two-worker proxy, so on a loaded runner some guardrails used their whole 1s timeout before their request left the proxy, and the sink never saw them. A different set of providers failed on each run. Eight in flight still overlaps the waits * test(integration): count a killed Arize chaos worker as gone once it is a zombie psutil reports an unreaped zombie as still running, so the 10s death check failed whenever the supervisor was slow to reap the SIGKILLed worker, and the stop then overran the 90s default timeout. The test now accepts a zombie or a missing process and has the owned-proxy time budget * test(integration): send the Typesafe connection test its real model and endpoint The test passed the proxy alias as litellm_params.model, and /health/test_connection lays the request over the stored deployment, so the alias replaced the real typesafe/ model and the check failed with "LLM Provider NOT provided" on every run since it landed in #45481. It now sends the provider model, api_base and api_key directly * test(integration): keep DB-stored models out of the usage-routing Redis read test The owned proxy inherited store_model_in_db and the job's shared database, so a model another test left behind joined the router and added its cooldown key to the MGET the test compares exactly. Reproduced locally with one /model/new model present (3 failed), and green with model loading from the database turned off * test(vertex_ai): prove the batch upload streams by laziness instead of peak memory The two tracemalloc ratio tests flaked on unrelated PRs because peak memory on a shared xdist worker includes other threads' allocations and garbage from earlier tests. They are replaced by deterministic checks of the same property: the upload stream parses and maps a row only when it is pulled, so a body whose tail is not JSON yields its valid rows first, and a Path source reflects a row rewritten on disk after the upload started. Making the parse, the output, or the file read eager fails these tests * test(integration): keep the alias test_connection call as a known bug * test: type the counting mapper and the Bedrock rerun helper * test: annotate the new test locals as Final and build them as tuples * test(vertex_ai): prove pull-driven transforms with the garbage tail alone
This commit is contained in:
parent
0519882fef
commit
cddc7cde97
20 changed files with 174 additions and 103 deletions
|
|
@ -15,7 +15,7 @@ The (feature, provider) for this cell is inferred from the file path by
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Mapping, Optional
|
||||
from typing import Any, Dict, Final, Mapping, Optional, Sequence
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -23,6 +23,7 @@ from e2e_metadata import Capability, Domain, Mode, Provider, Route, Subject, met
|
|||
from claude_code._env import require_proxy
|
||||
from claude_code.cli_driver import (
|
||||
ClaudeCLIError,
|
||||
ModelResult,
|
||||
failure_diagnostic,
|
||||
run_claude_models_parallel,
|
||||
)
|
||||
|
|
@ -62,12 +63,21 @@ def test_prompt_caching_5m_bedrock_converse(compat_result):
|
|||
upstream usage block surfaces a non-zero cache token count."""
|
||||
base_url, api_key = require_proxy(compat_result)
|
||||
|
||||
outcomes = run_claude_models_parallel(
|
||||
models=BEDROCK_CONVERSE_MODELS,
|
||||
prompt="Reply with the single word 'pong' and nothing else.",
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
def run(models: Sequence[str]) -> Dict[str, ModelResult]:
|
||||
return run_claude_models_parallel(
|
||||
models=models,
|
||||
prompt="Reply with the single word 'pong' and nothing else.",
|
||||
base_url=base_url,
|
||||
api_key=api_key,
|
||||
)
|
||||
|
||||
first: Final = run(BEDROCK_CONVERSE_MODELS)
|
||||
uncached: Final = tuple(
|
||||
model
|
||||
for model, outcome in first.items()
|
||||
if not isinstance(outcome, ClaudeCLIError) and outcome.exit_code == 0 and _cache_tokens(outcome.usage) <= 0
|
||||
)
|
||||
outcomes: Final = {**first, **run(uncached)} if uncached else first
|
||||
|
||||
failures = []
|
||||
for model in BEDROCK_CONVERSE_MODELS:
|
||||
|
|
|
|||
|
|
@ -60,7 +60,7 @@ _E2E_TEST_RAN = pytest.StashKey[bool]()
|
|||
_CALL_PASSED = pytest.StashKey[bool]()
|
||||
_IDLE_RSS = pytest.StashKey[RssCapture]()
|
||||
|
||||
IDLE_RSS_READ_TIMEOUT_SECONDS: Final = 10.0
|
||||
IDLE_RSS_READ_TIMEOUT_SECONDS: Final = 30.0
|
||||
|
||||
OPT_IN_MARKERS: Final = MappingProxyType(
|
||||
{
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ from __future__ import annotations
|
|||
import contextlib
|
||||
import json
|
||||
import threading
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass, field
|
||||
from types import MappingProxyType
|
||||
from typing import Final, cast
|
||||
|
|
@ -1031,6 +1031,7 @@ class TestResponses:
|
|||
REASONING_BACKEND: Final = "openai/gpt-5.4-mini"
|
||||
SHELL_BACKEND: Final = "openai/gpt-5.5"
|
||||
TOOL_DATE: Final = "2025-01-15"
|
||||
REASONING_ATTEMPTS: Final = 3
|
||||
|
||||
GET_TODAY_TOOL: FunctionToolParam = {
|
||||
"type": "function",
|
||||
|
|
@ -1062,6 +1063,15 @@ class TodayReport(BaseModel):
|
|||
number_of_r: str
|
||||
|
||||
|
||||
def _has_reasoning_item(response: Response) -> bool:
|
||||
return any(isinstance(item, ResponseReasoningItem) for item in response.output)
|
||||
|
||||
|
||||
def _until_reasoning_emitted(create: Callable[[], Response]) -> Response:
|
||||
attempts: Final = (create() for _ in range(REASONING_ATTEMPTS - 1))
|
||||
return next(filter(_has_reasoning_item, attempts), None) or create()
|
||||
|
||||
|
||||
class TestResponsesOpenAIHostedFeatures:
|
||||
@meta(
|
||||
Subject(
|
||||
|
|
@ -1091,14 +1101,16 @@ class TestResponsesOpenAIHostedFeatures:
|
|||
),
|
||||
}
|
||||
|
||||
first = client.responses.create(
|
||||
model=model,
|
||||
input=[question],
|
||||
tools=[GET_TODAY_TOOL],
|
||||
tool_choice={"type": "function", "name": "get_today"},
|
||||
reasoning={"effort": "medium", "summary": "auto"},
|
||||
text={"format": TODAY_REPORT_FORMAT},
|
||||
extra_body=NO_PROXY_CACHE,
|
||||
first: Final = _until_reasoning_emitted(
|
||||
lambda: client.responses.create(
|
||||
model=model,
|
||||
input=[question],
|
||||
tools=[GET_TODAY_TOOL],
|
||||
tool_choice={"type": "function", "name": "get_today"},
|
||||
reasoning={"effort": "medium", "summary": "auto"},
|
||||
text={"format": TODAY_REPORT_FORMAT},
|
||||
extra_body=NO_PROXY_CACHE,
|
||||
)
|
||||
)
|
||||
assert any(isinstance(item, ResponseReasoningItem) for item in first.output), (
|
||||
f"reasoning model returned no reasoning item: {first.output!r}"
|
||||
|
|
|
|||
|
|
@ -139,6 +139,17 @@ class TestResponsesRetrieve:
|
|||
pytest.fail(f"invalid response id expected 404, got {other!r}")
|
||||
|
||||
|
||||
def _retrieve_until_gone(client: openai.OpenAI, response_id: str) -> openai.APIStatusError:
|
||||
deadline: Final = time.monotonic() + POLL_TIMEOUT
|
||||
while time.monotonic() < deadline:
|
||||
try:
|
||||
client.responses.retrieve(response_id)
|
||||
except openai.APIStatusError as error:
|
||||
return error
|
||||
time.sleep(POLL_INTERVAL)
|
||||
raise AssertionError(f"response {response_id!r} was still retrievable {POLL_TIMEOUT}s after delete")
|
||||
|
||||
|
||||
def _register_openai(proxy: ProxyClient, resources: ResourceManager, prefix: str) -> str:
|
||||
return _register_response_deployment(proxy, resources, "openai", prefix)
|
||||
|
||||
|
|
@ -236,8 +247,8 @@ class TestStoredResponseLifecycle:
|
|||
|
||||
client.responses.delete(created.id)
|
||||
|
||||
gone: Final = pytest.raises(openai.APIStatusError, client.responses.retrieve, created.id)
|
||||
assert 400 <= gone.value.status_code < 500, f"retrieve after delete expected a 4xx: {gone.value!r}"
|
||||
gone: Final = _retrieve_until_gone(client, created.id)
|
||||
assert 400 <= gone.status_code < 500, f"retrieve after delete expected a 4xx: {gone!r}"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"deployment",
|
||||
|
|
@ -269,8 +280,8 @@ class TestStoredResponseLifecycle:
|
|||
response_id: Final = completed.response.id
|
||||
|
||||
client.responses.delete(response_id)
|
||||
gone: Final = pytest.raises(openai.APIStatusError, client.responses.retrieve, response_id)
|
||||
assert 400 <= gone.value.status_code < 500, f"retrieve after streamed delete expected a 4xx: {gone.value!r}"
|
||||
gone: Final = _retrieve_until_gone(client, response_id)
|
||||
assert 400 <= gone.status_code < 500, f"retrieve after streamed delete expected a 4xx: {gone!r}"
|
||||
|
||||
|
||||
@pytest.mark.provider_live
|
||||
|
|
|
|||
|
|
@ -66,7 +66,7 @@ REPLAY_ATTEMPTS: Final = 3
|
|||
ARITHMETIC_PROMPT = "What is 17 + 26? Answer with just the number."
|
||||
PERSON_PROMPT = "Invent a fictional person."
|
||||
CACHE_PREFIX_FACTS: Final = 600
|
||||
CACHE_ATTEMPTS: Final = 3
|
||||
CACHE_ATTEMPTS: Final = 5
|
||||
|
||||
PERSON_RESPONSE_FORMAT: dict[str, object] = {
|
||||
"type": "json_schema",
|
||||
|
|
|
|||
|
|
@ -84,6 +84,7 @@ const test = base.extend<Fixtures>({
|
|||
url: localServerUrl,
|
||||
transport: "http",
|
||||
},
|
||||
timeout: 60_000,
|
||||
});
|
||||
expect(registration.status(), await registration.text()).toBe(201);
|
||||
|
||||
|
|
@ -129,6 +130,11 @@ test("Submission Rules panel shows saved rules, preloads the editor, and save ke
|
|||
|
||||
await page.reload();
|
||||
await page.getByRole("tab", { name: "Submitted MCPs" }).click();
|
||||
await expect(page.getByText(submittedServerName, { exact: true })).toBeVisible();
|
||||
await expect(page.getByText("3 passing, 1 failing", { exact: true })).toBeVisible();
|
||||
const submittedCard = page
|
||||
.locator("div")
|
||||
.filter({ has: page.getByRole("heading", { name: submittedServerName, exact: true }) })
|
||||
.filter({ hasText: /\d+ passing, \d+ failing/ })
|
||||
.last();
|
||||
await expect(submittedCard.getByRole("heading", { name: submittedServerName, exact: true })).toBeVisible();
|
||||
await expect(submittedCard.getByText("3 passing, 1 failing", { exact: true })).toBeVisible();
|
||||
});
|
||||
|
|
|
|||
|
|
@ -51,10 +51,11 @@ async def test_real_partition_ddl_survives_witnessed_lock_and_is_idempotent() ->
|
|||
operation: Final = asyncio.create_task(
|
||||
manager.ensure_partitions(PartitionConnection(database), lambda: 7000)
|
||||
)
|
||||
wait_deadline: Final = time.monotonic() + 3
|
||||
wait_deadline: Final = time.monotonic() + 10
|
||||
try:
|
||||
while True:
|
||||
witnesses: Final = read_rows(
|
||||
witnesses: Final = await asyncio.to_thread(
|
||||
read_rows,
|
||||
"SELECT a.pid, extract(epoch FROM "
|
||||
"clock_timestamp()-a.query_start)::double precision AS age "
|
||||
"FROM pg_stat_activity a WHERE %s = ANY(pg_blocking_pids(a.pid)) "
|
||||
|
|
|
|||
|
|
@ -44,6 +44,7 @@ from _openinference_support import (
|
|||
_sse_json_values,
|
||||
)
|
||||
from integration._support.client import Gateway, eventually
|
||||
from integration._support.process import graceful_stop_seconds
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
|
||||
|
||||
|
|
@ -333,6 +334,14 @@ def test_arize_otel_v2_f2_slow_sink_does_not_deadlock(gateway: Gateway, tmp_path
|
|||
_matching_marker_span(rig.destination, marker)
|
||||
|
||||
|
||||
def _has_exited(process: psutil.Process) -> bool:
|
||||
try:
|
||||
return process.status() == psutil.STATUS_ZOMBIE
|
||||
except psutil.NoSuchProcess:
|
||||
return True
|
||||
|
||||
|
||||
@pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
|
||||
def test_arize_otel_v2_f3_one_proxy_worker_can_die(gateway: Gateway, tmp_path: Path) -> None:
|
||||
markers: Final = tuple("f3-" + uuid.uuid4().hex for _ in range(8))
|
||||
release: Final = threading.Event()
|
||||
|
|
@ -380,7 +389,7 @@ def test_arize_otel_v2_f3_one_proxy_worker_can_die(gateway: Gateway, tmp_path: P
|
|||
)
|
||||
assert observed >= 2, observed
|
||||
workers[0].kill()
|
||||
assert eventually(lambda: not workers[0].is_running(), bool, seconds=10), workers[0]
|
||||
assert eventually(lambda: _has_exited(workers[0]), bool, seconds=10), workers[0]
|
||||
release.set()
|
||||
in_flight: Final = tuple(
|
||||
(marker, future.result(timeout=60)) for marker, future in zip(markers, futures, strict=True)
|
||||
|
|
|
|||
|
|
@ -4,7 +4,8 @@ Each guardrail is configured against an owned sink that records the request and
|
|||
~20s. With `timeout: 1` the outbound call must abort near the bound, so the chat round trip
|
||||
completes in seconds instead of waiting on the sink. A control guardrail without `timeout`
|
||||
points at a sink path that sleeps ~3s and must wait for the reply, proving unset keeps the
|
||||
handler default. All probes are sent concurrently so their waits overlap.
|
||||
handler default. Probes are sent eight at a time so their waits overlap without starving the
|
||||
proxy's two workers.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -35,6 +36,7 @@ from integration._support.wire import Reply, Request, wire_server
|
|||
SLOW_SECONDS: Final = 20
|
||||
FAST_SECONDS: Final = 3
|
||||
BOUND_SECONDS: Final = 8
|
||||
PROBE_CONCURRENCY: Final = 8
|
||||
TOKEN_PATH: Final = "/token"
|
||||
TOKEN_REPLY: Final = json.dumps(
|
||||
{"access_token": "synthetic-google-token", "expires_in": 3600, "token_type": "Bearer"}
|
||||
|
|
@ -410,7 +412,7 @@ def outcomes(rig: Rig) -> Mapping[str, Outcome]:
|
|||
values: Final = tuple(_provider_values())
|
||||
names: Final = (*(value[0] for value in values), "control-generic")
|
||||
exchanges: Final = (*(value[4] for value in values), False)
|
||||
with ThreadPoolExecutor(max_workers=len(names)) as pool:
|
||||
with ThreadPoolExecutor(max_workers=PROBE_CONCURRENCY) as pool:
|
||||
results: Final = tuple(pool.map(partial(_chat, rig), names, exchanges))
|
||||
return MappingProxyType(dict(zip(names, results, strict=True)))
|
||||
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ import pytest
|
|||
import yaml
|
||||
from integration._support.client import Gateway, eventually, object_value
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.process import owned_proxy_process
|
||||
from integration._support.process import graceful_stop_seconds, owned_proxy_process
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
|
|
@ -129,6 +129,7 @@ async def test_passthrough_upstream_outage_mid_burst_still_logs_errors_once(gate
|
|||
), response.text
|
||||
|
||||
|
||||
@pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
|
||||
async def test_passthrough_worker_sigkill_leaves_sibling_serving_and_logging(gateway: Gateway, tmp_path: Path) -> None:
|
||||
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
path: Final = tmp_path / "chaos-kill.yaml"
|
||||
|
|
|
|||
|
|
@ -256,6 +256,26 @@ def test_each_provider_gets_its_own_path_key_and_body_and_is_billed_from_the_cos
|
|||
|
||||
|
||||
def test_test_connection_evaluation_mode_uses_typesafe_decisions_path(gateway: Gateway) -> None:
|
||||
provider: Final = _PROVIDERS[1]
|
||||
with gateway.scenario() as scenario:
|
||||
handle: Final = _register(scenario, _answer_body(provider))
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/health/test_connection",
|
||||
{
|
||||
"litellm_params": {"model": provider.model, "api_base": handle.api_base(), "api_key": provider.api_key},
|
||||
"mode": "evaluation",
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json()["status"] == "success", response.text
|
||||
(call,) = _upstream_calls(gateway, handle)
|
||||
assert call["path"] == f"/{handle.scenario_id}{provider.path}"
|
||||
|
||||
|
||||
def test_test_connection_by_configured_alias_probes_the_stored_model(gateway: Gateway) -> None:
|
||||
pytest.skip("BUG: /health/test_connection given only a configured alias sends the alias as the model")
|
||||
provider: Final = _PROVIDERS[1]
|
||||
with gateway.scenario() as scenario:
|
||||
handle: Final = _register(scenario, _answer_body(provider))
|
||||
|
|
|
|||
|
|
@ -70,7 +70,7 @@ async def test_s3_sync_and_async_uploads_pass_independent_wire_verification(monk
|
|||
prior: Final = asyncio.all_tasks()
|
||||
logger: Final = S3Logger(s3_bucket_name="integration-bucket", s3_region_name="us-east-1", s3_endpoint_url=wire.url,
|
||||
s3_aws_access_key_id=ACCESS, s3_aws_secret_access_key=SECRET, s3_callback_params_override={})
|
||||
owned: Final = asyncio.all_tasks() - prior
|
||||
owned: Final = {task for task in asyncio.all_tasks() - prior if task.get_coro().__qualname__ == "CustomBatchLogger.periodic_flush"}
|
||||
assert len(owned) == 1
|
||||
try:
|
||||
for mode in ("sync", "async"):
|
||||
|
|
|
|||
|
|
@ -11,8 +11,9 @@ from typing import Final
|
|||
|
||||
import httpx
|
||||
import psutil
|
||||
import pytest
|
||||
from integration._support.client import Gateway, eventually, gateway_from_environment
|
||||
from integration._support.process import owned_proxy_process
|
||||
from integration._support.process import graceful_stop_seconds, owned_proxy_process
|
||||
|
||||
_RESPONSES_JSON_SCENARIO: Final = "audit-chaos-responses"
|
||||
_RESPONSES_STREAM_SCENARIO: Final = "audit-chaos-responses-stream"
|
||||
|
|
@ -145,6 +146,7 @@ def _worker_pids(log: Path) -> tuple[int, ...]:
|
|||
return tuple(int(match) for match in _WORKER_PID.findall(log.read_text()))
|
||||
|
||||
|
||||
@pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
|
||||
def test_proxy_survives_worker_kill_mid_burst(tmp_path: Path) -> None:
|
||||
with gateway_from_environment() as upstream_gateway:
|
||||
directory: Final = tmp_path
|
||||
|
|
|
|||
|
|
@ -191,6 +191,10 @@ def test_proxy_usage_routing_reads_cooldown_tpm_then_rpm_from_redis(
|
|||
)
|
||||
config: Final = {
|
||||
**configuration,
|
||||
"general_settings": {
|
||||
**JSON_OBJECT.validate_python(configuration["general_settings"]),
|
||||
"store_model_in_db": False,
|
||||
},
|
||||
"model_list": _deployment_list(model_name, f"{wire.url}/v1", deployment_ids),
|
||||
"router_settings": {
|
||||
"routing_strategy": "usage-based-routing-v2",
|
||||
|
|
@ -205,7 +209,7 @@ def test_proxy_usage_routing_reads_cooldown_tpm_then_rpm_from_redis(
|
|||
with owned_proxy(
|
||||
bootstrap,
|
||||
tmp_path,
|
||||
{"REDIS_HOST": cache.host, "REDIS_PORT": str(cache.port)},
|
||||
{"REDIS_HOST": cache.host, "REDIS_PORT": str(cache.port), "STORE_MODEL_IN_DB": "False"},
|
||||
config=config_path,
|
||||
) as candidate:
|
||||
eventually(
|
||||
|
|
|
|||
|
|
@ -341,6 +341,10 @@ def test_invalid_multiplier_is_refused_at_registration(gateway: Gateway) -> None
|
|||
assert wire.drain() == ()
|
||||
|
||||
|
||||
def _body_without_created(response: httpx.Response) -> dict[str, object]:
|
||||
return {key: value for key, value in response.json().items() if key != "created"}
|
||||
|
||||
|
||||
def test_identical_regional_generations_each_bill_the_uplift_once(gateway: Gateway) -> None:
|
||||
with wire_server(_scripted()) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(gateway, scenario, wire, location={"vertex_location": _REGION})
|
||||
|
|
@ -348,7 +352,7 @@ def test_identical_regional_generations_each_bill_the_uplift_once(gateway: Gatew
|
|||
second_id, second = _generate(gateway, model)
|
||||
assert first.status_code == 200, first.text
|
||||
assert second.status_code == 200, second.text
|
||||
assert first.content == second.content, (first.text, second.text)
|
||||
assert _body_without_created(first) == _body_without_created(second), (first.text, second.text)
|
||||
rows: Final = _spend_rows((first_id, second_id))
|
||||
assert rows == {first_id: pytest.approx(_REGIONAL_TOKEN_COST), second_id: pytest.approx(_REGIONAL_TOKEN_COST)}
|
||||
assert len(wire.drain()) == 2
|
||||
|
|
|
|||
|
|
@ -42,8 +42,10 @@ def test_xai_message_name_filtering():
|
|||
class TestXAIReasoningEffort(BaseReasoningLLMTests):
|
||||
def get_base_completion_call_args(self):
|
||||
return {
|
||||
"model": "xai/grok-4.7",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"model": "xai/grok-4.6",
|
||||
"messages": [
|
||||
{"role": "user", "content": "What is 17 times 23? Work it out step by step, then give the number."}
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -8,8 +8,8 @@ dicts, joined output), which keeps peak memory bounded on large uploads.
|
|||
These tests lock in the behaviour that would regress if the streaming path were
|
||||
replaced by a list-based pipeline:
|
||||
1. Byte-for-byte output parity with a list pipeline (wire format).
|
||||
2. The streaming transform peaks at a clear fraction of a list pipeline on the
|
||||
same input (relative differential, robust to GC noise).
|
||||
2. The upload stream parses and transforms a row only when the HTTP client
|
||||
pulls it, and a Path source is read from disk as the upload proceeds.
|
||||
3. ``get_object_name`` only parses the first JSONL row, so a payload whose
|
||||
later rows are not valid JSON does not raise.
|
||||
4. A tuple-wrapped file handle uploaded through the real create_file ordering
|
||||
|
|
@ -29,6 +29,8 @@ import json
|
|||
import tempfile
|
||||
import time
|
||||
import tracemalloc
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
|
@ -254,44 +256,21 @@ class TestGetObjectNameLazyParse:
|
|||
assert "gemini-2.5-flash" in object_name
|
||||
|
||||
|
||||
class TestStreamingPeakMemory:
|
||||
"""
|
||||
Differential guard: the streaming transform must stay well under the peak
|
||||
that a list pipeline incurs on the same input. If the hot path builds full
|
||||
intermediate lists, the streaming assertion fails.
|
||||
|
||||
The assertion that matters is the *relative* one: ``streaming_peak`` must be
|
||||
a clear fraction of ``list_peak`` on the identical input. Absolute
|
||||
``tracemalloc`` ratios drift with GC timing and the live set carried in from
|
||||
earlier tests, so they make poor CI gates; the relative comparison cancels
|
||||
that shared noise and is exactly what regresses (toward 1.0) when the hot
|
||||
path builds full intermediate lists. ``gc.collect()`` before each
|
||||
measurement removes any garbage the previous run left behind.
|
||||
"""
|
||||
|
||||
@pytest.mark.no_cover
|
||||
def test_streaming_peak_well_below_list_pipeline(self):
|
||||
cfg = VertexAIFilesConfig()
|
||||
raw = _make_openai_jsonl_bytes(8000)
|
||||
content_str = raw.decode("utf-8")
|
||||
|
||||
def drain_stream():
|
||||
# Consume the upload body one row at a time, as the chunked uploader
|
||||
# does, without accumulating it.
|
||||
for _ in _OpenAIToVertexBatchUploadStream(raw, cfg._map_openai_to_vertex_params).iter_bytes():
|
||||
pass
|
||||
|
||||
streaming_peak = _measure_peak(drain_stream)
|
||||
list_peak = _measure_peak(lambda: _reference_vertex_jsonl_string(cfg, content_str))
|
||||
|
||||
# Core guard: the lazily consumed streaming body peaks well under a list
|
||||
# pipeline that materializes every transformed row. Building full
|
||||
# intermediate lists in the hot path pushes this ratio back toward 1.0.
|
||||
assert streaming_peak < list_peak * 0.6, (
|
||||
f"streaming peak {streaming_peak} not a clear win over list pipeline "
|
||||
f"{list_peak} (ratio {streaming_peak / list_peak:.2f})"
|
||||
class TestStreamingLaziness:
|
||||
def test_upload_stream_transforms_rows_only_as_they_are_pulled(self) -> None:
|
||||
cfg: Final = VertexAIFilesConfig()
|
||||
valid_rows: Final = 3
|
||||
raw: Final = _make_openai_jsonl_bytes(valid_rows) + b"\n" + b"\n".join(b"not-json" for _ in range(8000))
|
||||
chunks: Final = _OpenAIToVertexBatchUploadStream(raw, cfg._map_openai_to_vertex_params).iter_bytes()
|
||||
pulled: Final = tuple(next(chunks) for _ in range(valid_rows))
|
||||
custom_ids: Final = tuple(
|
||||
_get_litellm_batch_custom_id_from_labels(json.loads(chunk)["request"]["labels"]) for chunk in pulled
|
||||
)
|
||||
|
||||
assert custom_ids == ("request-0", "request-1", "request-2")
|
||||
with pytest.raises(json.JSONDecodeError):
|
||||
next(chunks)
|
||||
|
||||
def test_get_object_name_does_not_scale_with_payload(self):
|
||||
cfg = VertexAIFilesConfig()
|
||||
raw = _make_openai_jsonl_bytes(8000)
|
||||
|
|
@ -346,35 +325,36 @@ class TestPathSourcedStreaming:
|
|||
first_labels = json.loads(lines[0])["request"]["labels"]
|
||||
assert _get_litellm_batch_custom_id_from_labels(first_labels) == "request-0"
|
||||
|
||||
@pytest.mark.no_cover
|
||||
def test_path_source_peak_stays_below_list_pipeline(self, tmp_path):
|
||||
cfg = VertexAIFilesConfig()
|
||||
path, raw = self._write_jsonl(tmp_path, 8000)
|
||||
data = self._batch_request(path)
|
||||
content_str = raw.decode("utf-8")
|
||||
|
||||
def drain_stream():
|
||||
cfg.get_complete_file_url(
|
||||
api_base=None,
|
||||
api_key=None,
|
||||
model="",
|
||||
optional_params={},
|
||||
litellm_params={"gcs_bucket_name": "test-bucket"},
|
||||
data=data,
|
||||
)
|
||||
out = cfg.transform_create_file_request(
|
||||
model="", create_file_data=data, optional_params={}, litellm_params={}
|
||||
)
|
||||
for _ in _upload_stream(out).iter_bytes():
|
||||
pass
|
||||
|
||||
streaming_peak = _measure_peak(drain_stream)
|
||||
list_peak = _measure_peak(lambda: _reference_vertex_jsonl_string(cfg, content_str))
|
||||
|
||||
assert streaming_peak < list_peak * 0.3, (
|
||||
f"path-sourced streaming peak {streaming_peak} not a clear win over list pipeline "
|
||||
f"{list_peak} (ratio {streaming_peak / list_peak:.2f})"
|
||||
def test_path_source_is_read_as_the_upload_is_pulled(self, tmp_path: Path) -> None:
|
||||
cfg: Final = VertexAIFilesConfig()
|
||||
n_rows: Final = 8000
|
||||
path, raw = self._write_jsonl(tmp_path, n_rows)
|
||||
data: Final = self._batch_request(path)
|
||||
cfg.get_complete_file_url(
|
||||
api_base=None,
|
||||
api_key=None,
|
||||
model="",
|
||||
optional_params={},
|
||||
litellm_params={"gcs_bucket_name": "test-bucket"},
|
||||
data=data,
|
||||
)
|
||||
out: Final = cfg.transform_create_file_request(
|
||||
model="", create_file_data=data, optional_params={}, litellm_params={}
|
||||
)
|
||||
chunks: Final = _upload_stream(out).iter_bytes()
|
||||
first: Final = next(chunks)
|
||||
|
||||
with open(path, "r+b") as handle:
|
||||
handle.seek(raw.index(b'"request-6000"'))
|
||||
handle.write(b'"rewrote-6000"')
|
||||
custom_ids: Final = tuple(
|
||||
_get_litellm_batch_custom_id_from_labels(json.loads(chunk)["request"]["labels"])
|
||||
for chunk in (first, *chunks)
|
||||
)
|
||||
|
||||
assert len(custom_ids) == n_rows
|
||||
assert custom_ids[0] == "request-0"
|
||||
assert custom_ids[6000] == "rewrote-6000"
|
||||
|
||||
def test_path_source_stream_is_reiterable(self, tmp_path):
|
||||
cfg = VertexAIFilesConfig()
|
||||
|
|
|
|||
|
|
@ -503,7 +503,7 @@ async def test_native_count_finishing_after_quarter_worker_budget_keeps_plan_and
|
|||
|
||||
async def count(model: str, api_key: str, body: Mapping[str, JsonValue]) -> int:
|
||||
if not release.is_set():
|
||||
asyncio.get_running_loop().call_later(2.3, release.set)
|
||||
asyncio.get_running_loop().call_later(2.05, release.set)
|
||||
await release.wait()
|
||||
return await _count(model, api_key, body)
|
||||
|
||||
|
|
|
|||
|
|
@ -926,7 +926,7 @@ async def test_run_budget_stops_the_loop_and_leaves_the_backlog_for_the_next_run
|
|||
|
||||
cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
|
||||
started_at = time.monotonic()
|
||||
result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, time.monotonic() + 0.25)
|
||||
result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, time.monotonic() + 1.0)
|
||||
elapsed = time.monotonic() - started_at
|
||||
|
||||
assert result.stop_reason == "budget_exhausted"
|
||||
|
|
|
|||
|
|
@ -63,6 +63,12 @@ async def _wait_for_shadow_successes(logger: _RecordingLogger, expected: int, ti
|
|||
await asyncio.sleep(0.05)
|
||||
|
||||
|
||||
async def _wait_for_all_successes(logger: _RecordingLogger, expected: int, timeout: float = 5.0) -> None:
|
||||
deadline: Final = time.monotonic() + timeout
|
||||
while len(logger.success_kwargs) < expected and time.monotonic() < deadline:
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
|
||||
def _wait_for_shadow_successes_sync(logger: _RecordingLogger, expected: int, timeout: float = 5.0) -> None:
|
||||
deadline: Final = time.monotonic() + timeout
|
||||
while len(logger.shadow_successes()) < expected and time.monotonic() < deadline:
|
||||
|
|
@ -242,6 +248,7 @@ async def test_multiple_shadow_targets_fan_out_async(recording_logger):
|
|||
)
|
||||
assert [chunk async for chunk in response]
|
||||
await _wait_for_shadow_successes(recording_logger, expected=2)
|
||||
await _wait_for_all_successes(recording_logger, expected=3)
|
||||
|
||||
shadow_successes = recording_logger.shadow_successes()
|
||||
model_groups = sorted(call["litellm_params"]["metadata"]["model_group"] for call in shadow_successes)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue