litellm/tests/e2e/perf/test_reliability_memory_e2e.py

"""Live e2e: a few hundred requests that fail before any provider answers must not
grow the proxy's resident memory past a loose ceiling once the proxy is warm. It reads
the proxy's own RSS, so it runs in the perf lane, where nothing else is hitting the stack.

The regression this guards shipped in v1.100.0: every retry breadcrumb copied the
whole request and the copies nested into one router-global list, so a proxy under
retry-heavy failing traffic grew until it was OOM-killed. The traffic here has that
shape: a model group whose deployments refuse at the socket (an unreachable base
URL, with a literal key since the perf stack carries no provider keys) with
cooldown_time 0 so the router keeps retrying them, per-request retries, and a
fallback group that refuses the same way, each request carrying a long chat
transcript so every whole-request copy costs hundreds of containers instead of a
handful. Under the stack's cooldown policy a deployment that fails a handful of
times in a row is benched (a bad-credential 401 included), the router answers "No
deployments available" without retrying, and the retry loop that leaks stops running;
cooldown_time 0 keeps it running.

Four identical batches run back to back. The first two are the warmup that grows
the proxy's caches and allocator arenas, the last two are the measured window, and
each worker's RSS may grow at most a loose ceiling over that window. RSS cannot be
held to a tight per-batch budget: on the x86 release gate (rc/1.107.0 and v1.105.1
on isolated stacks, 28 batches each) the allocator kept setting new high-water marks
for the first several batches, about 60 MB on the third batch of a fresh stack and
more after the rest of the suite, with the gc-tracked object count flat to 0.1%, and
only settled to a couple of MB per batch after twenty. So the ceiling catches a gross
leak and nothing slower. RSS is read through /debug/memory/summary on every
configured replica; a burst of failing calls leaves a transient bulge of garbage that
gc reclaims within seconds, so each checkpoint samples until no new worker has
answered for a settle window and keeps the lowest reading per worker. A checkpoint
runs before the warmup and after every batch, and each batch's settled growth per
worker rides out as a repeated `rss_batch_growth` JUnit property as soon as that
batch settles, so a run that fails keeps the curve up to the batch it failed in. The
ceiling is judged per worker (by replica address, hostname and pid, since pods in
their own pid namespaces report the same pids) so each worker is compared with
itself, and the checkpoints around the measured window must see the same workers: a
single load-balanced address reaches the workers behind it one answer at a time, and
a worker that answered only one checkpoint would otherwise drop out of the
comparison, which is where a leaking worker could hide.

The stored-request half of that regression (every failing request's spend log carrying a
copy of the request per failed attempt) is a deterministic size check, so it lives in
tests/integration/spend/test_failing_request_stored_snapshot.py.

A cheaper check holds the idle footprint: every worker's RSS as the harness
read it at collection time, before this pytest process sent any traffic (see
conftest.pytest_collection_finish), must sit under a fixed budget. On the
release gate that is a fresh stack right after its readiness gate, one worker per
gateway replica, so the reading is what a DB-backed boot costs on its own. A
v1.100.x worker with a database idled at 886 MB RSS where v1.101.0rc1 idled at
544 MB on the same database (1.3 GB against 560 MB at the pod level, under a
2 GiB limit): the generated Prisma client at prisma-client-py's default recursive
type depth, 91k TypedDict classes that v1.101.0 cut to 19k with
recursive_type_depth = -1. The budget starts at the rc1 reading plus headroom and
E2E_MEMORY_IDLE_RSS_BUDGET_MB overrides it; a later session on the same stack (the
changed-files workflow's repeat passes, a developer's local loop) measures a proxy
already warmed by traffic, which that headroom also has to cover.
"""

from __future__ import annotations

import json
import time
from collections.abc import Callable, Mapping, Sequence
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final

import pytest

