test(e2e): batch/file API e2e suite

Adds a black-box e2e suite under tests/e2e/batches/ for the batches and files
API. Nothing is imported from the litellm codebase; the tests drive the live
proxy over HTTP and verify state through the generated prisma client, so they
catch real regressions rather than re-asserting internal calls.

It covers the gemini managed-files upload, verified by reading the file back
ACTIVE from the provider; a managed-object poll-cap guard for the #23472 OOM
that seeds more than one page of rows into real Postgres and watches the poll
cycle in the proxy logs; and a vertex streaming-upload memory guard for the
LIT-3382 OOM (gated, for an environment with real memory headroom). The pure
helpers have unit coverage, and the memory sampler raises rather than passing
vacuously when the cgroup read is unavailable

tests/e2e/batches/batch_client.py

@@ -0,0 +1,186 @@
+"""Client for managed batch-file e2e tests: upload an OpenAI-format batch JSONL
+through the proxy and verify the proxy actually persisted a transformed file at the
+provider.
+
+The proxy returns a self-describing unified file id (urlsafe-base64) that packs the
+target model and the provider's file URI. The tests decode it and read the file
+back from the provider, so the guard asserts the real side effect (an ACTIVE,
+non-empty file at the provider) rather than just trusting the echoed id - a broken
+upload or transform returns a bogus id or fails to persist, and the read-back
+fails. Request/response models used only by this suite live here (composition over
+the shared Gateway, DI'd in).
+"""
+
+from __future__ import annotations
+
+import base64
+import json
+import os
+from dataclasses import dataclass
+
+from pydantic import BaseModel
+
+import e2e_http
+from e2e_gateway import Gateway, build_gateway
+from e2e_http import URL, MultipartFile, NoBody, unwrap
+from models import KeyGenerateBody
+
+GEMINI_FILES_PREFIX = "https://generativelanguage.googleapis.com/v1beta/files/"
+
+
+def gemini_api_key() -> str | None:
+    """The provider key the read-back needs; tests skip when it is absent."""
+    return os.getenv("GEMINI_API_KEY")
+
+
+class FileUploadForm(BaseModel):
+    """The text fields of the /v1/files multipart upload. `target_model_names`
+    (comma-separated) routes the file through the managed-files DB path and its
+    OpenAI->provider JSONL transform."""
+
+    purpose: str = "batch"
+    target_model_names: str
+
+
+class FileObject(BaseModel):
+    id: str
+    bytes: int | None = None
+    status: str | None = None
+    purpose: str | None = None
+    object: str | None = None
+    filename: str | None = None
+
+
+class FileDeleteResponse(BaseModel):
+    id: str | None = None
+    deleted: bool | None = None
+
+
+class ProviderKeyParam(BaseModel):
+    key: str
+
+
+class ProviderFile(BaseModel):
+    """The provider-side (Google AI Studio Files API) view of the uploaded file.
+    `sizeBytes` comes back as a string, hence the `size_bytes` accessor."""
+
+    name: str | None = None
+    state: str | None = None
+    sizeBytes: str | None = None
+    mimeType: str | None = None
+
+    @property
+    def size_bytes(self) -> int:
+        return int(self.sizeBytes) if self.sizeBytes is not None else 0
+
+
+@dataclass(frozen=True, slots=True)
+class UnifiedFileId:
+    """The fields the proxy packs into its base64 unified file id."""
+
+    target_models: tuple[str, ...]
+    provider_file_uri: str
+
+
+def parse_unified_file_id(file_id: str) -> UnifiedFileId:
+    """Decode the proxy's unified file id: urlsafe-base64 of a ';'-separated list of
+    'key,value' pairs, e.g.
+    'litellm_proxy:application/jsonl;...;target_model_names,gemini-2.5-flash;llm_output_file_id,https://.../files/abc'.
+    Raises if the id is not a managed unified id, so a bogus/echoed id fails the
+    test instead of silently passing."""
+    padded = file_id + "=" * (-len(file_id) % 4)
+    decoded = base64.urlsafe_b64decode(padded).decode("utf-8")
+    if not decoded.startswith("litellm_proxy:"):
+        raise ValueError(f"not a managed unified file id: {decoded[:60]!r}")
+    fields = dict(part.split(",", 1) for part in decoded.split(";") if "," in part)
+    uri = fields.get("llm_output_file_id", "")
+    if not uri:
+        raise ValueError(f"unified id has no provider file uri: {decoded[:80]!r}")
+    target = fields.get("target_model_names", "")
+    return UnifiedFileId(
+        target_models=tuple(model for model in target.split(",") if model),
+        provider_file_uri=uri,
+    )
+
+
+def batch_jsonl(model: str, *, lines: int, pad_bytes: int = 0) -> bytes:
+    """An OpenAI-format batch input file targeting `model`: `lines` chat-completion
+    requests, each optionally padded so the file reaches a chosen size. custom_id is
+    unique per line, so a transform that drops or merges lines is detectable."""
+    pad = "x" * pad_bytes
+    rows = (
+        json.dumps(
+            {
+                "custom_id": f"req-{i}",
+                "method": "POST",
+                "url": "/v1/chat/completions",
+                "body": {
+                    "model": model,
+                    "messages": [{"role": "user", "content": f"{pad}{i}"}],
+                    "max_tokens": 4,
+                },
+            }
+        )
+        for i in range(lines)
+    )
+    return ("\n".join(rows) + "\n").encode("utf-8")
+
+
+@dataclass(frozen=True, slots=True)
+class BatchFilesClient:
+    gateway: Gateway
+
+    # ---- generic key ops (satisfy lifecycle.ResourceClient via the Gateway) ----
+
+    def generate_key(self, *, models: list[str] | None = None) -> str:
+        return self.gateway.generate_key(KeyGenerateBody(models=models or []))
+
+    def delete_key(self, key: str) -> None:
+        self.gateway.delete_key(key)
+
+    def delete_customers(self, user_ids: list[str]) -> None:
+        self.gateway.delete_customers(user_ids)
+
+    # ---- managed batch files -------------------------------------------
+
+    def upload_batch_file(
+        self, key: str, content: bytes, *, target_model: str
+    ) -> FileObject:
+        return unwrap(
+            self.gateway.transport.upload(
+                "/v1/files",
+                headers=self.gateway.transport.bearer(key),
+                form=FileUploadForm(target_model_names=target_model),
+                file=MultipartFile(
+                    filename="batch.jsonl",
+                    content=content,
+                    content_type="application/jsonl",
+                ),
+                response_type=FileObject,
+            )
+        )
+
+    def delete_file(self, file_id: str) -> None:
+        _ = self.gateway.transport.delete(
+            f"/v1/files/{file_id}",
+            headers=self.gateway.transport.master,
+            json=NoBody(),
+            response_type=FileDeleteResponse,
+        )
+
+    def provider_file(self, uri: str, *, api_key: str) -> ProviderFile:
+        """Read the uploaded file back from the provider's Files API. Goes straight
+        to the provider URL (not the proxy base), so it proves the proxy really
+        created the file rather than echoing an id."""
+        return unwrap(
+            e2e_http.get(
+                URL(uri),
+                headers=NoBody(),
+                params=ProviderKeyParam(key=api_key),
+                response_type=ProviderFile,
+            )
+        )
+
+
+def build_client() -> BatchFilesClient:
+    return BatchFilesClient(gateway=build_gateway())

