diff --git a/tests/e2e/batches/batch_client.py b/tests/e2e/batches/batch_client.py new file mode 100644 index 00000000000..d52ecb32040 --- /dev/null +++ b/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()) diff --git a/tests/e2e/batches/conftest.py b/tests/e2e/batches/conftest.py new file mode 100644 index 00000000000..adaae0764c3 --- /dev/null +++ b/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() diff --git a/tests/e2e/batches/memory.py b/tests/e2e/batches/memory.py new file mode 100644 index 00000000000..1ea2fd1695a --- /dev/null +++ b/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) diff --git a/tests/e2e/batches/poll_cap.py b/tests/e2e/batches/poll_cap.py new file mode 100644 index 00000000000..9915a1820c6 --- /dev/null +++ b/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 ...". 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 ``-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() diff --git a/tests/e2e/batches/test_batch_file_upload_e2e.py b/tests/e2e/batches/test_batch_file_upload_e2e.py new file mode 100644 index 00000000000..8892c63ae97 --- /dev/null +++ b/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 diff --git a/tests/e2e/batches/test_batch_helpers.py b/tests/e2e/batches/test_batch_helpers.py new file mode 100644 index 00000000000..3f5856ec23e --- /dev/null +++ b/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) diff --git a/tests/e2e/batches/test_batch_upload_memory_e2e.py b/tests/e2e/batches/test_batch_upload_memory_e2e.py new file mode 100644 index 00000000000..8cb17c59444 --- /dev/null +++ b/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)" + ) diff --git a/tests/e2e/batches/test_managed_poll_cap_e2e.py b/tests/e2e/batches/test_managed_poll_cap_e2e.py new file mode 100644 index 00000000000..37c17fa4767 --- /dev/null +++ b/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() diff --git a/tests/e2e/batches/test_memory_sampler.py b/tests/e2e/batches/test_memory_sampler.py new file mode 100644 index 00000000000..e62a3951c57 --- /dev/null +++ b/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) diff --git a/tests/e2e/batches/test_poll_cap_helpers.py b/tests/e2e/batches/test_poll_cap_helpers.py new file mode 100644 index 00000000000..4bda9b6ccbe --- /dev/null +++ b/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}) diff --git a/tests/e2e/docker-compose.yml b/tests/e2e/docker-compose.yml new file mode 100644 index 00000000000..f33ec2d3b17 --- /dev/null +++ b/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 diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index 7458f316852..6fdfa02bab7 100644 --- a/tests/e2e/e2e_http.py +++ b/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) diff --git a/tests/e2e/gateway/litellm-config.yml b/tests/e2e/gateway/litellm-config.yml index f4ca48cfee0..fa37d6e9b02 100644 --- a/tests/e2e/gateway/litellm-config.yml +++ b/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 diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py index 37412fc0cf5..c9324bb5167 100644 --- a/tests/e2e/transport.py +++ b/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,