from e2e_config import (
    MEMORY_CONCURRENCY,
    MEMORY_IDLE_RSS_BUDGET_MB,
    MEMORY_REQUESTS_PER_PHASE,
    MEMORY_RETRIES_PER_REQUEST,
    MEMORY_RSS_CEILING_MB,
    MEMORY_RSS_SAMPLE_INTERVAL_SECONDS,
    MEMORY_RSS_SETTLE_SAMPLES,
    MEMORY_TRANSCRIPT_TURNS,
    PROXY_REPLICA_URLS,
    unique_marker,
)
from e2e_metadata import Domain, Mode, Provider, Subject, meta
from lifecycle import ResourceManager
from memory_readings import RssCapture, RssReading, WorkerKey, read_rss_everywhere
from models import ChatMessage, LiteLLMParamsBody, ReliabilityChatBody, RouterSettingsOverride
from proxy_client import ProxyClient

pytestmark = [pytest.mark.e2e, pytest.mark.perf]

REFUSING_MODEL: Final = "openai/gpt-5.5"

DEPLOYMENTS_PER_GROUP: Final = 2
RSS_SAMPLE_CAP: Final = 4 * MEMORY_RSS_SETTLE_SAMPLES
WARMUP_BATCHES: Final = 2
MEASURED_BATCHES: Final = 2
FIRST_MEASURED_BATCH: Final = WARMUP_BATCHES + 1
LAST_BATCH: Final = WARMUP_BATCHES + MEASURED_BATCHES


@dataclass(frozen=True, slots=True)
class FailedCall:
    status_code: int
    seconds: float
    body_head: str


@dataclass(frozen=True, slots=True)
class WorkerGrowth:
    warm: RssReading
    after: RssReading

    @property
    def growth_mb(self) -> float:
        return self.after.ram_usage_mb - self.warm.ram_usage_mb


def _create_never_benched_refusing_deployment(proxy: ProxyClient, name: str) -> str:
    return proxy.create_model(
        name,
        LiteLLMParamsBody(
            model=REFUSING_MODEL, api_key="sk-unreachable", api_base="http://127.0.0.1:9/v1", cooldown_time=0
        ),
    )


def _register_refusing_group(proxy: ProxyClient, resources: ResourceManager, name: str) -> None:
    for model_id in tuple(_create_never_benched_refusing_deployment(proxy, name) for _ in range(DEPLOYMENTS_PER_GROUP)):
        resources.defer(lambda model_id=model_id: proxy.delete_model(model_id))


def _transcript(turns: int) -> tuple[ChatMessage, ...]:
    return tuple(
        ChatMessage(role=role, content=f"turn {turn} {role}")
        for turn in range(turns)
        for role in ("user", "assistant")
    )


TRANSCRIPT: Final = _transcript(MEMORY_TRANSCRIPT_TURNS)


def _fail_once(proxy: ProxyClient, key: str, model: str, override: RouterSettingsOverride) -> FailedCall:
    started: Final = time.perf_counter()
    resp: Final = proxy.transport.send(
        "/chat/completions",
        headers=proxy.transport.bearer(key),
        json=ReliabilityChatBody(
            model=model,
            messages=[*TRANSCRIPT, ChatMessage(role="user", content=f"memory regression {unique_marker()}")],
            max_tokens=512,
            router_settings_override=override,
            cache={"no-cache": True},
        ),
    )
    return FailedCall(resp.status_code, time.perf_counter() - started, resp.body[:300])


def _fail_many(proxy: ProxyClient, key: str, model: str, override: RouterSettingsOverride) -> tuple[FailedCall, ...]:
    with ThreadPoolExecutor(max_workers=MEMORY_CONCURRENCY) as pool:
        futures: Final = tuple(
            pool.submit(_fail_once, proxy, key, model, override) for _ in range(MEMORY_REQUESTS_PER_PHASE)
        )
        return tuple(future.result() for future in futures)


def _read_rss_everywhere_after_pause(proxy: ProxyClient) -> tuple[RssReading, ...]:
    time.sleep(MEMORY_RSS_SAMPLE_INTERVAL_SECONDS)
    capture: Final = read_rss_everywhere(proxy)
    assert not capture.failures, (
        f"{len(capture.failures)} replica(s) gave no RSS reading mid-checkpoint, so their workers cannot be "
        f"compared with themselves: {'; '.join(capture.failures)}"
    )
    return capture.readings