tests/e2e/batches/conftest.py

@@ -0,0 +1,10 @@
+"""Batch-files suite's `client` fixture."""
+
+import pytest
+
+from batch_client import BatchFilesClient, build_client
+
+
+@pytest.fixture(scope="session")
+def client() -> BatchFilesClient:
+    return build_client()

tests/e2e/batches/memory.py

@@ -0,0 +1,92 @@
+"""Peak-memory sampling for the batch-upload memory guard.
+
+The OOM regression shows up as the proxy's resident memory growing with the
+uploaded file size. To catch it, sample the proxy's memory while an upload runs
+and keep the peak. The sampler is a seam (Protocol) so the environment decides how
+memory is read: the local docker-compose stack reads the container's cgroup, an
+EKS run would read the gateway pod (e.g. via `kubectl exec`/metrics) with the same
+interface.
+"""
+
+from __future__ import annotations
+
+import subprocess
+import threading
+import time
+from dataclasses import dataclass
+from typing import Callable, Protocol, TypeVar
+
+T = TypeVar("T")
+
+
+@dataclass(frozen=True, slots=True)
+class PeakMemory:
+    baseline_bytes: int
+    peak_bytes: int
+
+    @property
+    def growth_bytes(self) -> int:
+        return max(0, self.peak_bytes - self.baseline_bytes)
+
+
+class MemorySampler(Protocol):
+    def measure(self, during: Callable[[], T]) -> tuple[T, PeakMemory]:
+        """Run `during`, sampling memory throughout; return its result and the peak
+        memory observed against the pre-run baseline."""
+        ...
+
+
+@dataclass(frozen=True, slots=True)
+class DockerCgroupSampler:
+    """Reads the litellm container's anonymous (RSS) memory from its cgroup
+    (`anon` in `/sys/fs/cgroup/memory.stat`, cgroup v2) on a background thread.
+    Anonymous memory is what a buffered-in-memory copy of the upload shows up as;
+    `memory.current` is avoided because it also counts reclaimable page cache, which
+    a streamed-to-disk upload fills without ever risking an OOM."""
+
+    container: str
+    interval_seconds: float = 0.1
+
+    def _read_bytes(self) -> int:
+        try:
+            result = subprocess.run(
+                ["docker", "exec", self.container, "cat", "/sys/fs/cgroup/memory.stat"],
+                capture_output=True,
+                text=True,
+                timeout=30,
+            )
+        except (OSError, subprocess.SubprocessError):
+            return -1
+        for line in result.stdout.splitlines():
+            field, _, value = line.partition(" ")
+            if field == "anon" and value.strip().isdigit():
+                return int(value.strip())
+        return -1
+
+    def measure(self, during: Callable[[], T]) -> tuple[T, PeakMemory]:
+        baseline = self._read_bytes()
+        if baseline < 0:
+            raise RuntimeError(
+                f"could not read anon memory from container {self.container!r}; the "
+                "memory guard cannot run, so refusing to pass vacuously (needs cgroup "
+                "v2 and docker exec access)"
+            )
+        peak = baseline
+        stop = threading.Event()
+
+        def sample() -> None:
+            nonlocal peak
+            while not stop.is_set():
+                current = self._read_bytes()
+                if current > peak:
+                    peak = current
+                time.sleep(self.interval_seconds)
+
+        sampler_thread = threading.Thread(target=sample)
+        sampler_thread.start()
+        try:
+            result = during()
+        finally:
+            stop.set()
+            sampler_thread.join()
+        return result, PeakMemory(baseline_bytes=baseline, peak_bytes=peak)

