mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-26 06:03:19 +00:00
fc27dfdc7d
* test(integration): give each owned proxy a private database and Redis by default owned_proxy and owned_proxy_process now boot on a scratch database and a per-proxy redis-server unless the caller opts into the worker's state with state=OwnedState.SHARED_WITH_WORKER, passes its own DATABASE_URL, or strips DATABASE_URL. Under the replica profile the private database gets the writer and reader roles so the reader path stays covered. Gateway carries the database_url its proxy writes to, and the Scenario cleanup checks and delete_key_if_present read from it, so they keep checking something on a private database. The 189 test files with owned proxy sites that depend on the worker's database or Redis opt in with a shares_worker_state module marker, which a module-scoped conftest fixture turns into the shared state. The 7 helper modules with such sites opt in per call, since conftest fixtures and other test files reach them. * test(integration): reach session fixtures with the shared marker and keep the worker retry default The module-scoped fixture that applied shares_worker_state ran after session-scoped fixtures, so proxies those booted went private in marked files. The marker is now applied around each marked item's setup, call and teardown, which covers fixtures of every scope. Owned proxies on the worker database inherited the worker gateway's router_settings.num_retries of 0 from it. A private proxy now gets the same default in its config copy unless the config sets num_retries itself. Five providers files read the owned proxy's spend rows or response cache through the worker's database or Redis and opt into sharing.
77 lines
2.8 KiB
Python
77 lines
2.8 KiB
Python
import uuid
|
|
from pathlib import Path
|
|
from typing import Final
|
|
|
|
import pytest
|
|
import yaml
|
|
from integration._support.client import Gateway, eventually, string_value
|
|
from integration._support.database import read_rows
|
|
from integration._support.process import owned_proxy
|
|
from pydantic import JsonValue
|
|
|
|
pytestmark: Final = pytest.mark.shares_worker_state
|
|
|
|
WRITE_STATEMENT_MAX_BYTES: Final = 200_000
|
|
MESSAGES_PER_REQUEST: Final = 60
|
|
MESSAGE_CHARACTERS: Final = 2_000
|
|
REQUESTS: Final = 4
|
|
|
|
|
|
def _config_storing_prompts(tmp_path: Path) -> Path:
|
|
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
|
config["general_settings"]["store_prompts_in_spend_logs"] = True
|
|
path: Final = tmp_path / "store-prompts.yaml"
|
|
path.write_text(yaml.safe_dump(config))
|
|
return path
|
|
|
|
|
|
def _prompt_messages(marker: str) -> list[JsonValue]:
|
|
return [
|
|
{"role": "user", "content": f"{marker}-{index}-".ljust(MESSAGE_CHARACTERS, "x")}
|
|
for index in range(MESSAGES_PER_REQUEST)
|
|
]
|
|
|
|
|
|
def _persisted(request_ids: tuple[str, ...]) -> list[dict[str, JsonValue]]:
|
|
placeholders: Final = ", ".join("%s" for _ in request_ids)
|
|
return read_rows(
|
|
"SELECT request_id, xmin::text AS statement, octet_length(proxy_server_request::text) AS stored_bytes "
|
|
f'FROM "LiteLLM_SpendLogs" WHERE request_id IN ({placeholders}) ORDER BY request_id',
|
|
request_ids,
|
|
)
|
|
|
|
|
|
@pytest.mark.covers("quota_management.spend_tracking.prompt_rows_are_written_in_byte_bounded_statements")
|
|
def test_prompt_carrying_spend_rows_flushed_together_are_written_in_byte_bounded_statements(
|
|
gateway: Gateway, tmp_path: Path
|
|
) -> None:
|
|
with (
|
|
gateway.scenario() as scenario,
|
|
owned_proxy(
|
|
gateway,
|
|
tmp_path,
|
|
{
|
|
"SPEND_LOG_WRITE_BATCH_MAX_BYTES": str(WRITE_STATEMENT_MAX_BYTES),
|
|
"SPEND_LOG_QUEUE_POLL_INTERVAL": "15",
|
|
},
|
|
config=_config_storing_prompts(tmp_path),
|
|
) as owned,
|
|
):
|
|
model: Final = scenario.model()
|
|
request_ids: Final = tuple(
|
|
string_value(
|
|
owned.post(
|
|
"/v1/chat/completions",
|
|
{"model": model, "messages": _prompt_messages(f"integration-prompt-{uuid.uuid4().hex}")},
|
|
)["id"]
|
|
)
|
|
for _ in range(REQUESTS)
|
|
)
|
|
assert len(set(request_ids)) == REQUESTS, request_ids
|
|
rows: Final = eventually(lambda: _persisted(request_ids), lambda values: len(values) == REQUESTS, seconds=70)
|
|
stored_bytes: Final = tuple(row["stored_bytes"] for row in rows)
|
|
assert all(
|
|
isinstance(size, int) and WRITE_STATEMENT_MAX_BYTES // 2 < size < WRITE_STATEMENT_MAX_BYTES
|
|
for size in stored_bytes
|
|
), rows
|
|
assert len({row["statement"] for row in rows}) == REQUESTS, rows
|