def _readings_until_no_new_worker(
    proxy: ProxyClient, readings: tuple[RssReading, ...], samples: int, samples_since_new_worker: int
) -> tuple[RssReading, ...]:
    if samples >= RSS_SAMPLE_CAP or samples_since_new_worker >= MEMORY_RSS_SETTLE_SAMPLES:
        return readings
    sample: Final = _read_rss_everywhere_after_pause(proxy)
    known: Final = frozenset(reading.worker for reading in readings)
    new_worker_answered: Final = any(reading.worker not in known for reading in sample)
    return _readings_until_no_new_worker(
        proxy, readings + sample, samples + 1, 0 if new_worker_answered else samples_since_new_worker + 1
    )


def _settled_rss_per_worker(proxy: ProxyClient) -> Mapping[WorkerKey, RssReading]:
    readings: Final = _readings_until_no_new_worker(proxy, (), 0, 0)
    assert readings, "no /debug/memory/summary read carried ram_usage_mb, so the proxy cannot report its RSS"
    return MappingProxyType(
        {
            worker: min((reading for reading in readings if reading.worker == worker), key=lambda r: r.ram_usage_mb)
            for worker in {reading.worker for reading in readings}
        }
    )


def _record_batch_growth(
    record_property: Callable[[str, object], None],
    batch: int,
    before: Mapping[WorkerKey, RssReading],
    after: Mapping[WorkerKey, RssReading],
) -> None:
    for growth in (WorkerGrowth(before[worker], after[worker]) for worker in before.keys() & after.keys()):
        record_property(
            "rss_batch_growth",
            json.dumps(
                {
                    "batch": batch,
                    "worker": growth.warm.where,
                    "before_mb": growth.warm.ram_usage_mb,
                    "after_mb": growth.after.ram_usage_mb,
                    "growth_mb": round(growth.growth_mb, 2),
                }
            ),
        )


def _settled_rss_after_batches(
    proxy: ProxyClient,
    key: str,
    model: str,
    override: RouterSettingsOverride,
    fallback: str,
    record_property: Callable[[str, object], None],
    before: Mapping[WorkerKey, RssReading],
    batch: int,
    last_batch: int,
) -> Mapping[WorkerKey, RssReading]:
    calls: Final = _fail_many(proxy, key, model, override)
    _assert_every_call_failed_through_fallback(calls, fallback)
    after: Final = _settled_rss_per_worker(proxy)
    _record_batch_growth(record_property, batch, before, after)
    if batch == last_batch:
        return after
    return _settled_rss_after_batches(
        proxy, key, model, override, fallback, record_property, after, batch + 1, last_batch
    )


def _heaviest_worker_growth(
    warm: Mapping[WorkerKey, RssReading], after: Mapping[WorkerKey, RssReading]
) -> WorkerGrowth:
    assert warm.keys() == after.keys(), (
        f"the workers answering /debug/memory/summary changed around measured batches {FIRST_MEASURED_BATCH}-"
        f"{LAST_BATCH}, so not every "
        f"worker can be compared with itself: gone after it {sorted(warm.keys() - after.keys())} (a worker "
        f"that died or was restarted under failing traffic, which is what an OOM kill looks like), first seen "
        f"after it {sorted(after.keys() - warm.keys())} (the warm window never reached them, so they have no "
        f"baseline; raise E2E_MEMORY_RSS_SETTLE_SAMPLES if the stack has more workers than the window covers)"
    )
    return max((WorkerGrowth(warm[worker], after[worker]) for worker in warm), key=lambda growth: growth.growth_mb)


def _assert_every_call_failed_through_fallback(calls: Sequence[FailedCall], fallback: str) -> None:
    served: Final = tuple(call for call in calls if call.status_code == 200)
    assert not served, (
        f"{len(served)} of {len(calls)} calls came back 200, so they reached a provider and never "
        f"exercised the retry loop: {served[0].body_head}"
    )
    without_fallback: Final = tuple(call for call in calls if fallback not in call.body_head)
    assert not without_fallback, (
        f"{len(without_fallback)} of {len(calls)} failures never named the fallback group {fallback}, "
        f"so the request did not run through retries into the fallback: {without_fallback[0].body_head}"
    )