tests/e2e/batches/poll_cap.py

@@ -0,0 +1,133 @@
+"""Helpers for the managed-object poll-cap regression (#23472), black-box.
+
+The bug: CheckBatchCost paged its managed-object query unbounded, so each poll
+cycle pulled the whole table into pod memory and OOM'd. The fix caps the query at
+MAX_OBJECTS_PER_POLL_CYCLE rows per cycle. There is no API that exposes the query,
+but the cap is observable in the proxy's logs: rows the cycle selects but can't
+decode are logged as "Skipping job <unified_object_id> ...". So seed more than one
+page of selectable rows and watch which the cycle touches - a capped cycle touches
+exactly the oldest page, an unbounded one touches them all.
+
+No imports from the litellm codebase: rows are seeded directly in Postgres via the
+generated prisma client, and the proxy's logs are read via a swappable command
+(docker locally, a pod-log command on EKS).
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+import subprocess
+from dataclasses import dataclass
+from datetime import datetime, timedelta, timezone
+
+from prisma import Prisma
+
+PROXY_DB_URL = os.getenv(
+    "E2E_MANAGED_DB_URL",
+    "postgresql://llmproxy:dbpassword9090@localhost:5432/litellm",
+)
+CONTAINER = os.getenv("LITELLM_CONTAINER", "e2e-litellm-1")
+
+SEED_PREFIX = "pollcap"
+
+
+def proxy_poll_cap() -> int:
+    """The proxy's per-cycle row cap (MAX_OBJECTS_PER_POLL_CYCLE). Read from the
+    proxy container's env so the test matches the running proxy; defaults to the
+    proxy's own default of 50 when unset. Override with E2E_POLL_CAP."""
+    override = os.getenv("E2E_POLL_CAP")
+    if override is not None:
+        return int(override)
+    try:
+        result = subprocess.run(
+            ["docker", "exec", CONTAINER, "printenv", "MAX_OBJECTS_PER_POLL_CYCLE"],
+            capture_output=True,
+            text=True,
+            timeout=10,
+        )
+    except (OSError, subprocess.SubprocessError):
+        return 50
+    value = result.stdout.strip()
+    return int(value) if value.isdigit() else 50
+
+
+def parse_skipped_indices(log_text: str, prefix: str) -> frozenset[int]:
+    """The seeded-row indices the proxy logged as skipped. Each seeded row's
+    unified_object_id is ``<prefix>-NNN``; the poll logs it when it selects but
+    cannot decode the row."""
+    pattern = re.compile(re.escape(prefix) + r"-(\d{3})")
+    return frozenset(int(match) for match in pattern.findall(log_text))
+
+
+@dataclass(frozen=True, slots=True)
+class ProxyLog:
+    container: str
+
+    def skipped_indices(self, prefix: str, *, since_seconds: int) -> frozenset[int]:
+        result = subprocess.run(
+            ["docker", "logs", "--since", f"{since_seconds}s", self.container],
+            capture_output=True,
+            text=True,
+            timeout=20,
+        )
+        return parse_skipped_indices(result.stdout + result.stderr, prefix)
+
+
+@dataclass(frozen=True, slots=True)
+class ManagedObjectSeeder:
+    """Seeds invalid (undecodable) batch managed-objects so the poll selects them,
+    logs a skip, and leaves them in place to be re-selected next cycle - never
+    processed, never deleted by the proxy. created_at is recent (well inside the
+    staleness cutoff) and strictly increasing, so the oldest `count` rows are a
+    deterministic page."""
+
+    db_url: str
+    prefix: str
+
+    async def _client(self) -> Prisma:
+        client = Prisma(datasource={"url": self.db_url})
+        await client.connect()
+        return client
+
+    async def reset(self) -> None:
+        client = await self._client()
+        try:
+            await client.litellm_managedobjecttable.delete_many(
+                where={"unified_object_id": {"startswith": f"{SEED_PREFIX}-"}}
+            )
+        finally:
+            await client.disconnect()
+
+    async def seed(self, count: int) -> None:
+        client = await self._client()
+        try:
+            base = datetime.now(timezone.utc) - timedelta(seconds=count + 5)
+            for i in range(count):
+                unified_object_id = f"{self.prefix}-{i:03d}"
+                await client.litellm_managedobjecttable.create(
+                    data={
+                        "unified_object_id": unified_object_id,
+                        "model_object_id": f"{self.prefix}-mob-{i}",
+                        "file_object": json.dumps(
+                            {"id": unified_object_id, "object": "batch"}
+                        ),
+                        "file_purpose": "batch",
+                        "status": "validating",
+                        "batch_processed": False,
+                        "team_id": self.prefix,
+                        "created_at": base + timedelta(seconds=i),
+                    }
+                )
+        finally:
+            await client.disconnect()
+
+    async def delete(self) -> None:
+        client = await self._client()
+        try:
+            await client.litellm_managedobjecttable.delete_many(
+                where={"unified_object_id": {"startswith": f"{self.prefix}-"}}
+            )
+        finally:
+            await client.disconnect()

