Merge pull request #32166 from BerriAI/litellm_e2e_batches_ocr_model_registration

fix(e2e): register batch + rust OCR deployments via /model/new

tests/e2e/batches/COVERAGE.md

@@ -9,20 +9,21 @@ cancels, and lists a batch; everything created is deleted on teardown.
 
 Only supported cells are tested. The capability table in `capabilities.py` holds one
 row per supported (provider, scenario) pair, so there are no skipped cells in the
-parametrized run.
+parametrized run. The batches suite never skips: missing provider creds or upstream
+failures are hard test failures (see `tests/e2e/CLAUDE.md`).
 
 | Provider  | create | retrieve | cancel | list | file backing |
 |-----------|--------|----------|--------|------|--------------|
 | OpenAI    | yes | yes | yes | yes | OpenAI Files |
 | Azure     | yes | yes | yes | yes | Azure Files |
-| Vertex AI | yes | yes | yes | yes | GCS bucket |
-| Bedrock   | yes | yes | no (limited upstream) | no | S3 bucket |
-| Anthropic | no  | yes (env-gated) | no | no | Anthropic Files |
+| Vertex AI | yes | yes | yes | yes | GCS bucket (`GCS_BUCKET_NAME` via files_settings) |
+| Bedrock   | yes | yes | no (limited upstream) | no | S3 bucket (`AWS_BATCH_S3_BUCKET` + `AWS_BATCH_ROLE_ARN` on model) |
 
 Bedrock cancel is unreliable upstream and list is unsupported, so both are gated off
-(`can_cancel=False`, `can_list=False`). Anthropic cannot create/cancel/list through
-litellm, so it has a standalone retrieve test that skips unless `ANTHROPIC_BATCH_ID`
-points at a real Anthropic batch.
+(`can_cancel=False`, `can_list=False`) when that provider is enabled in the matrix.
+Bedrock file upload requires a model on the request (`encoded` / `unified` scenarios only);
+`model_param` and `provider_fallback` are omitted because `POST /bedrock/v1/files` has no
+model-less passthrough path.
 
 ## Routing scenarios (per `litellm/proxy/batches_endpoints/endpoints.py`)
 
@@ -67,9 +68,10 @@ File delete asserts `object=="file"` and `deleted==True`.
 
 | File | Covers |
 |------|--------|
-| `batch_client.py` | typed file upload/download + batch create/retrieve/cancel/list/delete over the shared Gateway; denial helpers |
-| `capabilities.py` | the provider x scenario matrix + id-shape classifiers + per-provider raw-id assertion |
-| `test_batches_e2e.py` | parametrized lifecycle with per-endpoint output assertions, file upload/delete outputs, key-model-access denial, anthropic retrieve |
+| `batch_client.py` | typed file upload/download + batch create/retrieve/cancel/list/delete over the shared Gateway; runtime batch model registration via /model/new; denial helpers |
+| `capabilities.py` | the provider x scenario matrix + per-provider /model/new params + id-shape classifiers + per-provider raw-id assertion |
+| `conftest.py` | session-scoped batch deployment registration and teardown |
+| `test_batches_e2e.py` | parametrized lifecycle with per-endpoint output assertions, file upload/delete outputs, key-model-access denial |
 
 ## Out of scope (intentionally)
 

tests/e2e/batches/batch_client.py

@@ -1,11 +1,13 @@
 """Client for the batches e2e suite: file upload/download and the batch
 operations (create / retrieve / cancel / list) over the shared Gateway.
 
-`create_batch` returns the raw HTTP outcome (StreamingResponse) so a 403 model
-access denial and a provider-native batch body both surface; the test parses
-BatchObject from the body. A `provider` arg routes a call to /{provider}/v1/...,
-which the provider-fallback scenario needs (its ids are raw, not model-encoded).
-The request/response models are co-located here because only this suite uses them.
+Batch deployments are registered at runtime via /model/new (see conftest.py),
+not baked into the proxy config. `create_batch` returns the raw HTTP outcome
+(StreamingResponse) so a 403 model access denial and a provider-native batch
+body both surface; the test parses BatchObject from the body. A `provider` arg
+routes a call to /{provider}/v1/..., which the provider-fallback scenario needs
+(its ids are raw, not model-encoded). The request/response models are
+co-located here because only this suite uses them.
 """
 
 from __future__ import annotations
@@ -22,6 +24,7 @@ from e2e_http import (
     StreamingResponse,
     UnknownApiError,
 )
+from models import LiteLLMParamsBody
 
 
 class FileObject(BaseModel):
