litellm/tests/e2e/perf/test_perf_cases.py

"""One perf run per `PerformanceTestCase` in `cases/`: check the gateway was deployed with the case's
model groups and guardrails, build its teams, users and keys, replay its deterministic plan from a separate Locust
process against the out-of-process mock providers, then check every `Limit` against what the
clients saw, what the mocks received and how the workers' memory moved."""

from __future__ import annotations

import time
from typing import Final

import pytest
from cases import CASES
from idp import Keycloak
from lifecycle import ResourceManager
from memory_readings import read_rss_everywhere
from models import SpendLogsParams
from proxy_client import ProxyClient
from setup.datastore import read_db_counters, read_redis_counters, redis_monitor
from setup.observe import (
    Observation,
    evaluate,
    junit_properties,
    read_guardrail_calls,
    read_mock_requests,
    render,
    reset_mock,
    set_mock_latency,
)
from setup.performance_test_case import PerformanceTestCase
from setup.plan import build_plan
from setup.podcpu import gateway_cpu_sampler
from setup.provision import MOCK_PROVIDERS, build_setup, guardrail_url, mock_url
from setup.run import run_plan

pytestmark = [pytest.mark.perf, pytest.mark.quiet_stack]
_SPEND_SAMPLE: Final = 40


def _spend_logged_share(proxy: ProxyClient, call_ids: tuple[str, ...]) -> float | None:
    """Share of a sample of served requests whose spend row the gateway wrote, read back by the
    request id the gateway returned in `x-litellm-call-id`."""
    if not call_ids:
        return None
    found = 0
    for call_id in call_ids:
        rows = proxy.spend_logs(SpendLogsParams(request_id=call_id))
        found += 1 if rows else 0
    return found / len(call_ids)


@pytest.mark.parametrize("case", CASES, ids=[case.name for case in CASES])
def test_perf_case(
    case: PerformanceTestCase,
    proxy: ProxyClient,
    idp: Keycloak,
    resources: ResourceManager,
    request: pytest.FixtureRequest,
) -> None:
    plan: Final = build_plan(case)
    mock_urls: Final = {provider: mock_url(provider) for provider in MOCK_PROVIDERS}
    for url in (*mock_urls.values(), guardrail_url()):
        _ = reset_mock(url)
    for url in mock_urls.values():
        _ = set_mock_latency(url, case.mock_latency)
    setup_started: Final = time.monotonic()
    setup: Final = build_setup(case, proxy=proxy, idp=idp, resources=resources)
    rss_before: Final = read_rss_everywhere(proxy)
    db_before: Final = read_db_counters()
    redis_before: Final = read_redis_counters()
    run_started: Final = time.monotonic()
    with gateway_cpu_sampler() as cpu, redis_monitor() as monitor:
        outcome: Final = run_plan(case, plan, setup)
    run_finished: Final = time.monotonic()
    db_after: Final = read_db_counters()
    redis_after: Final = read_redis_counters()
    rss_after: Final = read_rss_everywhere(proxy)
    rss_read: Final = time.monotonic()
    assert outcome.records, f"locust wrote no records (exit {outcome.locust_exit}):\n{outcome.locust_output}"
    served_call_ids: Final = tuple(
        record.call_id for record in outcome.records if record.ok and not record.blocked and record.call_id
    )
    spend_share: Final = _spend_logged_share(
        proxy, served_call_ids[:: max(1, len(served_call_ids) // _SPEND_SAMPLE)][:_SPEND_SAMPLE]
    )
    spend_read: Final = time.monotonic()
    mock_requests: Final = tuple(seen for url in mock_urls.values() for seen in read_mock_requests(url))
    guardrail_calls: Final = tuple(read_guardrail_calls(guardrail_url())) if case.guardrails else ()
    mocks_read: Final = time.monotonic()
    observation: Final = Observation(
        records=outcome.records,
        mock_requests=mock_requests,
        guardrail_calls=guardrail_calls,
        rss_before=rss_before,
        rss_after=rss_after,
        wall_seconds=outcome.wall_seconds,
        spend_logged_share=spend_share,
        db_before=db_before,
        db_after=db_after,
        redis_before=redis_before,
        redis_after=redis_after,
        redis_commands=monitor.commands,
        gateway_cpu_seconds=cpu.seconds,
    )
    rows, exceeded_limits = evaluate(case, observation)
    item: Final = request.node  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]  # pytest: untyped
    assert isinstance(item, pytest.Item)
    item.user_properties[:] = [
        entry for entry in item.user_properties if not entry[0].startswith("perf_")
    ]  # rebind-ok: a --reruns retry replaces the first attempt's numbers instead of appending a second set
    item.user_properties.extend(junit_properties(rows))
    print(render(case, observation, rows))
    print(
        f"  phases: setup {run_started - setup_started:.0f}s, locust {run_finished - run_started:.0f}s, "
        f"observe {time.monotonic() - run_finished:.0f}s (rss {rss_read - run_finished:.0f}s, spend sample "
        f"{spend_read - rss_read:.0f}s, mock logs {mocks_read - spend_read:.0f}s); locust exit {outcome.locust_exit}"
    )
    assert len(outcome.records) == len(plan), (
        f"{len(plan)} planned, {len(outcome.records)} recorded; locust said:\n{outcome.locust_output}"
    )
    assert not exceeded_limits, "limits exceeded (observed > max):\n" + "\n".join(exceeded_limits)