tests/e2e/batches/test_batch_file_upload_e2e.py

@@ -0,0 +1,58 @@
+"""Managed batch-file upload guard for the LIT-3382 batch-processing failures.
+
+The customer-facing symptom of LIT-3382 (and the OOM fixed in #31036) was that
+uploading a batch JSONL failed outright, blocking batch workflows. This guard
+drives the managed-files upload that works locally (the gemini AI Studio target),
+end to end: it uploads a real multi-line OpenAI-format batch JSONL through the
+proxy, decodes the unified file id the proxy returns, then reads the file back
+from the provider's Files API and asserts it is ACTIVE and non-empty. A broken
+upload or transform either 500s, returns an id with no backing provider file, or
+persists an empty file - each fails an assertion here.
+
+The faithful *memory*/OOM regression (the streaming transform #31036 actually
+fixed) lives on the vertex_ai handler, which needs GCS + billing + memory
+headroom; that guard is test_batch_upload_memory_e2e, enabled on EKS.
+
+Skips when no proxy answers (shared conftest) or when GEMINI_API_KEY is absent
+(the provider read-back needs it).
+"""
+
+import pytest
+
+from batch_client import (
+    GEMINI_FILES_PREFIX,
+    BatchFilesClient,
+    batch_jsonl,
+    gemini_api_key,
+    parse_unified_file_id,
+)
+from lifecycle import ResourceManager
+
+pytestmark = pytest.mark.e2e
+
+TARGET_MODEL = "gemini-2.5-flash"
+
+
+def test_managed_batch_upload_persists_active_file_at_provider(
+    client: BatchFilesClient, resources: ResourceManager, scoped_key: str
+) -> None:
+    api_key = gemini_api_key()
+    if api_key is None:
+        pytest.skip("GEMINI_API_KEY not set; cannot read the provider file back")
+
+    content = batch_jsonl(TARGET_MODEL, lines=5)
+    uploaded = client.upload_batch_file(scoped_key, content, target_model=TARGET_MODEL)
+    resources.defer(lambda: client.delete_file(uploaded.id))
+
+    assert uploaded.status == "uploaded"
+    assert uploaded.purpose == "batch"
+    assert uploaded.bytes is not None and uploaded.bytes > 0
+
+    parsed = parse_unified_file_id(uploaded.id)
+    assert TARGET_MODEL in parsed.target_models
+    assert parsed.provider_file_uri.startswith(GEMINI_FILES_PREFIX)
+
+    provider = client.provider_file(parsed.provider_file_uri, api_key=api_key)
+    assert provider.state == "ACTIVE"
+    assert provider.mimeType == "application/jsonl"
+    assert provider.size_bytes > 0

tests/e2e/batches/test_batch_helpers.py