@@ -84,6 +87,12 @@ def is_result_access_denied[R: BaseModel](result: Result[R]) -> bool:
 class BatchClient:
     gateway: Gateway
 
+    def create_model(self, model_name: str, litellm_params: LiteLLMParamsBody) -> str:
+        return self.gateway.create_model(model_name, litellm_params, mode="batch")
+
+    def delete_model(self, model_id: str) -> None:
+        self.gateway.delete_model(model_id)
+
     def upload_file(
         self,
         *,

tests/e2e/batches/capabilities.py

@@ -13,6 +13,8 @@ import base64
 from dataclasses import dataclass
 from typing import Literal
 
+from models import LiteLLMParamsBody
+
 Scenario = Literal["encoded", "unified", "model_param", "provider_fallback"]
 
 IdShape = Literal["managed", "model_encoded", "raw"]
@@ -33,6 +35,39 @@ class Provider:
     can_cancel: bool
     can_list: bool
 
+    def litellm_params(self) -> LiteLLMParamsBody:
+        match self.name:
+            case "openai":
+                return LiteLLMParamsBody(
+                    model="openai/gpt-4o-mini",
+                    api_key="os.environ/OPENAI_API_KEY",
+                )
+            case "azure":
+                return LiteLLMParamsBody(
+                    model="azure/gpt-4.1-mini-batch",
+                    api_base="os.environ/AZURE_API_BASE",
+                    api_key="os.environ/AZURE_API_KEY",
+                    api_version="2024-07-01-preview",
+                )
+            case "vertex_ai":
+                return LiteLLMParamsBody(
+                    model="vertex_ai/gemini-2.5-flash",
+                    vertex_project="os.environ/VERTEXAI_PROJECT",
+                    vertex_location="us-central1",
+                    vertex_credentials="os.environ/VERTEXAI_CREDENTIALS",
+                )
+            case "bedrock":
+                return LiteLLMParamsBody(
+                    model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
+                    s3_access_key_id="os.environ/AWS_ACCESS_KEY_ID",
+                    s3_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY",
+                    s3_region_name="os.environ/AWS_REGION",
+                    s3_bucket_name="os.environ/AWS_BATCH_S3_BUCKET",
+                    aws_batch_role_arn="os.environ/AWS_BATCH_ROLE_ARN",
+                )
+            case _:
+                raise ValueError(f"unknown batch provider: {self.name!r}")
+
 
 @dataclass(frozen=True, slots=True)
 class Capability:
@@ -66,19 +101,28 @@ PROVIDERS: tuple[Provider, ...] = (
     Provider(
         "vertex_ai", "vertex-batch", "gemini-2.5-flash", can_cancel=True, can_list=True
     ),
-    # Provider(
-    #     "bedrock",
-    #     "bedrock-batch",
-    #     "us.anthropic.claude-haiku-4-5-20251001-v1:0",
-    #     can_cancel=False,
-    #     can_list=False,
-    # ),
+    Provider(
+        "bedrock",
+        "bedrock-batch",
+        "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
+        can_cancel=False,
+        can_list=False,
+    ),
 )
 
+BEDROCK_SCENARIOS: tuple[Scenario, ...] = ("encoded", "unified")
+
+
+def scenarios_for_provider(provider: Provider) -> tuple[Scenario, ...]:
+    if provider.name == "bedrock":
+        return BEDROCK_SCENARIOS
+    return SCENARIOS
+
+
 CAPABILITIES: tuple[Capability, ...] = tuple(
     Capability(p.name, p.model, p.raw_model, scenario, p.can_cancel, p.can_list)
     for p in PROVIDERS
-    for scenario in SCENARIOS
+    for scenario in scenarios_for_provider(p)
 )
 
 
@@ -88,17 +132,13 @@ def raw_id_matches_provider(provider: str, batch_id: str) -> bool:
     if provider in ("openai", "azure"):
         return batch_id.startswith("batch")
     if provider == "vertex_ai":
-        # Vertex returns the batch prediction job id, which depending on the
-        # routing path arrives either as the full resource name
-        # (projects/.../batchPredictionJobs/<id>) or as just the trailing
-        # numeric id, so accept either form.
         return (
             batch_id.startswith("projects/")
             or "batchPredictionJobs" in batch_id
             or batch_id.isdigit()
         )
     if provider == "bedrock":
-        return batch_id.startswith("arn:aws")
+        return batch_id.startswith("arn:aws:bedrock:")
     return True
 
 

tests/e2e/llm_translation/test_ocr_rust_e2e.py

@@ -1,31 +1,32 @@
 """Live e2e: Rust-backed OCR is reachable through the gateway across providers.
 
-The gateway config declares one rust-ocr deployment per provider (mistral,
-azure_ai, azure document intelligence, vertex mistral, vertex deepseek). Start the
-proxy with the Rust OCR path enabled:
+Each provider's OCR deployment is registered at runtime via /model/new and deleted
+on teardown, so nothing is hardcoded into the gateway config. Every provider is its
+own typed OcrProvider below: it owns the model id and the os.environ/* credential
+references the proxy resolves at call time, so adding a provider is a new type
+rather than another inline body. Start the proxy with the Rust OCR path enabled:
 
-    LITELLM_USE_RUST_OCR=1 litellm --config tests/e2e/gateway/litellm-config.yml
-
-Three behaviors are checked: the config declares every provider's deployment (a
-pure config read, no proxy needed); the running proxy loaded them onto /model/info;
-and each one returns a well-formed OCR document over /v1/ocr. Per the e2e
-"skip on environment, fail on behavior" rule, the proxy-backed cases skip when no
-proxy answers but fail (never skip) once a request reaches it, so a provider whose
-credentials are missing surfaces as a hard failure rather than silent green.
+Each case creates its deployment, drives a real /v1/ocr call, and asserts a
+well-formed OCR document comes back. Per the e2e "skip on environment, fail on
+behavior" rule, a case skips when no proxy answers but fails (never skips) once a
+request reaches it: the proxy fetches each provider's referenced secrets, so a
+missing credential surfaces as a live provider error rather than silent green.
 """
 
 from __future__ import annotations
 
 from dataclasses import dataclass
-from pathlib import Path
+from typing import Protocol
 
 import pytest
-import yaml
-from pydantic import BaseModel
 
+from e2e_config import unique_marker
 from e2e_http import unwrap
-from models import OcrBody, OcrDocument, OcrResponse
-from passthrough_client import PassthroughClient
+from endpoints_client import EndpointsClient
+from lifecycle import ResourceManager
+from models import LiteLLMParamsBody, OcrBody, OcrDocument, OcrResponse
+
+pytestmark = pytest.mark.e2e
 
 # Tiny in-repo fixtures served via jsdelivr (sha-pinned, immutable) so the request
 # bodies stay stable across runs.
@@ -40,53 +41,99 @@ TEST_IMAGE_URL = (
     "/tests/image_gen_tests/test_image.png"
 )
 
-CONFIG_PATH = Path(__file__).resolve().parents[1] / "gateway" / "litellm-config.yml"
+
+class OcrProvider(Protocol):
+    """One OCR provider's deployment config: its model id plus the os.environ/*
+    credential references the proxy resolves at call time. Each provider owns which
+    env vars it reads, so a new provider is a new type, not another inline body."""
+
+    def litellm_params(self) -> LiteLLMParamsBody: ...
+
+
+@dataclass(frozen=True, slots=True)
+class MistralOcr:
+    model: str = "mistral/mistral-ocr-latest"
+
+    def litellm_params(self) -> LiteLLMParamsBody:
+        return LiteLLMParamsBody(model=self.model, api_key="os.environ/MISTRAL_API_KEY")
+
+
+@dataclass(frozen=True, slots=True)
+class AzureAiOcr:
+    """azure_ai (mistral) OCR. The rust OCR path resolves credentials itself from
+    AZURE_AI_API_BASE / AZURE_AI_API_KEY when the deployment leaves them unset; it
+    does NOT unwrap an `os.environ/*` reference passed as api_base (it would be sent
+    to Azure verbatim), so we omit them and let litellm read the env vars by name."""
+
+    model: str
+
+    def litellm_params(self) -> LiteLLMParamsBody:
+        return LiteLLMParamsBody(model=self.model)
+
+
+@dataclass(frozen=True, slots=True)
+class AzureDocIntelligenceOcr:
+    """azure_ai Document Intelligence OCR. A separate Azure resource from the
+    mistral one, so it has its own env vars: AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT /
+    AZURE_DOCUMENT_INTELLIGENCE_API_KEY, which the OCR config resolves from the
+    doc-intelligence model name when api_base/api_key are left unset."""
+
+    model: str = "azure_ai/doc-intelligence/prebuilt-layout"
+
+    def litellm_params(self) -> LiteLLMParamsBody:
+        return LiteLLMParamsBody(model=self.model)
+
+
+@dataclass(frozen=True, slots=True)
+class VertexOcr:
+    model: str
+    location: str
+
+    def litellm_params(self) -> LiteLLMParamsBody:
+        return LiteLLMParamsBody(
+            model=self.model,
+            vertex_project="os.environ/VERTEXAI_PROJECT",
+            vertex_location=self.location,
+            vertex_credentials="os.environ/VERTEXAI_CREDENTIALS",
+        )
 
 
 @dataclass(frozen=True, slots=True)
 class _OcrCase:
-    model: str
+    suffix: str
+    provider: OcrProvider
     document: OcrDocument
 
 
 RUST_OCR_CASES: tuple[_OcrCase, ...] = (
     _OcrCase(
-        "rust-ocr-mistral",
+        "mistral",
+        MistralOcr(),
         OcrDocument(type="document_url", document_url=TEST_PDF_URL),
     ),
     _OcrCase(
-        "rust-ocr-azure-ai",
+        "azure-ai",
+        AzureAiOcr("azure_ai/mistral-document-ai-2505"),
         OcrDocument(type="document_url", document_url=TEST_PDF_URL),
     ),
     _OcrCase(
-        "rust-ocr-azure-document-intelligence",
+        "azure-document-intelligence",
+        AzureDocIntelligenceOcr(),
         OcrDocument(type="document_url", document_url=TEST_PDF_URL),
     ),
     _OcrCase(
-        "rust-ocr-vertex-mistral",
+        "vertex-mistral",
+        VertexOcr("vertex_ai/mistral-ocr-2505", "us-central1"),
         OcrDocument(type="document_url", document_url=TEST_PDF_URL),
     ),
     _OcrCase(
-        "rust-ocr-vertex-deepseek",
+        "vertex-deepseek",
+        VertexOcr("vertex_ai/deepseek-ocr-maas", "global"),
         OcrDocument(type="image_url", image_url=TEST_IMAGE_URL),
     ),
 )
 
-_EXPECTED_MODELS = frozenset(case.model for case in RUST_OCR_CASES)
-_CASE_IDS = tuple(case.model.removeprefix("rust-ocr-") for case in RUST_OCR_CASES)
-
-
-class _ConfiguredModel(BaseModel):
-    model_name: str
-
-
-class _GatewayConfig(BaseModel):
-    model_list: list[_ConfiguredModel]
-
-
-def _configured_model_names() -> frozenset[str]:
-    config = _GatewayConfig.model_validate(yaml.safe_load(CONFIG_PATH.read_text()))
-    return frozenset(entry.model_name for entry in config.model_list)
+_CASE_IDS = tuple(case.suffix for case in RUST_OCR_CASES)
 
 
 def _assert_ocr_document(response: OcrResponse) -> None:
@@ -96,22 +143,15 @@ def _assert_ocr_document(response: OcrResponse) -> None:
     assert response.pages[0].markdown is not None, "first page has no markdown"
 
 
-def test_rust_ocr_models_declared_in_gateway_config() -> None:
-    """Pure config read (no proxy): every provider's rust-ocr deployment the suite
-    exercises is declared in the gateway config the proxy runs with. A case added
-    here without a matching deployment fails before any live call is attempted."""
-    missing = _EXPECTED_MODELS - _configured_model_names()
-    assert not missing, f"rust-ocr models absent from {CONFIG_PATH.name}: {missing}"
-
-
-@pytest.mark.e2e
 class TestRustOcrGateway:
-    def test_gateway_loaded_rust_ocr_models(self, client: PassthroughClient) -> None:
-        loaded = frozenset(entry.model_name for entry in client.gateway.model_info())
-        missing = _EXPECTED_MODELS - loaded
-        assert not missing, f"proxy did not load rust-ocr models: {missing}"
-
     @pytest.mark.parametrize("case", RUST_OCR_CASES, ids=_CASE_IDS)
-    def test_rust_ocr_response(self, client: PassthroughClient, scoped_key: str, case: _OcrCase) -> None:
-        response = unwrap(client.gateway.ocr(scoped_key, OcrBody(model=case.model, document=case.document)))
+    def test_rust_ocr_response(
+        self, endpoints_client: EndpointsClient, resources: ResourceManager, case: _OcrCase
+    ) -> None:
+        model = f"rust-ocr-{case.suffix}-{unique_marker()}"
+        model_id = endpoints_client.create_model(model, case.provider.litellm_params())
+        resources.defer(lambda: endpoints_client.delete_model(model_id))
+        key = resources.key()
+
+        response = unwrap(endpoints_client.gateway.ocr(key, OcrBody(model=model, document=case.document)))
         _assert_ocr_document(response)