class TestReliabilityMemory:
    @pytest.mark.covers("reliability.perf.idle_memory.under_slo")
    @meta(
        Subject(
            domain=Domain.DEPLOY_OPS,
        )
    )
    def test_workers_idle_under_rss_budget_before_traffic(self, idle_rss: RssCapture) -> None:
        assert not idle_rss.failures, (
            f"{len(idle_rss.failures)} replica(s) gave no RSS reading when the session started, so their idle "
            f"footprint went unmeasured: {'; '.join(idle_rss.failures)}"
        )
        unmeasured: Final = frozenset(PROXY_REPLICA_URLS) - frozenset(reading.replica for reading in idle_rss.readings)
        assert not unmeasured, (
            f"{len(unmeasured)} of {len(PROXY_REPLICA_URLS)} replica(s) gave neither an RSS reading nor a failure "
            f"reason when the session started, so their idle footprint went unmeasured: {', '.join(sorted(unmeasured))}"
        )
        heaviest: Final = idle_rss.heaviest
        assert heaviest is not None, "no replica was configured to read, so nothing was measured"
        assert heaviest.ram_usage_mb <= MEMORY_IDLE_RSS_BUDGET_MB, (
            f"{heaviest.where} sat at {heaviest.ram_usage_mb:.0f} MB RSS when the session started, before it sent "
            f"any traffic, past the {MEMORY_IDLE_RSS_BUDGET_MB:.0f} MB idle budget; a DB-backed v1.100.x worker idled "
            f"at 886 MB where v1.101.0rc1 idled at 544 MB, and at that size the release stack's 2 GiB pod limit "
            f"leaves the worker little room for traffic"
        )

    @pytest.mark.filterwarnings("ignore:record_property is incompatible with junit_family")
    @pytest.mark.covers("reliability.perf.memory.under_slo")
    @meta(
        Subject(
            domain=Domain.ROUTING,
            providers=(Provider.OPENAI,),
            models=(REFUSING_MODEL,),
            mode=Mode.NONSTREAM,
        )
    )
    def test_failing_requests_do_not_grow_rss(
        self,
        proxy: ProxyClient,
        resources: ResourceManager,
        scoped_key: str,
        record_property: Callable[[str, object], None],
    ) -> None:
        marker: Final = unique_marker()
        primary: Final = f"reliability-memory-{marker}"
        fallback: Final = f"reliability-memory-fb-{marker}"
        _register_refusing_group(proxy, resources, primary)
        _register_refusing_group(proxy, resources, fallback)
        override: Final = RouterSettingsOverride(
            num_retries=MEMORY_RETRIES_PER_REQUEST, fallbacks=[{primary: [fallback]}]
        )

        warm: Final = _settled_rss_after_batches(
            proxy,
            scoped_key,
            primary,
            override,
            fallback,
            record_property,
            _settled_rss_per_worker(proxy),
            1,
            WARMUP_BATCHES,
        )
        after: Final = _settled_rss_after_batches(
            proxy,
            scoped_key,
            primary,
            override,
            fallback,
            record_property,
            warm,
            FIRST_MEASURED_BATCH,
            LAST_BATCH,
        )

        heaviest: Final = _heaviest_worker_growth(warm, after)
        assert heaviest.growth_mb <= MEMORY_RSS_CEILING_MB, (
            f"proxy RSS grew {heaviest.growth_mb:.1f} MB over measured batches {FIRST_MEASURED_BATCH}-{LAST_BATCH} "
            f"({MEASURED_BATCHES * MEMORY_REQUESTS_PER_PHASE} failing requests) after {WARMUP_BATCHES} identical "
            f"warmup batches, past the {MEMORY_RSS_CEILING_MB:.0f} MB ceiling: worker pid {heaviest.warm.worker_pid} on "
            f"{heaviest.warm.hostname or 'an unnamed host'} behind {heaviest.warm.replica} settled at "
            f"{heaviest.warm.ram_usage_mb:.1f} MB warm and {heaviest.after.ram_usage_mb:.1f} MB after; failing "
            f"requests are leaking memory the way the v1.100.0 retry breadcrumbs did"
        )