@@ -0,0 +1,70 @@
+"""Unit coverage for the batch-files helpers (no proxy needed, always runs).
+
+These pin the pure logic the e2e guard relies on: decoding the proxy's unified
+file id, and building a well-formed OpenAI batch JSONL. If decoding silently
+accepted a bogus id, the e2e read-back assertion would be meaningless - so the
+reject cases matter as much as the happy path.
+"""
+
+from __future__ import annotations
+
+import base64
+import json
+
+import pytest
+
+from batch_client import batch_jsonl, parse_unified_file_id
+
+
+def _encode(decoded: str) -> str:
+    return base64.urlsafe_b64encode(decoded.encode("utf-8")).decode("utf-8").rstrip("=")
+
+
+def test_parse_unified_file_id_extracts_target_and_uri() -> None:
+    uri = "https://generativelanguage.googleapis.com/v1beta/files/abc123"
+    file_id = _encode(
+        "litellm_proxy:application/jsonl;unified_id,xyz;"
+        f"target_model_names,gemini-2.5-flash;llm_output_file_id,{uri}"
+    )
+    parsed = parse_unified_file_id(file_id)
+    assert parsed.target_models == ("gemini-2.5-flash",)
+    assert parsed.provider_file_uri == uri
+
+
+def test_parse_unified_file_id_handles_multiple_targets() -> None:
+    file_id = _encode(
+        "litellm_proxy:application/jsonl;"
+        "target_model_names,model-a,model-b;"
+        "llm_output_file_id,https://example.com/files/f1"
+    )
+    assert parse_unified_file_id(file_id).target_models == ("model-a", "model-b")
+
+
+def test_parse_unified_file_id_rejects_non_managed_id() -> None:
+    with pytest.raises(ValueError):
+        parse_unified_file_id(_encode("openai-file-abc123"))
+
+
+def test_parse_unified_file_id_rejects_id_without_provider_uri() -> None:
+    with pytest.raises(ValueError):
+        parse_unified_file_id(
+            _encode(
+                "litellm_proxy:application/jsonl;target_model_names,gemini-2.5-flash"
+            )
+        )
+
+
+def test_batch_jsonl_emits_one_request_per_line_with_unique_ids() -> None:
+    raw = batch_jsonl("gemini-2.5-flash", lines=5)
+    rows = [json.loads(line) for line in raw.splitlines()]
+    assert len(rows) == 5
+    assert [r["custom_id"] for r in rows] == [f"req-{i}" for i in range(5)]
+    assert {r["body"]["model"] for r in rows} == {"gemini-2.5-flash"}
+    assert all(r["method"] == "POST" for r in rows)
+
+
+def test_batch_jsonl_padding_grows_the_file() -> None:
+    small = batch_jsonl("m", lines=10)
+    padded = batch_jsonl("m", lines=10, pad_bytes=1000)
+    assert len(padded) - len(small) >= 10 * 1000
+    assert len(batch_jsonl("m", lines=10)) == len(small)

tests/e2e/batches/test_batch_upload_memory_e2e.py

@@ -0,0 +1,81 @@
+"""Faithful memory/OOM regression guard for the batch-upload transform (#31036).
+
+LIT-3382: the OpenAI->Vertex batch JSONL upload transform buffered the whole file
+and made several in-memory copies, so a large upload OOM-killed the pod. #31036
+streams the upload so memory stays bounded regardless of file size. This guard
+uploads a large batch JSONL through the vertex_ai managed-files path while sampling
+the proxy's memory, and asserts peak growth stays within a small multiple of the
+file size: the streaming fix passes, the buffering regression (memory scaling with
+file size) blows past the bound.
+
+Gated off by default. The vertex_ai files handler stages the JSONL in GCS, which
+needs a writable bucket on a billing-enabled project, and a meaningful memory
+assertion needs real memory headroom (a ~3.8 GB local docker VM OOMs the whole
+proxy before the signal is clean). Both hold on EKS, not on the local
+docker-compose stack. To enable: configure VERTEXAI creds + a GCS bucket
+(`gcs_bucket_name`) on the vertex deployment, set E2E_BATCH_MEMORY_ENABLED=1, and
+inject the MemorySampler for the environment (the docker cgroup sampler here, or a
+pod-based one on EKS).
+
+NOTE: not yet validated against a live streaming vertex path (billing-blocked in
+the dev environment). Treat MAX_GROWTH_RATIO as a starting point and confirm the
+pass/fail margin when first enabling it.
+"""
+
+from __future__ import annotations
+
+import os
+
+import pytest
+
+from batch_client import BatchFilesClient, batch_jsonl
+from lifecycle import ResourceManager
+from memory import DockerCgroupSampler, MemorySampler
+
+pytestmark = [
+    pytest.mark.e2e,
+    pytest.mark.skipif(
+        os.getenv("E2E_BATCH_MEMORY_ENABLED") != "1",
+        reason="needs the streaming vertex_ai files path: GCS bucket + billing + "
+        "memory headroom (enable on EKS with E2E_BATCH_MEMORY_ENABLED=1)",
+    ),
+]
+
+VERTEX_MODEL = os.getenv("E2E_BATCH_MEMORY_MODEL", "gemini-2.5-flash-vertex")
+FILE_BYTES = int(os.getenv("E2E_BATCH_MEMORY_FILE_MB", "200")) * 1024 * 1024
+MAX_GROWTH_RATIO = float(os.getenv("E2E_BATCH_MEMORY_MAX_RATIO", "3.0"))
+CONTAINER = os.getenv("E2E_LITELLM_CONTAINER", "e2e-litellm-1")
+
+
+@pytest.fixture
+def sampler() -> MemorySampler:
+    return DockerCgroupSampler(container=CONTAINER)
+
+
+def _sized_batch(model: str, target_bytes: int, *, pad_bytes: int = 2000) -> bytes:
+    per_line = len(batch_jsonl(model, lines=1, pad_bytes=pad_bytes))
+    lines = max(1, target_bytes // per_line)
+    return batch_jsonl(model, lines=lines, pad_bytes=pad_bytes)
+
+
+def test_large_vertex_batch_upload_memory_is_bounded(
+    client: BatchFilesClient,
+    resources: ResourceManager,
+    scoped_key: str,
+    sampler: MemorySampler,
+) -> None:
+    content = _sized_batch(VERTEX_MODEL, FILE_BYTES)
+    file_size = len(content)
+
+    uploaded, mem = sampler.measure(
+        lambda: client.upload_batch_file(scoped_key, content, target_model=VERTEX_MODEL)
+    )
+    resources.defer(lambda: client.delete_file(uploaded.id))
+
+    assert uploaded.status == "uploaded"
+    ratio = mem.growth_bytes / file_size
+    assert ratio < MAX_GROWTH_RATIO, (
+        f"proxy memory grew {mem.growth_bytes / 1e6:.0f}MB for a "
+        f"{file_size / 1e6:.0f}MB upload (ratio {ratio:.1f}x >= {MAX_GROWTH_RATIO}x); "
+        f"the streaming transform must keep upload memory bounded (#31036 OOM regression)"
+    )

tests/e2e/batches/test_managed_poll_cap_e2e.py

@@ -0,0 +1,83 @@
+"""Black-box regression guard for the unbounded managed-object poll OOM (#23472).
+
+CheckBatchCost pages its managed-object query with take=MAX_OBJECTS_PER_POLL_CYCLE.
+Before the fix the query was unbounded, so each poll cycle pulled the entire table
+into pod memory and OOM'd. This guard seeds one page + 10 selectable rows into the
+real Postgres, lets the live proxy run a poll cycle, and asserts the cycle touched
+exactly the oldest page (the cap) and never the 10 newest rows. Drop the take (the
+regression) and the cycle selects every seeded row, so the newest rows show up in
+the logs and the assertion fails.
+
+Pure black-box: rows are seeded via the generated prisma client and the cap is read
+from the proxy's logs - no imports from the litellm codebase, no calls into its
+internals. The proxy's batch poll runs every ~15-45s here (proxy_batch_polling_interval
+is shortened in the e2e config so a cycle lands inside the test).
+
+Skips unless a proxy answers (shared conftest) and the poll actually runs within the
+deadline; once a cycle is observed, behavior is asserted. Seeded rows carry a unique
+prefix and are deleted in a finally block.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import time
+import uuid
+
+import pytest
+
+from poll_cap import (
+    CONTAINER,
+    PROXY_DB_URL,
+    SEED_PREFIX,
+    ManagedObjectSeeder,
+    ProxyLog,
+    proxy_poll_cap,
+)
+
+pytestmark = [pytest.mark.e2e, pytest.mark.asyncio]
+
+EXTRA_ROWS = 10
+DEADLINE_SECONDS = 180
+STABLE_SECONDS = 50
+
+
+async def test_managed_object_poll_is_capped_per_cycle() -> None:
+    cap = proxy_poll_cap()
+    prefix = f"{SEED_PREFIX}-{uuid.uuid4().hex[:8]}"
+    seeder = ManagedObjectSeeder(db_url=PROXY_DB_URL, prefix=prefix)
+    log = ProxyLog(container=CONTAINER)
+    oldest_page = frozenset(range(cap))
+    newest_rows = frozenset(range(cap, cap + EXTRA_ROWS))
+
+    start = time.monotonic()
+    await seeder.reset()
+    await seeder.seed(cap + EXTRA_ROWS)
+    try:
+        skipped: frozenset[int] = frozenset()
+        first_full_at: float | None = None
+        while time.monotonic() - start < DEADLINE_SECONDS:
+            since = int(time.monotonic() - start) + 5
+            skipped = log.skipped_indices(prefix, since_seconds=since)
+            if skipped & newest_rows:
+                break
+            if len(skipped) >= cap and first_full_at is None:
+                first_full_at = time.monotonic()
+            if (
+                first_full_at is not None
+                and time.monotonic() - first_full_at >= STABLE_SECONDS
+            ):
+                break
+            await asyncio.sleep(5)
+
+        if not skipped:
+            pytest.skip("managed-object batch poll did not run within the deadline")
+
+        assert skipped == oldest_page, (
+            f"poll cycle touched {len(skipped)} of {cap + EXTRA_ROWS} seeded rows "
+            f"(expected exactly the oldest {cap}); newest rows it should never reach "
+            f"= {sorted(skipped & newest_rows)}. The per-cycle cap "
+            f"(MAX_OBJECTS_PER_POLL_CYCLE={cap}) must bound the query (#23472 OOM regression)"
+        )
+    finally:
+        await seeder.delete()

tests/e2e/batches/test_memory_sampler.py

@@ -0,0 +1,19 @@
+"""The memory sampler must fail loudly, never vacuously pass.
+
+If the cgroup read fails (no cgroup v2, docker exec unavailable, a different
+memory.stat layout) _read_bytes returns -1; without a guard, baseline and peak are
+both -1, growth is 0, and the OOM assertion passes as if the proxy used no memory -
+hiding the LIT-3382 regression. measure() must raise instead.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from memory import DockerCgroupSampler
+
+
+def test_measure_raises_when_cgroup_unreadable() -> None:
+    sampler = DockerCgroupSampler(container="e2e-nonexistent-container-xyz")
+    with pytest.raises(RuntimeError):
+        sampler.measure(lambda: None)

tests/e2e/batches/test_poll_cap_helpers.py

@@ -0,0 +1,35 @@
+"""Unit coverage for the poll-cap log parser (no proxy needed, always runs).
+
+The e2e guard's verdict rests entirely on which seeded indices it reads out of the
+proxy logs, so the parser is pinned here: it must extract exactly the indices of
+its own prefix, ignore other prefixes, and not be fooled by substrings.
+"""
+
+from __future__ import annotations
+
+from poll_cap import parse_skipped_indices
+
+PREFIX = "pollcap-abc123"
+
+
+def test_parses_indices_for_its_own_prefix() -> None:
+    text = "\n".join(
+        f'{{"message": "Skipping job {PREFIX}-{i:03d} because it is not a valid '
+        'unified object id"}'
+        for i in (0, 7, 49)
+    )
+    assert parse_skipped_indices(text, PREFIX) == frozenset({0, 7, 49})
+
+
+def test_ignores_other_prefixes() -> None:
+    text = f"Skipping job pollcap-other99-005 ...\nSkipping job {PREFIX}-012 ..."
+    assert parse_skipped_indices(text, PREFIX) == frozenset({12})
+
+
+def test_empty_when_prefix_absent() -> None:
+    assert parse_skipped_indices("nothing here", PREFIX) == frozenset()
+
+
+def test_deduplicates_repeated_cycles() -> None:
+    text = f"{PREFIX}-003 ...\n{PREFIX}-003 ...\n{PREFIX}-004 ..."
+    assert parse_skipped_indices(text, PREFIX) == frozenset({3, 4})

tests/e2e/docker-compose.yml

@@ -0,0 +1,64 @@
+# Local e2e proxy, built from the litellm_internal_staging branch.
+#
+# The image `litellm-staging:local` is built from a clean checkout of
+# litellm_internal_staging (the working tree's uv.lock has merge markers, so
+# build from a worktree, not the dirty tree):
+#
+#   git worktree add /tmp/litellm-staging litellm_internal_staging
+#   docker build -t litellm-staging:local /tmp/litellm-staging
+#
+# Then provide provider keys (copy .env.example -> .env, fill in keys) and run:
+#
+#   docker compose -f tests/e2e/docker-compose.yml up -d
+#
+# The proxy comes up on http://localhost:4000 with master key sk-1234, which is
+# what the e2e suites default to. Point the suites at it:
+#
+#   LITELLM_MASTER_KEY=sk-1234 LITELLM_PROXY_URL=http://localhost:4000 \
+#     .venv/bin/python -m pytest tests/e2e/ -v
+
+services:
+  litellm:
+    image: litellm-staging:local
+    ports:
+      - "4000:4000"
+    command: ["--config", "/app/config.yaml", "--port", "4000"]
+    volumes:
+      - ./gateway/litellm-config.yml:/app/config.yaml
+    environment:
+      DATABASE_URL: "postgresql://llmproxy:dbpassword9090@db:5432/litellm"
+      LITELLM_MASTER_KEY: "sk-1234"
+      STORE_MODEL_IN_DB: "True"
+      # Run the budget-reset job every ~15-20s (default ~10min) so short budget
+      # windows actually reset within an e2e test (multi-window + reset suites).
+      PROXY_BUDGET_RESCHEDULER_MIN_TIME: "15"
+      PROXY_BUDGET_RESCHEDULER_MAX_TIME: "20"
+      # LITELLM_LICENSE is injected from .env (gitignored), never hardcoded here.
+    env_file:
+      - .env
+    depends_on:
+      - db
+      - redis
+
+  db:
+    image: postgres:17
+    environment:
+      POSTGRES_DB: litellm
+      POSTGRES_USER: llmproxy
+      POSTGRES_PASSWORD: dbpassword9090
+    ports:
+      - "5432:5432"
+    volumes:
+      - postgres_data:/var/lib/postgresql/data
+
+  redis:
+    image: redis:7-alpine
+    # Exposed on host 6380 (not 6379, to avoid clashing with a host-local redis) so
+    # test_spend_counter_reseed_e2e can read the shared spend counter back and assert
+    # it equals the DB spend. The litellm container still reaches it as redis:6379.
+    ports:
+      - "6380:6379"
+
+volumes:
+  postgres_data:
+    driver: local

tests/e2e/e2e_http.py

@@ -10,6 +10,7 @@ requests itself imports.
 
 from __future__ import annotations
 
+from dataclasses import dataclass
 from typing import Generic, Iterator, Literal, NewType, TypeVar, cast
 
 import pytest
@@ -19,6 +20,18 @@ from pydantic import BaseModel, ConfigDict, Field
 URL = NewType("URL", str)
 
 
+@dataclass(frozen=True, slots=True)
+class MultipartFile:
+    """The file part of a multipart/form-data upload. Raw bytes (not a pydantic
+    model) since the content is passed straight to the encoder, never serialized to
+    JSON - and a batch JSONL can be large enough that a copy would matter."""
+
+    filename: str
+    content: bytes
+    content_type: str
+    field_name: str = "file"
+
+
 class Headers(BaseModel):
     """Base for header models. Subclasses may alias to hyphenated header names
     (e.g. ``x-litellm-api-key``); serialization uses by_alias."""
@@ -304,3 +317,28 @@ def stream(
     """Streaming (SSE) call: consumes the stream counting events, and captures the
     x-litellm-call-id + content-type headers. Body is elided."""
     return send(url, headers=headers, json=json, stream=True, timeout=timeout)
+
+
+def upload[R: BaseModel](
+    url: URL,
+    *,
+    headers: BaseModel,
+    form: BaseModel,
+    file: MultipartFile,
+    response_type: type[R],
+    timeout: float = 120.0,
+) -> Result[R]:
+    """multipart/form-data POST (e.g. /v1/files): the text fields come from `form`,
+    the file part from `file`. requests sets the multipart Content-Type + boundary,
+    so `headers` must carry only auth (no Content-Type)."""
+    try:
+        resp = requests.post(
+            str(url),
+            headers=_headers(headers),
+            data=form.model_dump(by_alias=True, exclude_none=True),
+            files={file.field_name: (file.filename, file.content, file.content_type)},
+            timeout=timeout,
+        )
+    except requests.RequestException as exc:
+        return NetworkError(message=str(exc))
+    return _classify(resp, response_type)

tests/e2e/gateway/litellm-config.yml

@@ -27,6 +27,10 @@ general_settings:
   alerts: ["email"]
   proxy_budget_rescheduler_min_time: 15
   proxy_budget_rescheduler_max_time: 20
+  # Run the managed-object batch-cost poll every ~15-45s (default 1h) so the poll's
+  # per-cycle row cap (MAX_OBJECTS_PER_POLL_CYCLE) is observable within a test -
+  # exercised by tests/e2e/batches/test_managed_poll_cap_e2e.py.
+  proxy_batch_polling_interval: 15
 
 # fallbacks: [{"gpt-4": ["anthropic.claude-3-5-sonnet-20240620-v1:0"]}] #Configure fallbacks for context window exeeded errors (In this example, we will fall back to Claude Sonnet if over 8000 tokens, which is gpt-4's limit)
   # default_fallbacks: ["anthropic.claude-3-haiku-20240307-v1:0"] #Configure fallbacks for any error for every model (the above fallback configurations override this one)
@@ -127,9 +131,9 @@ model_list:
       model: openai/text-embedding-3-small
       api_key: os.environ/OPENAI_API_KEY
 
-  - model_name: gemini-2-embedding
+  - model_name: gemini-embedding-2
     litellm_params:
-      model: gemini/gemini-2-embedding
+      model: gemini/gemini-embedding-2
       api_key: os.environ/GEMINI_API_KEY
 
   # realtime models

tests/e2e/transport.py

@@ -13,7 +13,14 @@ from typing import Protocol
 from pydantic import BaseModel
 
 import e2e_http
-from e2e_http import URL, AuthHeaders, ProbeResult, Result, StreamingResponse
+from e2e_http import (
+    URL,
+    AuthHeaders,
+    MultipartFile,
+    ProbeResult,
+    Result,
+    StreamingResponse,
+)
 
 
 class Transport(Protocol):
@@ -21,6 +28,16 @@ class Transport(Protocol):
         self, path: str, *, headers: BaseModel, json: BaseModel, response_type: type[R]
     ) -> Result[R]: ...
 
+    def upload[R: BaseModel](
+        self,
+        path: str,
+        *,
+        headers: BaseModel,
+        form: BaseModel,
+        file: MultipartFile,
+        response_type: type[R],
+    ) -> Result[R]: ...
+
     def stream(
         self, path: str, *, headers: BaseModel, json: BaseModel
     ) -> StreamingResponse: ...
@@ -83,6 +100,24 @@ class HttpTransport:
             timeout=self.request_timeout,
         )
 
+    def upload[R: BaseModel](
+        self,
+        path: str,
+        *,
+        headers: BaseModel,
+        form: BaseModel,
+        file: MultipartFile,
+        response_type: type[R],
+    ) -> Result[R]:
+        return e2e_http.upload(
+            self._url(path),
+            headers=headers,
+            form=form,
+            file=file,
+            response_type=response_type,
+            timeout=self.request_timeout,
+        )
+
     def get[R: BaseModel](
         self,
         path: str,
@@ -203,6 +238,19 @@ class SplitTransport:
             path, headers=headers, json=json, response_type=response_type
         )
 
+    def upload[R: BaseModel](
+        self,
+        path: str,
+        *,
+        headers: BaseModel,
+        form: BaseModel,
+        file: MultipartFile,
+        response_type: type[R],
+    ) -> Result[R]:
+        return self._route(path).upload(
+            path, headers=headers, form=form, file=file, response_type=response_type
+        )
+
     def get[R: BaseModel](
         self,
         path: str,