diff --git a/.circleci/config.yml b/.circleci/config.yml index c940716264..544a5a1eed 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -3689,6 +3689,114 @@ jobs: - store_test_results: path: test-results + proxy_e2e_azure_batches_tests: + machine: + image: ubuntu-2204:2023.10.1 + resource_class: xlarge + working_directory: ~/project + steps: + - checkout + - setup_google_dns + - run: + name: Install Docker CLI + command: | + curl -fsSL https://get.docker.com | sh + sudo usermod -aG docker $USER + docker version + - run: + name: Install Python 3.12 + command: | + curl https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh --output miniconda.sh + bash miniconda.sh -b -p $HOME/miniconda + export PATH="$HOME/miniconda/bin:$PATH" + conda init bash + source ~/.bashrc + conda create -n myenv python=3.12 -y + conda activate myenv + python --version + - run: + name: Install Poetry + command: | + export PATH="$HOME/miniconda/bin:$PATH" + source $HOME/miniconda/etc/profile.d/conda.sh + conda activate myenv + pip install poetry + - run: + name: Install dockerize + command: | + wget https://github.com/jwilder/dockerize/releases/download/v0.6.1/dockerize-linux-amd64-v0.6.1.tar.gz + sudo tar -C /usr/local/bin -xzvf dockerize-linux-amd64-v0.6.1.tar.gz + rm dockerize-linux-amd64-v0.6.1.tar.gz + - run: + name: Start PostgreSQL Database + command: | + docker run -d \ + --name postgres-db \ + -e POSTGRES_USER=llmproxy \ + -e POSTGRES_PASSWORD=dbpassword9090 \ + -e POSTGRES_DB=litellm \ + -p 5432:5432 \ + postgres:15 + - run: + name: Wait for PostgreSQL to be ready + command: dockerize -wait tcp://localhost:5432 -timeout 1m + - run: + name: Install system dependencies + command: | + sudo apt-get update -y + sudo apt-get install -y libpq-dev + - run: + name: Install Dependencies + command: | + export PATH="$HOME/miniconda/bin:$PATH" + source $HOME/miniconda/etc/profile.d/conda.sh + conda activate myenv + poetry config virtualenvs.in-project true + poetry install --with dev,proxy-dev --extras "proxy" + poetry run pip install psycopg2-binary uvicorn fastapi httpx tenacity + - run: + name: Setup litellm-enterprise + command: | + export PATH="$HOME/miniconda/bin:$PATH" + source $HOME/miniconda/etc/profile.d/conda.sh + conda activate myenv + poetry run pip install --force-reinstall --no-deps -e enterprise/ + - run: + name: Generate Prisma client + command: | + export PATH="$HOME/miniconda/bin:$PATH" + source $HOME/miniconda/etc/profile.d/conda.sh + conda activate myenv + poetry run prisma generate --schema litellm/proxy/schema.prisma + - run: + name: Run Prisma migrations + command: | + export PATH="$HOME/miniconda/bin:$PATH" + source $HOME/miniconda/etc/profile.d/conda.sh + conda activate myenv + export DATABASE_URL=postgresql://llmproxy:dbpassword9090@localhost:5432/litellm + cd litellm/proxy + poetry run prisma migrate deploy --schema schema.prisma + cd ../.. + - run: + name: Run Azure Batch E2E Tests + command: | + export PATH="$HOME/miniconda/bin:$PATH" + source $HOME/miniconda/etc/profile.d/conda.sh + conda activate myenv + export DATABASE_URL=postgresql://llmproxy:dbpassword9090@localhost:5432/litellm + export USE_LOCAL_LITELLM=true + export USE_MOCK_MODELS=true + export USE_STATE_TRACKER=true + export LITELLM_LOG=DEBUG + poetry run pytest tests/proxy_e2e_azure_batches_tests/test_proxy_e2e_azure_batches.py \ + -vv -s -k "test_e2e_managed_batch" \ + --tb=short \ + --maxfail=3 \ + --durations=10 \ + --junitxml=test-results/junit.xml + no_output_timeout: 30m + upload-coverage: docker: - image: cimg/python:3.9 @@ -4458,6 +4566,12 @@ workflows: only: - main - /litellm_.*/ + - proxy_e2e_azure_batches_tests: + filters: + branches: + only: + - main + - /litellm_.*/ - llm_translation_testing: filters: branches: diff --git a/.github/workflows/test-proxy-e2e-azure-batches.yml b/.github/workflows/test-proxy-e2e-azure-batches.yml new file mode 100644 index 0000000000..38d436dc1f --- /dev/null +++ b/.github/workflows/test-proxy-e2e-azure-batches.yml @@ -0,0 +1,90 @@ +name: Proxy E2E Azure Batches Tests + +on: + pull_request: + branches: [main] + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +jobs: + proxy_e2e_azure_batches_tests: + runs-on: ubuntu-latest + timeout-minutes: 30 + + services: + postgres: + image: postgres:15 + env: + POSTGRES_USER: llmproxy + POSTGRES_PASSWORD: dbpassword9090 + POSTGRES_DB: litellm + ports: + - 5432:5432 + options: >- + --health-cmd pg_isready + --health-interval 10s + --health-timeout 5s + --health-retries 5 + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install Poetry + uses: snok/install-poetry@v1 + + - name: Cache Poetry dependencies + uses: actions/cache@v4 + with: + path: | + ~/.cache/pypoetry + ~/.cache/pip + .venv + key: ${{ runner.os }}-poetry-e2e-batches-${{ hashFiles('poetry.lock') }} + restore-keys: | + ${{ runner.os }}-poetry-e2e-batches- + ${{ runner.os }}-poetry- + + - name: Install dependencies + run: | + poetry config virtualenvs.in-project true + poetry install --with dev,proxy-dev --extras "proxy" + poetry run pip install psycopg2-binary uvicorn fastapi httpx + + - name: Setup litellm-enterprise + run: | + poetry run pip install --force-reinstall --no-deps -e enterprise/ + + - name: Generate Prisma client + run: | + poetry run prisma generate --schema litellm/proxy/schema.prisma + + - name: Run Prisma migrations + env: + DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm + run: | + cd litellm/proxy + poetry run prisma migrate deploy --schema schema.prisma + cd ../.. + + - name: Run Azure Batch E2E Tests + env: + DATABASE_URL: postgresql://llmproxy:dbpassword9090@localhost:5432/litellm + USE_LOCAL_LITELLM: "true" + USE_MOCK_MODELS: "true" + USE_STATE_TRACKER: "true" + LITELLM_LOG: DEBUG + run: | + poetry run pytest tests/proxy_e2e_azure_batches_tests/test_proxy_e2e_azure_batches.py \ + -vv -s -k "test_e2e_managed_batch" \ + --tb=short \ + --maxfail=3 \ + --durations=10 + diff --git a/tests/proxy_e2e_azure_batches_tests/__init__.py b/tests/proxy_e2e_azure_batches_tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/proxy_e2e_azure_batches_tests/base_integration_test.py b/tests/proxy_e2e_azure_batches_tests/base_integration_test.py new file mode 100644 index 0000000000..c819fa7bf4 --- /dev/null +++ b/tests/proxy_e2e_azure_batches_tests/base_integration_test.py @@ -0,0 +1,494 @@ +"""Base class for LiteLLM integration tests. + +Supports both local (mock) and remote testing modes via environment variables: +- USE_LOCAL_LITELLM: When "true", uses local LiteLLM at localhost:4000 (default: false) +- USE_MOCK_MODELS: When "true", uses mock model names (default: false) +- LITELLM_API_KEY: API key for remote LiteLLM (required when USE_LOCAL_LITELLM=false) +- LITELLM_BASE_URL: Base URL for remote LiteLLM (required when USE_LOCAL_LITELLM=false) +""" + +import enum +import os +import time +import uuid +from abc import ABC +from collections import defaultdict +from typing import Any, Callable, Dict, List, Tuple, Union + +import httpx +import openai +import pytest +import requests +from urllib3.exceptions import InsecureRequestWarning + +requests.packages.urllib3.disable_warnings(category=InsecureRequestWarning) + +LOCAL_LITELLM_BASE_URL = "http://localhost:4000" +LOCAL_MOCK_SERVER_URL = "http://localhost:8090" + +if "USE_LOCAL_LITELLM" not in os.environ: + os.environ["USE_LOCAL_LITELLM"] = "true" +if "USE_MOCK_MODELS" not in os.environ: + os.environ["USE_MOCK_MODELS"] = "true" +if "USE_STATE_TRACKER" not in os.environ: + os.environ["USE_STATE_TRACKER"] = "true" +if "DATABASE_URL" not in os.environ: + os.environ["DATABASE_URL"] = "postgresql://llmproxy:dbpassword9090@localhost:5432/litellm" + + +def use_local_litellm() -> bool: + return os.environ.get("USE_LOCAL_LITELLM", "false").lower() == "true" + + +def use_remote_litellm() -> bool: + return not use_local_litellm() + + +def use_mock_models() -> bool: + return os.environ.get("USE_MOCK_MODELS", "false").lower() == "true" + + +def get_local_litellm_base_url() -> str: + return LOCAL_LITELLM_BASE_URL + + +def get_remote_litellm_base_url() -> str: + return os.environ.get("LITELLM_BASE_URL", "").rstrip("/") + + +def get_litellm_base_url() -> str: + if use_local_litellm(): + return get_local_litellm_base_url() + return get_remote_litellm_base_url() + + +def get_litellm_api_key() -> str: + if use_local_litellm(): + return "sk-1234" + return os.environ.get("LITELLM_API_KEY", "") + + +def get_mock_server_base_url() -> str: + return LOCAL_MOCK_SERVER_URL + + +def get_responses_model_name() -> str: + if use_mock_models(): + return "openai-fake-gpt-4o" + return "gpt-4o-mini-2024-07-18" + + +def model_id(param) -> str: + """Generate a test ID from a model name or tuple containing model name. + + Handles both: + - String: "gpt-4o-mini" -> "gpt_4o_mini" + - Tuple: ("gpt-4o", "openai/gpt-4o") -> "gpt_4o" + """ + if isinstance(param, tuple): + name = param[0] + else: + name = param + return name.replace("-", "_").replace(".", "_") + + +def generate_test_id( + params: Tuple[str, ...], + test_name: str = "test", +) -> str: + """Generate test ID from model parameters tuple. + + Handles two tuple formats: + - 6 elements: (provider, deployment, model_name, api_version, action, reason) + - 7 elements: (provider, deployment, model_name, api_version, model_id, action, reason) + + Uses model_id (position 4) if 7 elements, otherwise model_name (position 2). + """ + provider = params[0] + deployment = params[1] + api_version = params[3] + + if len(params) == 7: + identifier = params[4] # model_id + else: + identifier = params[2] # model_name + + test_id = "/".join([provider, deployment, api_version, identifier, test_name]) + return test_id.replace("-", "_").replace(".", "_") + + +class ModelTestAction(enum.Enum): + NOT_APPLICABLE = 1 + SKIP = 2 + RUN = 3 + WARN_ON_FAIL = 4 + + def applicable(self) -> bool: + return self.value != ModelTestAction.NOT_APPLICABLE.value + + +class BaseLiteLLMIntegrationTest(ABC): + """Base class for all LiteLLM integration tests. + + Supports both local/mock and remote testing based on environment variables. + """ + + @staticmethod + def get_api_key() -> str: + return get_litellm_api_key() + + @staticmethod + def get_base_url() -> str: + return get_litellm_base_url() + + @staticmethod + def get_ca_bundle_path() -> str: + current_dir = os.path.dirname(os.path.abspath(__file__)) + # change if needed + + @classmethod + def _get_ssl_verify_setting(cls) -> Union[bool, str]: + """Get the appropriate SSL verification setting based on mode. + + Returns path string (not SSLContext) for compatibility with both + requests and httpx libraries. + """ + if use_local_litellm(): + return False + ca_bundle_path = cls.get_ca_bundle_path() + if os.path.exists(ca_bundle_path): + return ca_bundle_path + return True + + @classmethod + def setup_class(cls): + cls.api_key = cls.get_api_key() + cls.base_url = cls.get_base_url() + + if not cls.api_key: + pytest.fail( + "API key is not available. Set LITELLM_API_KEY or USE_LOCAL_LITELLM=true", + ) + if not cls.base_url: + pytest.fail( + "Base URL is not available. Set LITELLM_BASE_URL or USE_LOCAL_LITELLM=true", + ) + + verify_setting = cls._get_ssl_verify_setting() + + if use_remote_litellm() and isinstance(verify_setting, str): + os.environ["REQUESTS_CA_BUNDLE"] = verify_setting + os.environ["CURL_CA_BUNDLE"] = verify_setting + print(f"Using CA bundle: {verify_setting}") + + cls.openai_client = openai.OpenAI( + base_url=cls.base_url, + api_key=cls.api_key, + http_client=httpx.Client(verify=verify_setting), + ) + + @classmethod + def make_request( + cls, + method: str, + endpoint: str, + timeout_secs: int, + **kwargs, + ) -> requests.Response: + headers = kwargs.get("headers", {}) + headers["Authorization"] = f"Bearer {cls.api_key}" + kwargs["headers"] = headers + kwargs.setdefault("timeout", timeout_secs) + kwargs.setdefault("verify", cls._get_ssl_verify_setting()) + + url = f"{cls.base_url}{endpoint}" + return requests.request(method, url, **kwargs) + + @staticmethod + def generate_request_id() -> str: + return f"req-{uuid.uuid4().hex[:8]}" + + @staticmethod + def get_timeout_secs(model_name: str) -> int: + model_lower = model_name.lower() + slow_models = ["gpt-5", "gpt_5", "o1", "claude-opus", "claude_opus", "o3", "o4"] + + if any(slow_model in model_lower for slow_model in slow_models): + return 300 + return 60 + + @staticmethod + def generate_unique_filename(extension: str = "txt") -> str: + return f"test_{time.time()}.{extension}" + + @staticmethod + def extract_model_params(model_data: Dict[str, Any]) -> Tuple[str, str, str, str]: + """Extract standardized parameters from model data.""" + model_name = model_data.get("model_name", "") + model_info = model_data.get("model_info", {}) + provider = model_info.get("litellm_provider", "unknown") + litellm_params = model_data.get("litellm_params", {}) + + if provider == "azure": + api_base = litellm_params.get("api_base", "unknown") + if api_base != "unknown" and "//" in api_base: + domain_name = api_base.split("//")[1] + deployment = domain_name.split(".")[0] + else: + deployment = "unknown" + api_version = litellm_params.get("api_version", "unknown") + elif provider in ["bedrock", "bedrock_converse"]: + deployment = litellm_params.get("aws_region_name", "unknown") + api_version = "unknown" + else: + deployment = "unknown" + api_version = "unknown" + + return provider, deployment, model_name, api_version + + @classmethod + def _fetch_all_models_from_litellm(cls) -> List[Dict[str, Any]]: + base_url = cls.get_base_url() + api_key = cls.get_api_key() + + if not api_key or not base_url: + return [] + + verify_setting = cls._get_ssl_verify_setting() + + response = requests.get( + f"{base_url}/model/info", + headers={"Authorization": f"Bearer {api_key}"}, + verify=verify_setting, + timeout=30, + ) + + if response.status_code != 200: + raise RuntimeError( + f"Failed to fetch all models from {base_url}. Response code: {response.status_code}", + ) + + data = response.json() + return data.get("data", []) + + @classmethod + def _fetch_all_approved_models(cls) -> List[Dict[str, Any]]: + return cls._fetch_all_models_from_litellm() + + @classmethod + def build_model_test_params( + cls, + should_skip_model: Callable[ + [str, str, str, str, Dict[str, Any]], + Tuple["ModelTestAction", str], + ], + include_model_id: bool = False, + include_load_balanced: bool = False, + ) -> List[Tuple[str, ...]]: + """Build test parameters from all approved models. + + Args: + should_skip_model: Callback that determines if a model should be skipped. + Signature: (provider, deployment, model_name, api_version, model_info) -> (action, reason) + include_model_id: If True, includes model_id in tuple (7 elements), else 6 elements. + include_load_balanced: If True, adds extra tests for load-balanced model groups. + + Returns: + List of tuples with model test parameters. + - 6-element: (provider, deployment, model_name, api_version, action, reason) + - 7-element: (provider, deployment, model_name, api_version, model_id, action, reason) + """ + models = cls._fetch_all_approved_models() + test_params: List[Tuple[str, ...]] = [] + models_by_model_name: Dict[str, List[Tuple[str, ...]]] = defaultdict(list) + + for model_data in models: + model_info = model_data.get("model_info", {}) or {} + + provider, deployment, model_name, api_version = cls.extract_model_params( + model_data, + ) + + model_test_action, model_test_action_reason = should_skip_model( + provider, + deployment, + model_name, + api_version, + model_info, + ) + + if model_test_action.applicable(): + if include_model_id: + model_id = str(model_info.get("id")) + params_tuple: Tuple[str, ...] = ( + provider, + deployment, + model_name, + api_version, + model_id, + model_test_action, + model_test_action_reason, + ) + else: + params_tuple = ( + provider, + deployment, + model_name, + api_version, + model_test_action, + model_test_action_reason, + ) + + test_params.append(params_tuple) + + if include_load_balanced: + models_by_model_name[model_name].append(params_tuple) + + if include_load_balanced and include_model_id: + for load_balanced_model_name, deployments in models_by_model_name.items(): + if len(deployments) <= 1: + continue + + first_deployment = deployments[0] + test_params.append( + ( + first_deployment[0], # provider + "load_balanced", + load_balanced_model_name, + "load_balanced", + load_balanced_model_name, # model_id = model_name for LB + first_deployment[5], # model_test_action + first_deployment[6], # model_test_action_reason + ), + ) + + return test_params + + +class UserKeyTestMixin: + """Mixin for tests that need to create users and API keys.""" + + allowed_routes: list[str] = [] + + _base_url: str = None + _master_api_key: str = None + admin_client: httpx.Client = None + + @classmethod + def setup_admin_client(cls): + cls._base_url = get_litellm_base_url() + cls._master_api_key = get_litellm_api_key() + verify_setting = ( + False + if use_local_litellm() + else BaseLiteLLMIntegrationTest._get_ssl_verify_setting() + ) + cls.admin_client = httpx.Client(base_url=cls._base_url, verify=verify_setting) + + @classmethod + def teardown_admin_client(cls): + if cls.admin_client: + cls.admin_client.close() + + @staticmethod + def unique_suffix() -> str: + return f"{time.strftime('%Y%m%d%H%M%S')}{int(time.time() * 1000) % 1000:03d}" + + @classmethod + def create_user_and_key(cls, user_suffix: str) -> tuple[str, str, str]: + user_email = f"test-user-{user_suffix}-{cls.unique_suffix()}@test.com" + user_response = cls.admin_client.post( + "/user/new", + json={ + "user_email": user_email, + "user_alias": user_email, + "user_role": "internal_user", + "auto_create_key": "false", + }, + headers={ + "Authorization": f"Bearer {cls._master_api_key}", + "Content-Type": "application/json", + }, + timeout=30, + ) + assert user_response.status_code == 200, ( + f"Failed to create user: {user_response.status_code} - {user_response.text}" + ) + user_id = user_response.json().get("user_id") + + key_alias = user_email.replace("@", "-at-").replace(".", "-") + key_response = cls.admin_client.post( + "/key/generate", + json={ + "user_id": user_id, + "key_alias": key_alias, + "allowed_routes": cls.allowed_routes, + }, + headers={ + "Authorization": f"Bearer {cls._master_api_key}", + "Content-Type": "application/json", + }, + timeout=30, + ) + assert key_response.status_code == 200, ( + f"Failed to create key: {key_response.status_code} - {key_response.text}" + ) + api_key = key_response.json().get("key") + + print(f"Created user {user_email}") + return user_id, api_key, user_email + + @classmethod + def create_user_key_and_client( + cls, + user_suffix: str, + ) -> tuple[str, str, str, openai.OpenAI]: + user_id, api_key, user_email = cls.create_user_and_key(user_suffix) + verify_setting = ( + False + if use_local_litellm() + else BaseLiteLLMIntegrationTest._get_ssl_verify_setting() + ) + client = openai.OpenAI( + base_url=cls._base_url, + api_key=api_key, + http_client=httpx.Client(verify=verify_setting), + ) + return user_id, api_key, user_email, client + + @classmethod + def create_key_and_client( + cls, + user_id: str, + key_suffix: str, + ) -> tuple[str, openai.OpenAI]: + key_alias = f"additional-key-{key_suffix}-{cls.unique_suffix()}" + key_response = cls.admin_client.post( + "/key/generate", + json={ + "user_id": user_id, + "key_alias": key_alias, + "allowed_routes": cls.allowed_routes, + }, + headers={ + "Authorization": f"Bearer {cls._master_api_key}", + "Content-Type": "application/json", + }, + timeout=30, + ) + assert key_response.status_code == 200, ( + f"Failed to create additional key: {key_response.status_code} - {key_response.text}" + ) + api_key = key_response.json().get("key") + verify_setting = ( + False + if use_local_litellm() + else BaseLiteLLMIntegrationTest._get_ssl_verify_setting() + ) + client = openai.OpenAI( + base_url=cls._base_url, + api_key=api_key, + http_client=httpx.Client(verify=verify_setting), + ) + print(f"Created additional key for user {user_id}") + return api_key, client \ No newline at end of file diff --git a/tests/proxy_e2e_azure_batches_tests/conftest.py b/tests/proxy_e2e_azure_batches_tests/conftest.py new file mode 100644 index 0000000000..1bad010a20 --- /dev/null +++ b/tests/proxy_e2e_azure_batches_tests/conftest.py @@ -0,0 +1,311 @@ +""" +Pytest configuration for Azure Batch E2E Tests. + +This conftest manages: +1. Mock Azure Batch server (FastAPI on port 8090) +2. LiteLLM proxy server (port 4000) +3. PostgreSQL database setup +""" + +import asyncio +import os +import subprocess +import sys +import time +from pathlib import Path +from typing import Generator + +import httpx +import pytest + +_test_dir = Path(__file__).parent +sys.path.insert(0, str(_test_dir.parent.parent)) # litellm root +sys.path.insert(0, str(_test_dir)) # test directory for local imports + +LOG_DIR = _test_dir + + +def pytest_configure(config): + """Ensure test directory is in Python path before collection.""" + test_dir = Path(__file__).parent + if str(test_dir) not in sys.path: + sys.path.insert(0, str(test_dir)) + + +MOCK_SERVER_PORT = 8090 +MOCK_SERVER_URL = f"http://localhost:{MOCK_SERVER_PORT}" +LITELLM_PROXY_PORT = 4000 +LITELLM_PROXY_URL = f"http://localhost:{LITELLM_PROXY_PORT}" +DATABASE_URL = "postgresql://llmproxy:dbpassword9090@localhost:5432/litellm" + + +def kill_process_on_port(port: int) -> None: + """Kill any process using the specified port.""" + try: + result = subprocess.run( + ["lsof", "-ti", f":{port}"], + capture_output=True, + text=True, + timeout=5, + ) + if result.stdout.strip(): + pids = result.stdout.strip().split("\n") + for pid in pids: + try: + subprocess.run(["kill", "-9", pid.strip()], timeout=5) + except Exception: + pass + time.sleep(1) + except Exception: + pass + + +def wait_for_server(url: str, max_attempts: int = 30, delay: float = 1.0) -> bool: + """Wait for a server to become available at url/health. + + Any HTTP response (including 401) means the server is up. + Only connection errors count as "not ready yet". + """ + for attempt in range(max_attempts): + try: + response = httpx.get(f"{url}/health", timeout=2.0) + return True + except (httpx.ConnectError, httpx.TimeoutException, httpx.NetworkError): + pass + except Exception: + pass + if attempt < max_attempts - 1: + time.sleep(delay) + return False + + +def _read_log_tail(log_path: Path, max_lines: int = 80) -> str: + """Read the last N lines of a log file, returning empty string if not found.""" + if not log_path.exists(): + return "(log file not found)" + try: + text = log_path.read_text() + lines = text.strip().splitlines() + if len(lines) > max_lines: + return f"... ({len(lines) - max_lines} lines truncated) ...\n" + "\n".join( + lines[-max_lines:] + ) + return text + except Exception as e: + return f"(error reading log: {e})" + + +def _check_process_alive(process: subprocess.Popen, label: str, log_path: Path): + """Check if a subprocess crashed immediately after starting. + Raises pytest.fail with log output if the process has already exited. + """ + time.sleep(1) + exit_code = process.poll() + if exit_code is not None: + log_output = _read_log_tail(log_path) + pytest.fail( + f"{label} exited immediately with code {exit_code}.\n" + f"--- {label} log ({log_path}) ---\n{log_output}\n" + f"--- end log ---" + ) + + +def setup_database() -> bool: + """Ensure PostgreSQL database exists and is accessible.""" + try: + import psycopg2 + + conn = psycopg2.connect( + host="localhost", + port=5432, + database="litellm", + user="llmproxy", + password="dbpassword9090", + connect_timeout=5, + ) + conn.close() + return True + except ImportError: + print("WARNING: psycopg2 not installed — cannot verify database") + return False + except Exception: + return False + + +@pytest.fixture(scope="session") +def mock_azure_server() -> Generator[str, None, None]: + """Start mock Azure batch server as a subprocess.""" + print(f"\n{'=' * 60}") + print("Setting up Mock Azure Batch Server") + print(f"{'=' * 60}") + + kill_process_on_port(MOCK_SERVER_PORT) + + runner_script = Path(__file__).parent / "fixtures" / "run_mock_server.py" + runner_script.write_text( + """ +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from fixtures.mock_azure_batch_server import create_mock_azure_batch_server +import uvicorn + +if __name__ == "__main__": + app = create_mock_azure_batch_server() + uvicorn.run(app, host="0.0.0.0", port=8090, log_level="info", access_log=False) +""" + ) + + mock_log = LOG_DIR / "mock_server.log" + log_file = open(mock_log, "w") + + print(f"Starting mock server on port {MOCK_SERVER_PORT}...") + print(f"Log file: {mock_log}") + process = subprocess.Popen( + [sys.executable, str(runner_script)], + stdout=log_file, + stderr=subprocess.STDOUT, + cwd=Path(__file__).parent, + ) + + _check_process_alive(process, "Mock server", mock_log) + + if not wait_for_server(MOCK_SERVER_URL, max_attempts=30, delay=1.0): + log_output = _read_log_tail(mock_log) + exit_code = process.poll() + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + log_file.close() + pytest.fail( + f"Mock server failed to start on port {MOCK_SERVER_PORT} " + f"(process exit_code={exit_code}).\n" + f"--- mock server log ---\n{log_output}\n--- end log ---\n" + f"Hint: ensure 'uvicorn' and 'fastapi' are installed." + ) + + print(f"Mock Azure server ready at {MOCK_SERVER_URL}") + yield MOCK_SERVER_URL + + print("\nShutting down mock server...") + try: + process.terminate() + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + log_file.close() + print("Mock server stopped") + + +@pytest.fixture(scope="session") +def litellm_proxy_server(mock_azure_server: str) -> Generator[str, None, None]: + """Start LiteLLM proxy server for the test session.""" + print(f"\n{'=' * 60}") + print("Setting up LiteLLM Proxy Server") + print(f"{'=' * 60}") + + if not setup_database(): + pytest.skip( + "PostgreSQL database not available at localhost:5432. " + "Start PostgreSQL and create a 'litellm' database:\n" + " docker run -d --name litellm-db -p 5432:5432 " + '-e POSTGRES_USER=llmproxy -e POSTGRES_PASSWORD=dbpassword9090 ' + "-e POSTGRES_DB=litellm postgres:15\n" + "Then run: prisma db push --schema=litellm/proxy/schema.prisma" + ) + print("Database connection verified") + + config_path = Path(__file__).parent / "fixtures" / "config.yml" + if not config_path.exists(): + pytest.fail(f"Config file not found: {config_path}") + print("Config file found") + + kill_process_on_port(LITELLM_PROXY_PORT) + + os.environ["MOCK_SERVER_URL_V1"] = f"{mock_azure_server}/v1" + os.environ["MOCK_SERVER_URL_OPENAI_V1"] = f"{mock_azure_server}/openai/v1" + os.environ["DATABASE_URL"] = DATABASE_URL + os.environ["USE_LOCAL_LITELLM"] = "true" + os.environ["USE_MOCK_MODELS"] = "true" + os.environ["USE_STATE_TRACKER"] = "true" + os.environ["PROXY_BATCH_POLLING_INTERVAL"] = "10" + + print("Environment configured") + + print(f"Starting LiteLLM proxy on port {LITELLM_PROXY_PORT}...") + litellm_root = Path(__file__).parent.parent.parent + + cmd = [ + sys.executable, + "-m", + "litellm.proxy.proxy_cli", + "--config", + str(config_path), + "--port", + str(LITELLM_PROXY_PORT), + "--detailed_debug", + ] + + proxy_log = LOG_DIR / "proxy_server.log" + log_file = open(proxy_log, "w") + print(f"Log file: {proxy_log}") + + process = subprocess.Popen( + cmd, + stdout=log_file, + stderr=subprocess.STDOUT, + env=os.environ.copy(), + cwd=litellm_root, + ) + + _check_process_alive(process, "LiteLLM proxy", proxy_log) + + if not wait_for_server(LITELLM_PROXY_URL, max_attempts=60, delay=1.0): + log_output = _read_log_tail(proxy_log) + exit_code = process.poll() + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + log_file.close() + pytest.fail( + f"LiteLLM proxy failed to start on port {LITELLM_PROXY_PORT} " + f"(process exit_code={exit_code}).\n" + f"--- proxy log (last 80 lines) ---\n{log_output}\n--- end log ---\n" + f"Hints:\n" + f" 1. Ensure Prisma client is generated: " + f"cd {litellm_root} && prisma generate --schema=litellm/proxy/schema.prisma\n" + f" 2. Ensure DB migrations are applied: " + f"prisma db push --schema=litellm/proxy/schema.prisma\n" + f" 3. Check the full log at: {proxy_log}" + ) + + print(f"LiteLLM proxy ready at {LITELLM_PROXY_URL}") + yield LITELLM_PROXY_URL + + print("\nShutting down LiteLLM proxy...") + try: + process.terminate() + process.wait(timeout=10) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + log_file.close() + print("LiteLLM proxy stopped") + + +@pytest.fixture(scope="session") +def event_loop(): + """Provide an event loop for async tests.""" + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + yield loop + loop.close() diff --git a/tests/proxy_e2e_azure_batches_tests/fixtures/__init__.py b/tests/proxy_e2e_azure_batches_tests/fixtures/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/proxy_e2e_azure_batches_tests/fixtures/config.yml b/tests/proxy_e2e_azure_batches_tests/fixtures/config.yml new file mode 100644 index 0000000000..c991a32aab --- /dev/null +++ b/tests/proxy_e2e_azure_batches_tests/fixtures/config.yml @@ -0,0 +1,56 @@ +model_list: + - model_name: openai-fake-gpt-3.5-turbo + litellm_params: + model: openai/openai-fake-gpt-3.5-turbo + api_base: os.environ/MOCK_SERVER_URL_V1 + api_key: fake-key + - model_name: openai-fake-gpt-4 + litellm_params: + model: openai/openai-fake-gpt-4 + api_base: os.environ/MOCK_SERVER_URL_V1 + api_key: fake-key + - model_name: openai-fake-gpt-4o + litellm_params: + model: openai/openai-fake-gpt-4o + api_base: os.environ/MOCK_SERVER_URL_V1 + api_key: fake-key + - model_name: fake-text-embedding-3-small + litellm_params: + model: openai/fake-text-embedding-3-small + api_base: os.environ/MOCK_SERVER_URL_V1 + api_key: fake-key + - model_name: o3-mini-batch-2025-01-31 + litellm_params: + model: openai/o3-mini-batch-2025-01-31 + api_base: os.environ/MOCK_SERVER_URL_OPENAI_V1 + api_key: fake-key + model_info: + mode: batch + - model_name: azure-fake-gpt-5-batch-2025-08-07 + litellm_params: + api_base: http://0.0.0.0:8090 + api_key: fake-key + api_version: 2025-03-01-preview + base_model: azure/gpt-5 + model: azure/gpt-5-mini + custom_llm_provider: azure + +general_settings: + master_key: sk-1234 + database_url: os.environ/DATABASE_URL + proxy_batch_polling_interval: 10 + +litellm_settings: + drop_params: true + set_verbose: true + json_logs: true + # S3 callback for batch completion logging (points to mock server) + callbacks: ["s3_v2"] + s3_callback_params: + s3_bucket_name: litellm-test-bucket + s3_region_name: us-east-1 + s3_endpoint_url: http://0.0.0.0:8090 + s3_aws_access_key_id: fake-key + s3_aws_secret_access_key: fake-secret + s3_use_ssl: false + s3_verify: false \ No newline at end of file diff --git a/tests/proxy_e2e_azure_batches_tests/fixtures/mock_azure_batch_server/__init__.py b/tests/proxy_e2e_azure_batches_tests/fixtures/mock_azure_batch_server/__init__.py new file mode 100644 index 0000000000..3452b3aa50 --- /dev/null +++ b/tests/proxy_e2e_azure_batches_tests/fixtures/mock_azure_batch_server/__init__.py @@ -0,0 +1,3 @@ +from .server import create_mock_azure_batch_server + +__all__ = ["create_mock_azure_batch_server"] diff --git a/tests/proxy_e2e_azure_batches_tests/fixtures/mock_azure_batch_server/mock_azure_batch.py b/tests/proxy_e2e_azure_batches_tests/fixtures/mock_azure_batch_server/mock_azure_batch.py new file mode 100644 index 0000000000..940f32f595 --- /dev/null +++ b/tests/proxy_e2e_azure_batches_tests/fixtures/mock_azure_batch_server/mock_azure_batch.py @@ -0,0 +1,517 @@ +import asyncio +import io +import json +import logging +import time +import uuid +from typing import Dict, List, Optional + +from fastapi import FastAPI, HTTPException, Query, Request, UploadFile +from fastapi.responses import StreamingResponse +from pydantic import BaseModel + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + + +class FileObject(BaseModel): + id: str + object: str = "file" + bytes: int + created_at: int + filename: str + purpose: str + status: str = "processed" + status_details: Optional[str] = None + expires_at: Optional[int] = None + + +class BatchObject(BaseModel): + id: str + object: str = "batch" + endpoint: str + errors: Optional[Dict] = None + input_file_id: str + completion_window: str + status: str + output_file_id: Optional[str] = None + error_file_id: Optional[str] = None + created_at: int + in_progress_at: Optional[int] = None + expires_at: Optional[int] = None + finalizing_at: Optional[int] = None + completed_at: Optional[int] = None + failed_at: Optional[int] = None + expired_at: Optional[int] = None + cancelling_at: Optional[int] = None + cancelled_at: Optional[int] = None + request_counts: Optional[Dict[str, int]] = None + metadata: Optional[Dict] = None + + +class BatchListResponse(BaseModel): + object: str = "list" + data: List[Dict] + first_id: Optional[str] = None + last_id: Optional[str] = None + has_more: bool = False + + +file_storage: Dict[str, Dict] = {} +batch_storage: Dict[str, BatchObject] = {} +batch_results: Dict[str, List[Dict]] = {} + +PROCESSING_DELAY_SECONDS = float(1) +VALIDATING_DELAY_SECONDS = float(3) + + +async def process_batch(batch_id: str): + logger.info(f"Starting batch processing for {batch_id}") + try: + batch = batch_storage[batch_id] + + await asyncio.sleep(VALIDATING_DELAY_SECONDS) + batch.status = "in_progress" + batch.in_progress_at = int(time.time()) + logger.info(f"Batch {batch_id} status: in_progress") + + await process_batch_requests(batch_id) + await asyncio.sleep(PROCESSING_DELAY_SECONDS) + + batch.status = "finalizing" + batch.finalizing_at = int(time.time()) + logger.info(f"Batch {batch_id} status: finalizing") + await asyncio.sleep(PROCESSING_DELAY_SECONDS) + + await create_output_file(batch_id) + + batch.status = "completed" + batch.completed_at = int(time.time()) + logger.info(f"Batch {batch_id} status: completed") + + except Exception as e: + logger.error(f"Batch {batch_id} failed: {e}") + batch = batch_storage[batch_id] + batch.status = "failed" + batch.failed_at = int(time.time()) + batch.errors = { + "object": "list", + "data": [{"code": "processing_error", "message": str(e)}], + } + + +async def process_batch_requests(batch_id: str): + batch = batch_storage[batch_id] + input_file = file_storage[batch.input_file_id] + + requests = [] + for line in input_file["content"].split("\n"): + if line.strip(): + try: + requests.append(json.loads(line)) + except json.JSONDecodeError as e: + logger.warning(f"Invalid JSON line in batch {batch_id}: {e}") + + logger.info(f"Batch {batch_id} has {len(requests)} requests") + + results = [] + failed_count = 0 + for req in requests: + result = await process_single_request(req) + if result.get("error"): + failed_count += 1 + results.append(result) + + batch_results[batch_id] = results + batch.request_counts = { + "total": len(requests), + "completed": len(results) - failed_count, + "failed": failed_count, + } + + +async def process_single_request(request_data: Dict) -> Dict: + custom_id = request_data.get("custom_id") + url = request_data.get("url", "/v1/chat/completions") + body = request_data.get("body", {}) + + if "/chat/completions" in url: + response_body = { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": int(time.time()), + "model": body.get("model", "gpt-4o"), + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Mock batch response."}, + "finish_reason": "stop", + }, + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + status_code = 200 + else: + response_body = {"error": {"message": f"Unsupported endpoint: {url}"}} + status_code = 400 + + return { + "id": f"batch_req_{uuid.uuid4().hex[:12]}", + "custom_id": custom_id, + "response": { + "status_code": status_code, + "request_id": f"req_{uuid.uuid4().hex[:12]}", + "body": response_body, + }, + "error": None, + } + + +async def create_output_file(batch_id: str): + results = batch_results.get(batch_id, []) + output_lines = [json.dumps(result) for result in results] + output_content = "\n".join(output_lines) + + output_file_id = f"file-batch-output-{uuid.uuid4().hex[:12]}" + file_storage[output_file_id] = { + "content": output_content, + "filename": f"batch_output_{batch_id}.jsonl", + "purpose": "batch_output", + "bytes": len(output_content.encode()), + "created_at": int(time.time()), + } + + batch = batch_storage[batch_id] + batch.output_file_id = output_file_id + logger.info(f"Created output file {output_file_id} for batch {batch_id}") + + +def validate_batch_input(content: str) -> tuple[bool, str, List[Dict]]: + requests = [] + custom_ids = set() + + lines = content.strip().split("\n") + if not lines or all(not line.strip() for line in lines): + return False, "empty_batch", [] + + for line_num, line in enumerate(lines, 1): + if not line.strip(): + continue + try: + req = json.loads(line) + except json.JSONDecodeError: + return False, "invalid_json_line", [] + + for field in ["custom_id", "method", "url", "body"]: + if field not in req: + return False, "invalid_request", [] + + if req["custom_id"] in custom_ids: + return False, "duplicate_custom_id", [] + custom_ids.add(req["custom_id"]) + + requests.append(req) + + if len(requests) > 100000: + return False, "too_many_tasks", [] + + return True, "", requests + + +def setup_batch_routes(app: FastAPI): + # Files endpoints (OpenAI and Azure paths) + @app.post("/openai/v1/files") + @app.post("/openai/files") + @app.post("/v1/files") + @app.post("/files") + async def create_file(request: Request): + form = await request.form() + logger.info(f"File upload form fields: {list(form.keys())}") + + file: UploadFile = form.get("file") + purpose: str = form.get("purpose", "batch") + + if not file: + raise HTTPException(status_code=400, detail="No file provided") + + logger.info(f"Uploading file: {file.filename}, purpose: {purpose}") + + content = await file.read() + content_str = content.decode("utf-8") + + file_id = f"file-{uuid.uuid4().hex[:24]}" + created_at = int(time.time()) + + expires_at = None + expires_after_seconds = form.get("expires_after[seconds]") + if expires_after_seconds: + try: + seconds = int(expires_after_seconds) + logger.info(f"expires_after[seconds] = {seconds}") + if seconds < 259200 or seconds > 2592000: + raise HTTPException( + status_code=400, + detail={ + "error": { + "code": "invalidPayload", + "message": "Value for Seconds must be between 259200 and 2592000.", + }, + }, + ) + expires_at = created_at + seconds + logger.info(f"Calculated expires_at: {expires_at}") + except ValueError as e: + logger.warning(f"Failed to parse expires_after[seconds]: {e}") + + file_storage[file_id] = { + "content": content_str, + "filename": file.filename or "batch_input.jsonl", + "purpose": purpose, + "bytes": len(content), + "created_at": created_at, + "expires_at": expires_at, + } + + logger.info(f"Created file {file_id}, expires_at={expires_at}") + return FileObject( + id=file_id, + bytes=len(content), + created_at=created_at, + filename=file.filename or "batch_input.jsonl", + purpose=purpose, + expires_at=expires_at, + ).model_dump() + + @app.get("/openai/v1/files/{file_id}") + @app.get("/openai/files/{file_id}") + @app.get("/v1/files/{file_id}") + @app.get("/files/{file_id}") + async def get_file(file_id: str): + logger.info(f"Getting file: {file_id}") + if file_id not in file_storage: + raise HTTPException(status_code=404, detail="File not found") + + file_data = file_storage[file_id] + return FileObject( + id=file_id, + bytes=file_data["bytes"], + created_at=file_data["created_at"], + filename=file_data["filename"], + purpose=file_data["purpose"], + expires_at=file_data.get("expires_at"), + ).model_dump() + + @app.get("/openai/v1/files/{file_id}/content") + @app.get("/openai/files/{file_id}/content") + @app.get("/v1/files/{file_id}/content") + @app.get("/files/{file_id}/content") + async def get_file_content(file_id: str): + logger.info(f"Getting file content: {file_id}") + if file_id not in file_storage: + raise HTTPException(status_code=404, detail="File not found") + + file_data = file_storage[file_id] + content = file_data["content"] + + return StreamingResponse( + io.StringIO(content), + media_type="application/octet-stream", + headers={ + "Content-Disposition": f"attachment; filename={file_data['filename']}", + }, + ) + + @app.delete("/openai/v1/files/{file_id}") + @app.delete("/openai/files/{file_id}") + @app.delete("/v1/files/{file_id}") + @app.delete("/files/{file_id}") + async def delete_file(file_id: str): + logger.info(f"Deleting file: {file_id}") + if file_id not in file_storage: + raise HTTPException(status_code=404, detail="File not found") + + del file_storage[file_id] + return {"id": file_id, "object": "file", "deleted": True} + + @app.get("/openai/v1/files") + @app.get("/openai/files") + @app.get("/v1/files") + @app.get("/files") + async def list_files( + purpose: Optional[str] = None, + limit: int = Query(10000, le=10000), + ): + logger.info(f"Listing files, purpose: {purpose}, limit: {limit}") + files = [] + for file_id, file_data in file_storage.items(): + if purpose is None or file_data.get("purpose") == purpose: + files.append( + FileObject( + id=file_id, + bytes=file_data["bytes"], + created_at=file_data["created_at"], + filename=file_data["filename"], + purpose=file_data["purpose"], + expires_at=file_data.get("expires_at"), + ).model_dump(), + ) + return {"object": "list", "data": files[:limit]} + + # Batches endpoints (OpenAI and Azure paths) + @app.post("/openai/v1/batches") + @app.post("/openai/batches") + @app.post("/v1/batches") + @app.post("/batches") + async def create_batch(request_data: dict): + input_file_id = request_data.get("input_file_id") + endpoint = request_data.get("endpoint", "/v1/chat/completions") + completion_window = request_data.get("completion_window", "24h") + metadata = request_data.get("metadata", {}) + output_expires_after = request_data.get("output_expires_after") + + logger.info( + f"Creating batch with input_file: {input_file_id}, endpoint: {endpoint}, output_expires_after: {output_expires_after}", + ) + + if not input_file_id or input_file_id not in file_storage: + raise HTTPException(status_code=400, detail="Input file not found") + + input_file = file_storage[input_file_id] + is_valid, error_code, _ = validate_batch_input(input_file["content"]) + if not is_valid: + raise HTTPException( + status_code=400, + detail={ + "error": { + "code": error_code, + "message": f"Validation failed: {error_code}", + }, + }, + ) + + batch_id = f"batch_{uuid.uuid4()}" + created_at = int(time.time()) + + if output_expires_after: + seconds = ( + output_expires_after.get("seconds", 0) + if isinstance(output_expires_after, dict) + else 0 + ) + expires_at = created_at + seconds + logger.info( + f"Using output_expires_after: {seconds}s, expires_at: {expires_at}", + ) + elif completion_window == "24h": + expires_at = created_at + (24 * 60 * 60) + else: + expires_at = created_at + (24 * 60 * 60) + + batch = BatchObject( + id=batch_id, + endpoint=endpoint, + input_file_id=input_file_id, + completion_window=completion_window, + status="validating", + created_at=created_at, + expires_at=expires_at, + request_counts={"total": 0, "completed": 0, "failed": 0}, + metadata=metadata, + ) + + batch_storage[batch_id] = batch + logger.info(f"Created batch {batch_id}") + + asyncio.create_task(process_batch(batch_id)) + + return batch.model_dump() + + @app.get("/openai/v1/batches/{batch_id}") + @app.get("/openai/batches/{batch_id}") + @app.get("/v1/batches/{batch_id}") + @app.get("/batches/{batch_id}") + async def get_batch(batch_id: str): + logger.info(f"Getting batch: {batch_id}") + if batch_id not in batch_storage: + raise HTTPException(status_code=404, detail="Batch not found") + + return batch_storage[batch_id].model_dump() + + @app.get("/openai/v1/batches") + @app.get("/openai/batches") + @app.get("/v1/batches") + @app.get("/batches") + async def list_batches( + after: Optional[str] = Query(None), + limit: int = Query(20, le=100), + ): + logger.info(f"Listing batches, after: {after}, limit: {limit}") + batches = list(batch_storage.values()) + batches.sort(key=lambda x: x.created_at, reverse=True) + + if after: + after_index = next((i for i, b in enumerate(batches) if b.id == after), -1) + if after_index >= 0: + batches = batches[after_index + 1 :] + + batches = batches[:limit] + + return BatchListResponse( + data=[batch.model_dump() for batch in batches], + first_id=batches[0].id if batches else None, + last_id=batches[-1].id if batches else None, + has_more=len(batches) == limit, + ).model_dump() + + @app.post("/openai/v1/batches/{batch_id}/cancel") + @app.post("/openai/batches/{batch_id}/cancel") + @app.post("/v1/batches/{batch_id}/cancel") + @app.post("/batches/{batch_id}/cancel") + async def cancel_batch(batch_id: str): + logger.info(f"Cancelling batch: {batch_id}") + if batch_id not in batch_storage: + raise HTTPException(status_code=404, detail="Batch not found") + + batch = batch_storage[batch_id] + if batch.status in ["completed", "failed", "cancelled", "expired"]: + raise HTTPException( + status_code=400, + detail=f"Cannot cancel batch in {batch.status} status", + ) + + batch.status = "cancelled" + batch.cancelled_at = int(time.time()) + logger.info(f"Batch {batch_id} cancelled") + + return batch.model_dump() + + # Debug endpoints + @app.get("/debug/batches") + async def debug_list_batches(): + return { + "batches": { + batch_id: batch.model_dump() + for batch_id, batch in batch_storage.items() + }, + "files": { + file_id: {k: v for k, v in data.items() if k != "content"} + for file_id, data in file_storage.items() + }, + } + + @app.post("/reset") + @app.post("/debug/clear") + async def reset_all(): + file_storage.clear() + batch_storage.clear() + batch_results.clear() + logger.info("All data cleared") + return {"message": "All data cleared"} + + @app.get("/debug/status") + async def debug_status(): + return { + "files_count": len(file_storage), + "batches_count": len(batch_storage), + "batch_statuses": {bid: b.status for bid, b in batch_storage.items()}, + } diff --git a/tests/proxy_e2e_azure_batches_tests/fixtures/mock_azure_batch_server/mock_chat.py b/tests/proxy_e2e_azure_batches_tests/fixtures/mock_azure_batch_server/mock_chat.py new file mode 100644 index 0000000000..c33523579a --- /dev/null +++ b/tests/proxy_e2e_azure_batches_tests/fixtures/mock_azure_batch_server/mock_chat.py @@ -0,0 +1,124 @@ +import json +import time +import uuid +from datetime import datetime + +from fastapi import FastAPI, Request +from fastapi.responses import StreamingResponse + + +def get_request_details(request: Request, body: dict = None) -> str: + details = { + "method": request.method, + "url": str(request.url), + "path": request.url.path, + "headers": dict(request.headers), + "query_params": dict(request.query_params), + } + return json.dumps(details, indent=2) + + +def data_generator(response_details: str, model: str): + response_id = uuid.uuid4().hex + content = response_details + chunk_size = 50 + for i in range(0, len(content), chunk_size): + text_chunk = content[i : i + chunk_size] + chunk = { + "id": f"chatcmpl-{response_id}", + "object": "chat.completion.chunk", + "created": int(time.time()), + "model": model, + "choices": [{"index": 0, "delta": {"content": text_chunk}}], + } + yield f"data: {json.dumps(chunk)}\n\n" + final_chunk = { + "id": f"chatcmpl-{response_id}", + "object": "chat.completion.chunk", + "created": int(time.time()), + "model": model, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + } + yield f"data: {json.dumps(final_chunk)}\n\n" + yield "data: [DONE]\n\n" + + +def setup_chat_routes(app: FastAPI): + @app.post("/chat/completions") + @app.post("/v1/chat/completions") + @app.post("/openai/deployments/{model:path}/chat/completions") + async def completion(request: Request): + data = await request.json() + model = data.get("model", "unknown") + request_details = get_request_details(request, data) + timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S") + response_details = f"Request:{request_details}, Canned Response:{timestamp}" + + if data.get("stream"): + return StreamingResponse( + content=data_generator(response_details, model), + media_type="text/event-stream", + ) + else: + response_id = uuid.uuid4().hex + response = { + "id": f"chatcmpl-{response_id}", + "object": "chat.completion", + "created": int(time.time()), + "model": model, + "system_fingerprint": "fp_mock_server", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": response_details, + }, + "logprobs": None, + "finish_reason": "stop", + }, + ], + "usage": { + "prompt_tokens": 9, + "completion_tokens": 12, + "total_tokens": 21, + }, + } + return response + + @app.post("/completions") + @app.post("/v1/completions") + async def text_completion(request: Request): + data = await request.json() + model = data.get("model", "unknown") + request_details = get_request_details(request, data) + timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S") + response_details = f"Request:{request_details}, Canned Response:{timestamp}" + + if data.get("stream"): + return StreamingResponse( + content=data_generator(response_details, model), + media_type="text/event-stream", + ) + else: + response = { + "id": f"cmpl-{uuid.uuid4().hex}", + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "logprobs": None, + "text": response_details, + }, + ], + "created": int(time.time()), + "model": model, + "object": "text_completion", + "system_fingerprint": None, + "usage": { + "completion_tokens": 16, + "prompt_tokens": 10, + "total_tokens": 26, + }, + } + return response diff --git a/tests/proxy_e2e_azure_batches_tests/fixtures/mock_azure_batch_server/mock_embeddings.py b/tests/proxy_e2e_azure_batches_tests/fixtures/mock_azure_batch_server/mock_embeddings.py new file mode 100644 index 0000000000..f31b1ad4b8 --- /dev/null +++ b/tests/proxy_e2e_azure_batches_tests/fixtures/mock_azure_batch_server/mock_embeddings.py @@ -0,0 +1,23 @@ +from fastapi import FastAPI, Request + + +def setup_embeddings_routes(app: FastAPI): + @app.post("/embeddings") + @app.post("/v1/embeddings") + @app.post("/openai/deployments/{model:path}/embeddings") + async def embeddings(request: Request): + data = await request.json() + model = data.get("model", "unknown") + _small_embedding = [ + -0.006929283495992422, + -0.005336422007530928, + -4.547132266452536e-05, + -0.024047505110502243, + ] + big_embedding = _small_embedding * 100 + return { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": big_embedding}], + "model": model, + "usage": {"prompt_tokens": 5, "total_tokens": 5}, + } diff --git a/tests/proxy_e2e_azure_batches_tests/fixtures/mock_azure_batch_server/mock_responses.py b/tests/proxy_e2e_azure_batches_tests/fixtures/mock_azure_batch_server/mock_responses.py new file mode 100644 index 0000000000..94cb25794b --- /dev/null +++ b/tests/proxy_e2e_azure_batches_tests/fixtures/mock_azure_batch_server/mock_responses.py @@ -0,0 +1,170 @@ +import json +import re +import time +import uuid +from datetime import datetime + +from typing import Any + +from fastapi import FastAPI, Request, HTTPException + + +# Header to identify which model/deployment this request targets (simulates Azure model-specific encryption). +# When set, the mock validates that encrypted_content in input was produced by this model. +MOCK_AZURE_MODEL_HEADER = "X-Mock-Azure-Model" + +# Prefix we use in mock encrypted_content: gAAA_model__<32hex uuid> +# Model id can contain underscores (e.g. gpt-5.1-codex-openai-2). +ENCRYPTED_CONTENT_MODEL_PREFIX = re.compile(r"^gAAA_model_(.+)_[0-9a-f]{32}$") + + +def _extract_model_from_encrypted_content(encrypted: str) -> str | None: + """Extract model id from our mock encrypted_content format, or None if not our format.""" + if not isinstance(encrypted, str) or not encrypted.startswith("gAAA"): + return None + m = ENCRYPTED_CONTENT_MODEL_PREFIX.match(encrypted) + return m.group(1) if m else None + + +def _collect_encrypted_contents(obj, out: list[str]) -> None: + """Recursively collect all encrypted_content string values from input structure.""" + if isinstance(obj, dict): + if "encrypted_content" in obj and obj["encrypted_content"]: + out.append(obj["encrypted_content"]) + for v in obj.values(): + _collect_encrypted_contents(v, out) + elif isinstance(obj, list): + for item in obj: + _collect_encrypted_contents(item, out) + + +def _validate_encrypted_content_model(request_model: str | None, input_data: Any) -> str | None: + """ + If request_model is set, check that all encrypted_content in input was produced by this model. + Returns error message if validation fails, else None. + Content with our format (gAAA_model__) must match request_model. + """ + if not request_model: + return None + encrypted_values: list[str] = [] + _collect_encrypted_contents(input_data, encrypted_values) + for enc in encrypted_values: + content_model = _extract_model_from_encrypted_content(enc) + if content_model is not None and content_model != request_model: + err = enc[:50] + "..." if len(enc) > 50 else enc + return f"The encrypted content {err} could not be verified." + return None + + +def get_request_details(request: Request, body: dict = None) -> str: + details = { + "method": request.method, + "url": str(request.url), + "path": request.url.path, + "headers": dict(request.headers), + "query_params": dict(request.query_params), + } + return json.dumps(details, indent=2) + + +def setup_responses_routes(app: FastAPI): + @app.post("/responses") + @app.post("/v1/responses") + @app.post("/openai/responses") + async def responses_api(request: Request): + data = await request.json() + model = data.get("model", "unknown") + + # Simulate Azure: encrypted content from one model cannot be verified by another. + input_data = data.get("input") + err_msg = _validate_encrypted_content_model(model, input_data) + if err_msg is not None: + raise HTTPException( + status_code=400, + detail={ + "error": { + "message": err_msg, + "type": "invalid_request_error", + "param": None, + "code": "invalid_encrypted_content", + } + }, + ) + + request_details = get_request_details(request, data) + timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S") + response_details = f"Request:{request_details}, Canned Response:{timestamp}" + response_id = uuid.uuid4().hex + message_id = f"msg_{uuid.uuid4().hex[:34]}" + reasoning_id = f"rs_{uuid.uuid4().hex[:34]}" + + output_items: list[dict[str, Any]] = [ + { + "id": message_id, + "content": [ + { + "annotations": [], + "text": response_details, + "type": "output_text", + "logprobs": [], + }, + ], + "role": "assistant", + "status": "completed", + "type": "message", + }, + ] + + if model: + output_items.append( + { + "id": reasoning_id, + "type": "reasoning", + "status": "completed", + "encrypted_content": f"gAAA_model_{model}_{uuid.uuid4().hex}", + } + ) + + return { + "id": f"resp_{response_id}", + "created_at": int(time.time()), + "error": None, + "incomplete_details": None, + "instructions": None, + "metadata": {}, + "model": model, + "object": "response", + "output": output_items, + "parallel_tool_calls": True, + "temperature": data.get("temperature", 1.0), + "tool_choice": data.get("tool_choice", "auto"), + "tools": data.get("tools", []), + "top_p": data.get("top_p", 1.0), + "max_output_tokens": data.get("max_output_tokens"), + "previous_response_id": None, + "reasoning": {"effort": None, "summary": None}, + "status": "completed", + "text": {"format": {"type": "text"}, "verbosity": "medium"}, + "truncation": "disabled", + "usage": { + "input_tokens": 11, + "input_tokens_details": { + "audio_tokens": None, + "cached_tokens": 0, + "text_tokens": None, + }, + "output_tokens": 19, + "output_tokens_details": {"reasoning_tokens": 0, "text_tokens": None}, + "total_tokens": 30, + "cost": None, + }, + "user": None, + "store": True, + "background": False, + "content_filters": None, + "max_tool_calls": None, + "prompt_cache_key": None, + "safety_identifier": None, + "service_tier": "default", + "top_logprobs": 0, + } diff --git a/tests/proxy_e2e_azure_batches_tests/fixtures/mock_azure_batch_server/mock_s3_callback.py b/tests/proxy_e2e_azure_batches_tests/fixtures/mock_azure_batch_server/mock_s3_callback.py new file mode 100644 index 0000000000..8cc99a75b2 --- /dev/null +++ b/tests/proxy_e2e_azure_batches_tests/fixtures/mock_azure_batch_server/mock_s3_callback.py @@ -0,0 +1,98 @@ +""" +Mock S3 callback receiver for testing LiteLLM S3 callbacks. + +This module provides S3-compatible endpoints that capture callback data +sent by LiteLLM's s3_v2 callback handler after batch completion. +""" + +import json +import logging +import time +from typing import Any, Dict, List, Optional + +from fastapi import FastAPI, Request +from pydantic import BaseModel + +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + + +class S3CallbackRecord(BaseModel): + key: str + bucket: str + content: Dict[str, Any] + timestamp: int + content_type: Optional[str] = None + + +callback_storage: List[S3CallbackRecord] = [] + + +def setup_s3_callback_routes(app: FastAPI): + @app.put("/{bucket}/{key:path}") + async def s3_put_object(bucket: str, key: str, request: Request): + content_type = request.headers.get("content-type", "application/json") + body = await request.body() + + try: + content = json.loads(body.decode("utf-8")) + except (json.JSONDecodeError, UnicodeDecodeError): + content = {"raw": body.decode("utf-8", errors="replace")} + + record = S3CallbackRecord( + key=key, + bucket=bucket, + content=content, + timestamp=int(time.time()), + content_type=content_type, + ) + callback_storage.append(record) + + logger.info(f"S3 callback received: bucket={bucket}, key={key}") + logger.debug(f"Callback content: {json.dumps(content, indent=2)[:500]}") + + return { + "ETag": f'"{hash(body)}"', + "VersionId": None, + } + + @app.get("/mock-s3/callbacks") + async def list_callbacks( + bucket: Optional[str] = None, + key_prefix: Optional[str] = None, + limit: int = 100, + ): + results = callback_storage + + if bucket: + results = [r for r in results if r.bucket == bucket] + + if key_prefix: + results = [r for r in results if r.key.startswith(key_prefix)] + + return { + "count": len(results), + "callbacks": [r.model_dump() for r in results[-limit:]], + } + + @app.get("/mock-s3/callbacks/count") + async def count_callbacks(bucket: Optional[str] = None): + if bucket: + count = sum(1 for r in callback_storage if r.bucket == bucket) + else: + count = len(callback_storage) + + return {"count": count} + + @app.get("/mock-s3/callbacks/latest") + async def get_latest_callback(): + if not callback_storage: + return {"callback": None} + return {"callback": callback_storage[-1].model_dump()} + + @app.delete("/mock-s3/callbacks") + async def clear_callbacks(): + count = len(callback_storage) + callback_storage.clear() + logger.info(f"Cleared {count} S3 callbacks") + return {"cleared": count} diff --git a/tests/proxy_e2e_azure_batches_tests/fixtures/mock_azure_batch_server/server.py b/tests/proxy_e2e_azure_batches_tests/fixtures/mock_azure_batch_server/server.py new file mode 100644 index 0000000000..a0bda6a186 --- /dev/null +++ b/tests/proxy_e2e_azure_batches_tests/fixtures/mock_azure_batch_server/server.py @@ -0,0 +1,33 @@ +from fastapi import FastAPI, Request +from fastapi.middleware.cors import CORSMiddleware + +from .mock_azure_batch import setup_batch_routes +from .mock_chat import setup_chat_routes +from .mock_embeddings import setup_embeddings_routes +from .mock_responses import setup_responses_routes +from .mock_s3_callback import setup_s3_callback_routes + + +def create_mock_azure_batch_server() -> FastAPI: + """Create a FastAPI app that mocks Azure Batch API and S3 callbacks.""" + app = FastAPI() + + app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], + ) + + @app.get("/health") + async def health(): + return {"status": "ok"} + + setup_chat_routes(app) + setup_responses_routes(app) + setup_embeddings_routes(app) + setup_batch_routes(app) + setup_s3_callback_routes(app) + + return app diff --git a/tests/proxy_e2e_azure_batches_tests/fixtures/run_mock_server.py b/tests/proxy_e2e_azure_batches_tests/fixtures/run_mock_server.py new file mode 100644 index 0000000000..8804c47b7d --- /dev/null +++ b/tests/proxy_e2e_azure_batches_tests/fixtures/run_mock_server.py @@ -0,0 +1,12 @@ + +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent.parent)) + +from fixtures.mock_azure_batch_server import create_mock_azure_batch_server +import uvicorn + +if __name__ == "__main__": + app = create_mock_azure_batch_server() + uvicorn.run(app, host="0.0.0.0", port=8090, log_level="info", access_log=False) diff --git a/tests/proxy_e2e_azure_batches_tests/test_fixtures_smoke.py b/tests/proxy_e2e_azure_batches_tests/test_fixtures_smoke.py new file mode 100644 index 0000000000..eeb1796371 --- /dev/null +++ b/tests/proxy_e2e_azure_batches_tests/test_fixtures_smoke.py @@ -0,0 +1,41 @@ +""" +Smoke test to verify fixtures start and stop correctly. +Run this first to ensure the infrastructure works before running full E2E tests. +""" + +import httpx +import pytest + + +pytestmark = pytest.mark.usefixtures("mock_azure_server", "litellm_proxy_server") + + +def test_mock_server_health(mock_azure_server): + """Verify mock Azure server is running and healthy.""" + response = httpx.get(f"{mock_azure_server}/health", timeout=5.0) + assert response.status_code == 200 + assert response.json() == {"status": "ok"} + print(f"✓ Mock Azure server is healthy at {mock_azure_server}") + + +def test_litellm_proxy_health(litellm_proxy_server): + """Verify LiteLLM proxy is running and healthy.""" + response = httpx.get(f"{litellm_proxy_server}/health", timeout=5.0) + assert response.status_code == 200 + print(f"✓ LiteLLM proxy is healthy at {litellm_proxy_server}") + + +def test_litellm_proxy_model_list(litellm_proxy_server): + """Verify LiteLLM proxy can list models.""" + response = httpx.get( + f"{litellm_proxy_server}/v1/models", + headers={"Authorization": "Bearer sk-1234"}, + timeout=5.0, + ) + assert response.status_code == 200 + data = response.json() + assert "data" in data + models = [m["id"] for m in data["data"]] + print(f"✓ LiteLLM proxy has {len(models)} models configured") + assert "azure-fake-gpt-5-batch-2025-08-07" in models + print(f"✓ Azure batch model is configured") diff --git a/tests/proxy_e2e_azure_batches_tests/test_managed_files_base.py b/tests/proxy_e2e_azure_batches_tests/test_managed_files_base.py new file mode 100644 index 0000000000..79e7e58f39 --- /dev/null +++ b/tests/proxy_e2e_azure_batches_tests/test_managed_files_base.py @@ -0,0 +1,1085 @@ +"""Base class for managed files and batch API tests.""" + +import json +import os +import sys +import time +from datetime import datetime +from typing import Optional +from urllib.parse import urlparse + +import httpx +import openai +import psycopg2 +import pytest +from tenacity import Retrying, stop_after_delay, wait_fixed + +sys.path.insert(0, os.path.abspath("../..")) + +from base_integration_test import ( + BaseLiteLLMIntegrationTest, + get_mock_server_base_url, + use_mock_models, +) + + +class ManagedFilesState: + """Query and pretty print the state of managed files and objects tables.""" + + def __init__(self, database_url: Optional[str] = None): + self.database_url = database_url or os.environ.get("DATABASE_URL") + if not self.database_url: + raise ValueError("DATABASE_URL not provided and not in environment") + + def _get_connection(self): + parsed = urlparse(self.database_url) + return psycopg2.connect( + host=parsed.hostname, + port=parsed.port or 5432, + user=parsed.username, + password=parsed.password, + dbname=parsed.path.lstrip("/"), + ) + + def _shorten_id(self, id_str: str, max_len: int = 24) -> str: + if id_str is None: + return "None" + if len(id_str) <= max_len: + return id_str + return id_str[:10] + "..." + id_str[-10:] + + def _format_timestamp(self, ts) -> str: + if ts is None: + return "None" + if isinstance(ts, datetime): + return ts.strftime("%Y-%m-%d %H:%M:%S") + return str(ts) + + def get_managed_files(self, limit: int = 20) -> list: + query = """ + SELECT unified_file_id, file_purpose, created_by, created_at, + updated_at, model_mappings, storage_backend + FROM "LiteLLM_ManagedFileTable" + ORDER BY created_at DESC + LIMIT %s + """ + with self._get_connection() as conn: + with conn.cursor() as cur: + cur.execute(query, (limit,)) + columns = [desc[0] for desc in cur.description] + return [dict(zip(columns, row)) for row in cur.fetchall()] + + def get_managed_objects( + self, + limit: int = 20, + status: Optional[str] = None, + ) -> list: + query = """ + SELECT id, unified_object_id, status, file_purpose, + created_by, created_at, updated_at + FROM "LiteLLM_ManagedObjectTable" + """ + params = [] + if status: + query += " WHERE status = %s" + params.append(status) + query += " ORDER BY created_at DESC LIMIT %s" + params.append(limit) + + with self._get_connection() as conn: + with conn.cursor() as cur: + cur.execute(query, params) + columns = [desc[0] for desc in cur.description] + return [dict(zip(columns, row)) for row in cur.fetchall()] + + def print_managed_files(self, limit: int = 20): + files = self.get_managed_files(limit) + print(f"\n{'=' * 80}") + print(f"MANAGED FILES TABLE ({len(files)} rows)") + print(f"{'=' * 80}") + + if not files: + print(" (no rows)") + return + + for i, f in enumerate(files, 1): + print(f"\n[{i}] unified_file_id: {self._shorten_id(f['unified_file_id'])}") + print(f" purpose: {f['file_purpose']}") + print(f" created_by: {f['created_by']}") + print(f" created_at: {self._format_timestamp(f['created_at'])}") + print(f" storage_backend: {f.get('storage_backend', 'None')}") + if f.get("model_mappings"): + mappings = f["model_mappings"] + if isinstance(mappings, dict): + print(f" model_mappings: {len(mappings)} model(s)") + for model_id, file_id in list(mappings.items())[:3]: + print( + f" - {self._shorten_id(model_id)}: {self._shorten_id(file_id)}", + ) + if len(mappings) > 3: + print(f" ... and {len(mappings) - 3} more") + + def print_managed_objects(self, limit: int = 20, status: Optional[str] = None): + """Pretty print the managed objects table.""" + objects = self.get_managed_objects(limit, status) + status_filter = f" (status={status})" if status else "" + print(f"\n{'=' * 80}") + print(f"MANAGED OBJECTS TABLE{status_filter} ({len(objects)} rows)") + print(f"{'=' * 80}") + + if not objects: + print(" (no rows)") + return + + for i, o in enumerate(objects, 1): + print(f"\n[{i}] id: {o['id']}") + print(f" unified_object_id: {self._shorten_id(o['unified_object_id'])}") + print(f" status: {o['status']}") + print(f" file_purpose: {o['file_purpose']}") + print(f" created_by: {o['created_by']}") + print(f" created_at: {self._format_timestamp(o['created_at'])}") + + def print_validating_batches(self): + """Print batches that are stuck in validating state.""" + self.print_managed_objects(status="validating") + + def print_all(self, limit: int = 10): + """Print both tables.""" + self.print_managed_files(limit) + self.print_managed_objects(limit) + + def count_by_status(self) -> dict: + """Count managed objects by status.""" + query = """ + SELECT status, COUNT(*) as count + FROM "LiteLLM_ManagedObjectTable" + GROUP BY status + ORDER BY count DESC + """ + with self._get_connection() as conn: + with conn.cursor() as cur: + cur.execute(query) + return {row[0]: row[1] for row in cur.fetchall()} + + def print_summary(self): + """Print a summary of table states.""" + print(f"\n{'=' * 80}") + print("DATABASE STATE SUMMARY") + print(f"{'=' * 80}") + + with self._get_connection() as conn: + with conn.cursor() as cur: + cur.execute('SELECT COUNT(*) FROM "LiteLLM_ManagedFileTable"') + file_count = cur.fetchone()[0] + + cur.execute('SELECT COUNT(*) FROM "LiteLLM_ManagedObjectTable"') + object_count = cur.fetchone()[0] + + print(f"\nManaged Files: {file_count} total") + print(f"Managed Objects: {object_count} total") + + status_counts = self.count_by_status() + if status_counts: + print("\nObjects by status:") + for status, count in status_counts.items(): + print(f" - {status}: {count}") + + def get_file_by_unified_id(self, unified_file_id: str) -> Optional[dict]: + """Get a managed file by its unified file ID.""" + query = """ + SELECT unified_file_id, file_object, created_by, created_at, + updated_at, model_mappings, storage_backend + FROM "LiteLLM_ManagedFileTable" + WHERE unified_file_id = %s + """ + with self._get_connection() as conn: + with conn.cursor() as cur: + cur.execute(query, (unified_file_id,)) + row = cur.fetchone() + if row: + columns = [desc[0] for desc in cur.description] + return dict(zip(columns, row)) + return None + + def get_batch_by_unified_id(self, unified_object_id: str) -> Optional[dict]: + """Get a managed batch/object by its unified object ID.""" + query = """ + SELECT id, unified_object_id, model_object_id, status, file_purpose, + created_by, created_at, updated_at + FROM "LiteLLM_ManagedObjectTable" + WHERE unified_object_id = %s + """ + with self._get_connection() as conn: + with conn.cursor() as cur: + cur.execute(query, (unified_object_id,)) + row = cur.fetchone() + if row: + columns = [desc[0] for desc in cur.description] + return dict(zip(columns, row)) + return None + + def get_batch_by_id(self, batch_id: int) -> Optional[dict]: + """Get a managed batch/object by its integer ID.""" + query = """ + SELECT id, unified_object_id, status, file_purpose, + created_by, created_at, updated_at + FROM "LiteLLM_ManagedObjectTable" + WHERE id = %s + """ + with self._get_connection() as conn: + with conn.cursor() as cur: + cur.execute(query, (batch_id,)) + row = cur.fetchone() + if row: + columns = [desc[0] for desc in cur.description] + return dict(zip(columns, row)) + return None + + +MIN_EXPIRY_SECONDS = 259200 + + +class _BaseSubTracker: + """Shared helpers for sub-trackers.""" + + def _shorten_id(self, id_str: str, max_len: int = 20) -> str: + if id_str is None: + return "None" + if len(id_str) <= max_len: + return id_str + return id_str[:8] + "..." + id_str[-8:] + + def _format_timestamp(self, ts) -> str: + if ts is None: + return "None" + if isinstance(ts, datetime): + return ts.strftime("%H:%M:%S") + if isinstance(ts, int): + return datetime.fromtimestamp(ts).strftime("%H:%M:%S") + return str(ts) + + +class BatchDbStateTracker(_BaseSubTracker): + """Tracks batch/file state in the LiteLLM database.""" + + def __init__(self, db_state: ManagedFilesState): + self.db_state = db_state + + def get_file_state(self, file_id: str) -> Optional[dict]: + return self.db_state.get_file_by_unified_id(file_id) + + def get_batch_state(self, batch_id: str) -> Optional[dict]: + return self.db_state.get_batch_by_unified_id(batch_id) + + def format_file_lines(self, file_id: str) -> tuple[str, list[str]]: + """Return (header, detail_lines) for the DB file state.""" + db_file = self.get_file_state(file_id) + header_id = ( + self._shorten_id(db_file.get("unified_file_id")) if db_file else "N/A" + ) + header = f"FILE (DB): {header_id}" + + if not db_file: + return header, [" (not found in DB)"] + + file_obj = db_file.get("file_object") or {} + if isinstance(file_obj, str): + try: + file_obj = json.loads(file_obj) + except Exception: + file_obj = {} + lines = [ + f" purpose: {file_obj.get('purpose', 'N/A')}", + f" storage: {db_file.get('storage_backend', 'N/A')}", + f" created: {self._format_timestamp(db_file.get('created_at'))}", + f" updated: {self._format_timestamp(db_file.get('updated_at'))}", + ] + mappings = db_file.get("model_mappings") + if mappings and isinstance(mappings, dict): + lines.append(f" mappings: {len(mappings)} model(s)") + return header, lines + + def format_batch_lines(self, batch_id: str) -> tuple[str, list[str]]: + """Return (header, detail_lines) for the DB batch state.""" + db_batch = self.get_batch_state(batch_id) + header_id = ( + self._shorten_id(db_batch.get("unified_object_id")) if db_batch else "N/A" + ) + header = f"BATCH (DB): {header_id}" + + if not db_batch: + return header, [" (not found in DB)"] + + lines = [ + f" status: {db_batch.get('status', 'N/A')}", + f" purpose: {db_batch.get('file_purpose', 'N/A')}", + f" created: {self._format_timestamp(db_batch.get('created_at'))}", + f" updated: {self._format_timestamp(db_batch.get('updated_at'))}", + ] + return header, lines + + +class BatchProviderStateTracker(_BaseSubTracker): + """Tracks batch/file state as reported by the LLM provider (via OpenAI client).""" + + def __init__(self, openai_client: openai.OpenAI): + self.client = openai_client + + def get_file_state(self, file_id: str) -> Optional[dict]: + try: + file_obj = self.client.files.retrieve(file_id) + return { + "id": file_obj.id, + "status": file_obj.status, + "purpose": file_obj.purpose, + "bytes": file_obj.bytes, + "filename": file_obj.filename, + "created_at": file_obj.created_at, + "expires_at": file_obj.expires_at, + } + except Exception as e: + return {"error": str(e)} + + def get_batch_state(self, batch_id: str) -> Optional[dict]: + try: + batch = self.client.batches.retrieve(batch_id) + return { + "id": batch.id, + "status": batch.status, + "input_file_id": batch.input_file_id, + "output_file_id": batch.output_file_id, + "error_file_id": batch.error_file_id, + "created_at": batch.created_at, + "completed_at": batch.completed_at, + "request_counts": batch.request_counts, + } + except Exception as e: + return {"error": str(e)} + + def format_file_lines( + self, + file_id: str, + db_state: Optional[BatchDbStateTracker] = None, + ) -> tuple[str, list[str]]: + """Return (header, detail_lines) for the provider file state.""" + raw_file_id = "N/A" + if db_state: + db_file = db_state.get_file_state(file_id) + if db_file: + mappings = db_file.get("model_mappings") + if mappings and isinstance(mappings, dict) and mappings: + first_file_id = next(iter(mappings.values()), None) + raw_file_id = ( + self._shorten_id(first_file_id) if first_file_id else "N/A" + ) + header = f"FILE (RAW): {raw_file_id}" + + provider_file = self.get_file_state(file_id) + if provider_file and "error" not in provider_file: + lines = [ + f" status: {provider_file.get('status', 'N/A')}", + f" purpose: {provider_file.get('purpose', 'N/A')}", + f" bytes: {provider_file.get('bytes', 0)}", + f" created: {self._format_timestamp(provider_file.get('created_at'))}", + f" expires: {self._format_timestamp(provider_file.get('expires_at'))}", + ] + elif provider_file and "error" in provider_file: + lines = [f" ERROR: {provider_file['error'][:35]}"] + else: + lines = [" (not found)"] + return header, lines + + def format_batch_lines( + self, + batch_id: str, + db_state: Optional[BatchDbStateTracker] = None, + ) -> tuple[str, list[str]]: + """Return (header, detail_lines) for the provider batch state.""" + raw_prov_id = "N/A" + if db_state: + db_batch = db_state.get_batch_state(batch_id) + if db_batch: + raw_prov_id = self._shorten_id(db_batch.get("model_object_id")) + header = f"BATCH (RAW): {raw_prov_id}" + + provider_batch = self.get_batch_state(batch_id) + if provider_batch and "error" not in provider_batch: + lines = [ + f" status: {provider_batch.get('status', 'N/A')}", + f" input: {self._shorten_id(provider_batch.get('input_file_id'))}", + f" output: {self._shorten_id(provider_batch.get('output_file_id'))}", + f" created: {self._format_timestamp(provider_batch.get('created_at'))}", + f" completed: {self._format_timestamp(provider_batch.get('completed_at'))}", + ] + req_counts = provider_batch.get("request_counts") + if req_counts: + lines.append( + f" requests: {req_counts.total} total, {req_counts.completed} done", + ) + elif provider_batch and "error" in provider_batch: + lines = [f" ERROR: {provider_batch['error'][:35]}"] + else: + lines = [" (not found)"] + return header, lines + + +class BatchS3StateTracker(_BaseSubTracker): + """Tracks S3 callback state from the mock S3 server.""" + + def __init__(self, mock_server_base_url: str): + self.mock_server_base_url = mock_server_base_url + + def get_callbacks(self, limit: int = 100) -> list[dict]: + try: + response = httpx.get( + f"{self.mock_server_base_url}/mock-s3/callbacks", + params={"limit": limit}, + timeout=5, + ) + if response.status_code == 200: + return response.json().get("callbacks", []) + return [] + except Exception: + return [] + + def get_batch_callbacks(self) -> list[dict]: + """Return only callbacks related to batch operations.""" + batch_call_types = { + "acreate_batch", + "aretrieve_batch", + "acreate_file", + "afile_content", + } + return [ + cb + for cb in self.get_callbacks() + if cb.get("content", {}).get("call_type", "") in batch_call_types + ] + + def get_cost_callbacks(self) -> list[dict]: + """Return CheckBatchCost callbacks (aretrieve_batch with no user_api_key_hash).""" + result = [] + for cb in self.get_callbacks(): + content = cb.get("content", {}) + if content.get("call_type") != "aretrieve_batch": + continue + metadata = content.get("metadata") or {} + if metadata.get("user_api_key_hash") is None: + result.append(cb) + return result + + def format_batch_lines(self, batch_id: str) -> tuple[str, list[str]]: + """Return (header, detail_lines) summarising S3 callback state for this batch.""" + all_cbs = self.get_callbacks() + batch_cbs = self.get_batch_callbacks() + cost_cbs = self.get_cost_callbacks() + + header = f"S3 CALLBACKS: {len(all_cbs)} total" + lines = [ + f" batch-related: {len(batch_cbs)}", + f" cost events: {len(cost_cbs)}", + ] + + # Summarise call_type breakdown for batch callbacks + type_counts: dict[str, int] = {} + for cb in batch_cbs: + ct = cb.get("content", {}).get("call_type", "unknown") + type_counts[ct] = type_counts.get(ct, 0) + 1 + for ct, count in sorted(type_counts.items()): + lines.append(f" {ct}: {count}") + + # Show cost info from the latest cost callback (if any) + if cost_cbs: + latest = cost_cbs[-1].get("content", {}) + lines.append(f" latest cost event:") + lines.append(f" model: {latest.get('model', 'N/A')}") + lines.append(f" response_cost: {latest.get('response_cost', 'N/A')}") + lines.append(f" total_tokens: {latest.get('total_tokens', 0)}") + + return header, lines + + def print_all_callbacks(self): + """Print every S3 callback object in detail, ordered by S3 key timestamp.""" + callbacks = self.get_callbacks() + + # Sort by the timestamp embedded in the S3 key (e.g. "2026-02-15/time-13-01-31-269789_...") + callbacks.sort(key=lambda cb: cb.get("key", "")) + + print(f"\n{'=' * 90}") + print( + f"S3 CALLBACK DETAIL — {len(callbacks)} object(s), ordered by received time", + ) + print(f"{'=' * 90}") + + if not callbacks: + print(" (no callbacks)") + return + + for i, cb in enumerate(callbacks, 1): + content = cb.get("content", {}) + metadata = content.get("metadata") or {} + hidden = content.get("hidden_params") or {} + + print(f"\n[{i}] call_type: {content.get('call_type', 'N/A')}") + print( + f" s3_received_at: {cb.get('received_at', cb.get('timestamp', 'N/A'))}", + ) + print(f" id: {self._shorten_id(content.get('id', ''))}") + print(f" model: {content.get('model', 'N/A')}") + print(f" status: {content.get('status', 'N/A')}") + print(f" response_cost: {content.get('response_cost', 'N/A')}") + print(f" total_tokens: {content.get('total_tokens', 0)}") + print(f" prompt_tokens: {content.get('prompt_tokens', 0)}") + print(f" completion_tokens: {content.get('completion_tokens', 0)}") + print( + f" custom_llm_provider: {content.get('custom_llm_provider', 'N/A')}", + ) + print(f" api_base: {self._shorten_id(content.get('api_base', ''), 40)}") + print(f" cache_hit: {content.get('cache_hit', 'N/A')}") + + print(f" metadata:") + print( + f" user_api_key_hash: {self._shorten_id(metadata.get('user_api_key_hash', 'None'))}", + ) + print( + f" user_api_key_alias: {metadata.get('user_api_key_alias', 'None')}", + ) + print( + f" user_api_key_team_id: {metadata.get('user_api_key_team_id', 'None')}", + ) + print( + f" user_api_key_team_alias: {metadata.get('user_api_key_team_alias', 'None')}", + ) + print( + f" user_api_key_user_id: {metadata.get('user_api_key_user_id', 'None')}", + ) + + batch_models = hidden.get("batch_models") + if batch_models: + print(f" batch_models: {batch_models}") + + response = content.get("response") or {} + if isinstance(response, dict) and response.get("status"): + print(f" response.status: {response.get('status')}") + req_counts = response.get("request_counts") or {} + if req_counts: + print( + f" response.request_counts: total={req_counts.get('total', 0)}, completed={req_counts.get('completed', 0)}, failed={req_counts.get('failed', 0)}", + ) + out_file = response.get("output_file_id") + if out_file: + print(f" response.output_file_id: {self._shorten_id(out_file)}") + + s3_key = cb.get("key", "") + if s3_key: + print(f" s3_key: {s3_key}") + + print(f"\n{'=' * 90}\n") + + +class NoOpStateTracker: + """No-op tracker used when state tracking is disabled.""" + + def set_file_id(self, file_id: str): + pass + + def set_batch_id(self, batch_id: str): + pass + + def print_state(self, step_name: str): + pass + + def wait_and_print_s3_callbacks(self): + pass + + def assert_batch_cost_callback(self): + pass + + +class StateTracker: + """Tracks and prints DB, Provider, and S3 state after each step.""" + + def __init__( + self, + db_tracker: BatchDbStateTracker, + provider_tracker: BatchProviderStateTracker, + s3_tracker: Optional[BatchS3StateTracker] = None, + ): + self.db_tracker = db_tracker + self.provider_tracker = provider_tracker + self.s3_tracker = s3_tracker + self.current_file_id: Optional[str] = None + self.current_batch_id: Optional[str] = None + self.step_number = 0 + + def set_file_id(self, file_id: str): + """Set the file ID to track.""" + self.current_file_id = file_id + + def set_batch_id(self, batch_id: str): + """Set the batch ID to track.""" + self.current_batch_id = batch_id + + def print_state(self, step_name: str): + """Print DB, provider, and S3 state for tracked file and batch.""" + self.step_number += 1 + has_s3 = self.s3_tracker is not None + col_width = 40 + num_cols = 3 if has_s3 else 2 + total_width = (col_width + 3) * num_cols + + print(f"\n{'─' * total_width}") + print(f"│ STEP {self.step_number}: {step_name}") + print(f"{'─' * total_width}") + + col_headers = [ + f"{'DATABASE STATE':<{col_width}}", + f"{'PROVIDER STATE':<{col_width}}", + ] + if has_s3: + col_headers.append(f"{'S3 STATE':<{col_width}}") + print("│ " + " │ ".join(col_headers)) + print(f"{'─' * total_width}") + + if self.current_file_id: + self._print_file_state(col_width, has_s3) + + if self.current_batch_id: + self._print_batch_state(col_width, has_s3) + + print(f"{'─' * total_width}\n") + + def _has_completed_batch_cost_callback(self) -> bool: + """Check if an aretrieve_batch callback with completed status and cost>0 exists.""" + for cb in self.s3_tracker.get_callbacks(): + content = cb.get("content", {}) + if content.get("call_type") != "aretrieve_batch": + continue + response = content.get("response") or {} + if not isinstance(response, dict) or response.get("status") != "completed": + continue + cost = content.get("response_cost", 0) + if cost and cost > 0: + return True + return False + + def wait_and_print_s3_callbacks(self): + """Wait for the S3 v2 logger to flush, then print all callbacks in detail. + + Waits until the cost callback arrives or max_wait is reached. + After detecting the cost callback, waits one extra flush interval + for the proxy to finalize batch_processed before returning. + """ + if not self.s3_tracker: + return + + s3_flush_interval = int(os.environ.get("DEFAULT_S3_FLUSH_INTERVAL_SECONDS", 10)) + batch_poll_interval = int(os.environ.get("PROXY_BATCH_POLLING_INTERVAL", 10)) + max_wait = batch_poll_interval * 3 + s3_flush_interval * 5 + prev_count = len(self.s3_tracker.get_callbacks()) + waited = 0 + cost_detected = False + while waited < max_wait: + print( + f"Waiting for {s3_flush_interval} secs for S3 callbacks to be flushed", + ) + time.sleep(s3_flush_interval) + waited += s3_flush_interval + curr_count = len(self.s3_tracker.get_callbacks()) + print( + f"[S3 flush wait] {waited}s/{max_wait}s — " + f"callbacks: {prev_count} → {curr_count}", + ) + prev_count = curr_count + + if not cost_detected and self._has_completed_batch_cost_callback(): + print( + "Cost callback detected — waiting one more interval " + "for batch_processed finalization" + ) + cost_detected = True + elif cost_detected: + break + + self.s3_tracker.print_all_callbacks() + + def assert_batch_cost_callback(self): + """Assert that a completed-batch S3 callback with non-zero cost exists.""" + if not self.s3_tracker: + return + + callbacks = self.s3_tracker.get_callbacks() + valid_callbacks = [] + for cb in callbacks: + content = cb.get("content", {}) + if content.get("call_type") != "aretrieve_batch": + continue + response = content.get("response") or {} + if not isinstance(response, dict) or response.get("status") != "completed": + continue + cost = content.get("response_cost", 0) + if cost and cost > 0: + valid_callbacks.append(cb) + + if len(valid_callbacks) != 1: + print( + f"\n❌ Assertion failed: Found {len(valid_callbacks)} valid callbacks (expected 1)", + ) + print( + "\nAll valid callbacks with call_type=aretrieve_batch, status=completed, cost>0:", + ) + for idx, cb in enumerate(valid_callbacks, 1): + content = cb.get("content", {}) + print(f"\n[{idx}] Callback:") + print(f" id: {content.get('id', 'N/A')}") + print(f" response_cost: {content.get('response_cost', 0)}") + print(f" litellm_call_id: {content.get('litellm_call_id', 'N/A')}") + response = content.get("response", {}) + print(f" response.id: {response.get('id', 'N/A')}") + print(f" response.status: {response.get('status', 'N/A')}") + metadata = content.get("metadata", {}) + print( + f" user_api_key_user_id: {metadata.get('user_api_key_user_id', 'N/A')}", + ) + print( + f" user_api_key_alias: {metadata.get('user_api_key_alias', 'N/A')}", + ) + print( + f" user_api_key_hash: {metadata.get('user_api_key_hash', 'N/A')}", + ) + print(f" source: {metadata.get('source', 'NOT SET')}") + raise AssertionError( + f"Expected 1 valid callback with call_type=aretrieve_batch, " + f"response.status=completed, and response_cost > 0. " + f"Found {len(valid_callbacks)} valid callbacks.", + ) + + valid_callback = valid_callbacks[0] + callback_user_alias = ( + valid_callback.get("content", {}) + .get("metadata", {}) + .get("user_api_key_alias") + ) + if not callback_user_alias: + raise AssertionError( + f"Expected user_api_key_alias to be set. Found {callback_user_alias}.", + ) + + if callback_user_alias == "default_user_alias": + raise AssertionError( + f"Expected user_api_key_alias to be set to the user who created the batch. " + f"Expected user_api_key_alias to be 'default_user_alias'. " + f"Found {callback_user_alias}.", + ) + + def _print_columns(self, columns: list[list[str]], col_width: int): + """Print multiple columns side-by-side.""" + max_lines = max(len(col) for col in columns) + for i in range(max_lines): + parts = [] + for col in columns: + line = col[i] if i < len(col) else "" + parts.append(f"{line:<{col_width}}") + print("│ " + " │ ".join(parts)) + + def _print_file_state(self, col_width: int, has_s3: bool): + db_header, db_lines = self.db_tracker.format_file_lines(self.current_file_id) + prov_header, prov_lines = self.provider_tracker.format_file_lines( + self.current_file_id, + db_state=self.db_tracker, + ) + + headers = [db_header, prov_header] + columns = [db_lines, prov_lines] + if has_s3: + headers.append("") + columns.append([]) + + header_parts = [f"{h:<{col_width}}" for h in headers] + print("│ " + " │ ".join(header_parts)) + self._print_columns(columns, col_width) + + def _print_batch_state(self, col_width: int, has_s3: bool): + db_header, db_lines = self.db_tracker.format_batch_lines(self.current_batch_id) + prov_header, prov_lines = self.provider_tracker.format_batch_lines( + self.current_batch_id, + db_state=self.db_tracker, + ) + + headers = [db_header, prov_header] + columns = [db_lines, prov_lines] + if has_s3: + s3_header, s3_lines = self.s3_tracker.format_batch_lines( + self.current_batch_id, + ) + headers.append(s3_header) + columns.append(s3_lines) + + # blank separator row + blank = [f"{'':<{col_width}}"] * len(headers) + print("│ " + " │ ".join(blank)) + + header_parts = [f"{h:<{col_width}}" for h in headers] + print("│ " + " │ ".join(header_parts)) + self._print_columns(columns, col_width) + + +def get_batch_model_names(): + if use_mock_models(): + return [ + "azure-fake-gpt-5-batch-2025-08-07", + ] + return [ + "gpt-5-batch-2025-08-07", + ] + + +class ManagedFilesBase(BaseLiteLLMIntegrationTest): + """Base class with shared helpers for managed files and batch tests.""" + + @pytest.fixture(autouse=True) + def setup_test(self, request): + print( + f"Base URL: {self.base_url}, Using mock models: {use_mock_models()}\n", + ) + + def create_state_tracker(self) -> "StateTracker | NoOpStateTracker": + """Create a StateTracker for observing DB, Provider, and S3 state. + + Returns a NoOpStateTracker if USE_STATE_TRACKER is not 'true' or + if DATABASE_URL is not set. + """ + use_tracker = os.environ.get("USE_STATE_TRACKER", "").lower() == "true" + if not use_tracker: + return NoOpStateTracker() + + database_url = os.environ.get("DATABASE_URL") + if not database_url: + print("Warning: DATABASE_URL not set, state tracking disabled") + return NoOpStateTracker() + try: + db_state = ManagedFilesState(database_url) + db_tracker = BatchDbStateTracker(db_state) + provider_tracker = BatchProviderStateTracker(self.openai_client) + + s3_tracker = None + try: + mock_url = get_mock_server_base_url() + s3_tracker = BatchS3StateTracker(mock_url) + except Exception: + pass + + return StateTracker(db_tracker, provider_tracker, s3_tracker) + except Exception as e: + print(f"Warning: Could not create state tracker: {e}") + return NoOpStateTracker() + + def create_openai_client_with_key(self, api_key: str) -> openai.OpenAI: + """Create an OpenAI client with a specific API key.""" + return openai.OpenAI( + base_url=self.base_url, + api_key=api_key, + http_client=httpx.Client(verify=self._get_ssl_verify_setting()), + ) + + def create_batch_request_file_on_disk(self, tmpdir, model: str): + request_id = self.generate_request_id() + batch_request = { + "custom_id": request_id, + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": model, + "messages": [ + {"role": "user", "content": "What is 2+2?"}, + ], + }, + } + + request_file = os.path.join(tmpdir, f"request-{request_id}.jsonl") + with open(request_file, "w") as f: + f.write(json.dumps(batch_request)) + + return request_file + + def create_batch_input_file( + self, + client: openai.OpenAI, + request_file: str, + expiry_seconds: int = MIN_EXPIRY_SECONDS, + target_model_names: str = None, + ): + extra_body = { + "expires_after": { + "seconds": expiry_seconds, + "anchor": "created_at", + }, + } + if target_model_names: + extra_body["target_model_names"] = target_model_names + + batch_input_file = client.files.create( + file=open(request_file, "rb"), + purpose="batch", + extra_body=extra_body, + ) + return batch_input_file + + def create_batch( + self, + client: openai.OpenAI, + input_file_id: str, + expiry_seconds: int = MIN_EXPIRY_SECONDS, + ): + batch = client.batches.create( + input_file_id=input_file_id, + endpoint="/v1/chat/completions", + completion_window="24h", + extra_body={ + "output_expires_after": { + "seconds": expiry_seconds, + "anchor": "created_at", + }, + }, + ) + return batch + + def wait_for_batch_state( + self, + client: openai.OpenAI, + batch_id: str, + expected_status: str, + max_seconds: int = 60, + wait_seconds: int = 5, + state_tracker: "StateTracker | NoOpStateTracker | None" = None, + ): + if state_tracker is None: + state_tracker = NoOpStateTracker() + poll_count = 0 + for attempt in Retrying( + stop=stop_after_delay(max_seconds), + wait=wait_fixed(wait_seconds), + ): + with attempt: + poll_count += 1 + batch_response = client.batches.retrieve(batch_id=batch_id) + print( + f"[{time.strftime('%H:%M:%S')}] Poll #{poll_count}: Batch status: {batch_response.status}, expected: {expected_status}", + ) + state_tracker.print_state( + f"Poll #{poll_count} - status: {batch_response.status}", + ) + if batch_response.status == expected_status: + return batch_response + if batch_response.status in ["failed", "expired", "cancelled"]: + raise Exception( + f"Batch failed with status: {batch_response.status}", + ) + raise Exception(f"Batch not in {expected_status} state yet") + return None + + def wait_for_batch_completed( + self, + client: openai.OpenAI, + batch_id: str, + max_seconds: int = 120, + wait_seconds: int = 5, + ): + return self.wait_for_batch_state( + client, + batch_id, + "completed", + max_seconds, + wait_seconds, + ) + + def shorten_id(self, id_str: str) -> str: + if id_str is None: + return "None" + if len(id_str) <= 20: + return id_str + return id_str[:8] + "..." + id_str[-8:] + + def reset_mock_server(self): + if not use_mock_models(): + return + print("Resetting mock server state...") + reset_response = httpx.post(f"{get_mock_server_base_url()}/reset") + assert reset_response.status_code == 200, f"Reset failed: {reset_response.text}" + + def print_file_metadata(self, file_obj, label="File"): + print(f"{label} metadata:") + print(f"\tid={self.shorten_id(file_obj.id)}") + print(f"\tobject={file_obj.object}") + print(f"\tbytes={file_obj.bytes}") + print(f"\tfilename={file_obj.filename}") + print(f"\tpurpose={file_obj.purpose}") + print(f"\tstatus={file_obj.status}") + print(f"\tcreated_at={file_obj.created_at}") + print(f"\texpires_at={file_obj.expires_at}") + if file_obj.status_details: + print(f"\tstatus_details={file_obj.status_details}") + + def print_batch_metadata(self, batch): + print("Batch metadata:") + print(f"\tid={self.shorten_id(batch.id)}") + print(f"\tstatus={batch.status}") + print(f"\tendpoint={batch.endpoint}") + print(f"\tcompletion_window={batch.completion_window}") + print(f"\tinput_file_id={self.shorten_id(batch.input_file_id)}") + print(f"\tcreated_at={batch.created_at}") + print(f"\texpires_at={batch.expires_at}") + print(f"\tin_progress_at={batch.in_progress_at}") + print(f"\tcompleted_at={batch.completed_at}") + print(f"\toutput_file_id={self.shorten_id(batch.output_file_id)}") + print(f"\trequest_counts={batch.request_counts}") + + def wait_for_batch_list(self, model_name, max_seconds=90, wait_seconds=10): + for attempt in Retrying( + stop=stop_after_delay(max_seconds), + wait=wait_fixed(wait_seconds), + ): + with attempt: + batches_list = self.openai_client.batches.list( + limit=10, + # extra query is not supported by managed batches + # extra_query={"target_model_names": model_name}, + ) + print( + f"Batches in list: {len(batches_list.data)}", + ) + if len(batches_list.data) == 0: + raise Exception("No batches found in list yet") + print("Batches in list:") + for batch in batches_list.data: + print( + f" ID: {self.shorten_id(batch.id)} Status: {batch.status}, Created at: {batch.created_at}, Completed at: {batch.completed_at}", + ) + return batches_list + return None + + def wait_for_batch_in_list( + self, + client: openai.OpenAI, + batch_id: str, + max_seconds: int = 10, + wait_seconds: float = 0.5, + ): + """Wait for a specific batch to appear in the batch list. + + This handles the race condition where batch creation returns before + the database insert completes (due to asyncio.create_task). + """ + for attempt in Retrying( + stop=stop_after_delay(max_seconds), + wait=wait_fixed(wait_seconds), + ): + with attempt: + batches_list = client.batches.list(limit=20) + batch_ids = [b.id for b in batches_list.data] + if batch_id not in batch_ids: + raise Exception( + f"Batch {self.shorten_id(batch_id)} not found in list yet", + ) + return batches_list + return None \ No newline at end of file diff --git a/tests/proxy_e2e_azure_batches_tests/test_proxy_e2e_azure_batches.py b/tests/proxy_e2e_azure_batches_tests/test_proxy_e2e_azure_batches.py new file mode 100644 index 0000000000..262c55efc5 --- /dev/null +++ b/tests/proxy_e2e_azure_batches_tests/test_proxy_e2e_azure_batches.py @@ -0,0 +1,323 @@ +import base64 +import os +import sys +import time +import warnings + +import httpx +import openai +import pytest +from tenacity import RetryError + +sys.path.insert(0, os.path.abspath("../..")) + +from base_integration_test import ( + get_mock_server_base_url, + model_id, + use_mock_models, + UserKeyTestMixin, +) +from test_managed_files_base import ( + ManagedFilesBase, + MIN_EXPIRY_SECONDS, + get_batch_model_names, +) + +MANAGED_FILE_ID_PREFIX = "litellm_proxy" + +pytestmark = [ + pytest.mark.usefixtures("mock_azure_server", "litellm_proxy_server"), + pytest.mark.skipif( + os.environ.get("SKIP_E2E_TESTS", "false").lower() == "true", + reason="E2E tests disabled via SKIP_E2E_TESTS env var" + ), +] + + +def is_managed_id(file_id: str) -> bool: + """Check if a file ID is a base64-encoded LiteLLM managed/unified ID.""" + try: + padded = file_id + "=" * (-len(file_id) % 4) + decoded = base64.urlsafe_b64decode(padded).decode() + return decoded.startswith(MANAGED_FILE_ID_PREFIX) + except Exception: + return False + + +def assert_managed_id(file_id: str, label: str): + assert is_managed_id(file_id), f"{label} should be a managed ID, got raw: {file_id}" + + +def wip_features_enabled() -> bool: + return os.environ.get("WIP_FEATURES", "").lower() == "true" + + +class TestManagedFilesAPI(ManagedFilesBase, UserKeyTestMixin): + @classmethod + def setup_class(cls): + super().setup_class() + cls.setup_admin_client() + + @classmethod + def teardown_class(cls): + cls.teardown_admin_client() + + @pytest.fixture(autouse=True) + def setup_test(self): + print( + f"\nBase URL: {self.base_url}, Using mock models: {use_mock_models()}", + ) + self.clear_s3_callbacks() + + user_id, api_key, user_email, client = self.create_user_key_and_client( + "e2e-batch", + ) + self.test_user_id = user_id + self.openai_client = client + print(f"Using user {user_email} (id={user_id})") + + def _create_and_verify_batch_input_file(self, tmp_path, model_name): + request_file = self.create_batch_request_file_on_disk(tmp_path, model_name) + + print("Creating batch input file...") + batch_input_file = self.create_batch_input_file( + self.openai_client, + request_file, + MIN_EXPIRY_SECONDS, + target_model_names=model_name, + ) + print(f"Created batch input file: {self.shorten_id(batch_input_file.id)}") + assert_managed_id(batch_input_file.id, "batch_input_file.id") + + print("Retrieving batch input file metadata...") + metadata = self.openai_client.files.retrieve(batch_input_file.id) + assert_managed_id(metadata.id, "files.retrieve(input).id") + assert metadata.id == batch_input_file.id, ( + f"Input file ID mismatch: retrieve returned '{metadata.id}' but expected '{batch_input_file.id}'" + ) + assert metadata.object == "file" + assert metadata.bytes > 0, "bytes not set" + assert metadata.filename == "modified_file.jsonl" + assert metadata.purpose == "batch" + assert metadata.status in ["uploaded", "processed", "error"] + assert metadata.created_at > 0 + if wip_features_enabled(): + assert metadata.expires_at > 0, "expires_at not set" + self.print_file_metadata(metadata, "Input file") + + return batch_input_file + + def _create_and_verify_batch(self, input_file_id): + print("\nCreating batch...") + batch = self.create_batch( + self.openai_client, + input_file_id, + MIN_EXPIRY_SECONDS, + ) + print(f"Created batch: {self.shorten_id(batch.id)}") + + assert batch.id, "No batch ID returned" + assert_managed_id(batch.id, "batch.id") + assert_managed_id(batch.input_file_id, "batch.input_file_id") + assert batch.input_file_id == input_file_id, "batch.input_file_id mismatch" + assert batch.status in ["validating", "in_progress", "finalizing", "completed"] + if not batch.expires_at: + warnings.warn("batch expires_at not set") + else: + assert batch.expires_at > 0 + if not batch.endpoint: + warnings.warn("batch.endpoint empty - Azure API quirk, not a bug") + else: + assert batch.endpoint == "/v1/chat/completions" + assert batch.completion_window == "24h" + assert batch.created_at > 0 + self.print_batch_metadata(batch) + + return batch + + def _list_batches(self, batch_id, model_name): + if not wip_features_enabled(): + return + print("\nListing batches...") + try: + batches_list = self.wait_for_batch_list( + model_name, + max_seconds=30, + wait_seconds=5, + ) + batch_ids = [b.id for b in (batches_list.data if batches_list else [])] + if batch_id not in batch_ids: + warnings.warn( + f"Batch {batch_id} not found in list. " + f"batches.list returns raw IDs, not encoded IDs. raw IDs: {batch_ids}", + ) + except openai.APIError as e: + pytest.fail(f"batches.list() failed: {e}") + + def _wait_for_batch_completion(self, batch_id, tracker): + print(f"\nWaiting for batch {self.shorten_id(batch_id)} to complete...") + try: + batch_response = self.wait_for_batch_state( + self.openai_client, + batch_id, + "completed", + max_seconds=25 * 60, + wait_seconds=15, + state_tracker=tracker, + ) + except RetryError: + tracker.print_state("Timeout waiting for batch completion") + raise TimeoutError("Timed out waiting for batch to be in state: completed") + + assert_managed_id(batch_response.id, "batch_response.id") + assert batch_response.id == batch_id, ( + f"batch_response.id mismatch: got '{batch_response.id}' but expected '{batch_id}'" + ) + assert_managed_id(batch_response.input_file_id, "batch_response.input_file_id") + assert_managed_id( + batch_response.output_file_id, + "batch_response.output_file_id", + ) + + return batch_response + + def _get_and_verify_batch_output(self, output_file_id): + print("\nRetrieving batch output file metadata...") + metadata = self.openai_client.files.retrieve(output_file_id) + assert_managed_id(metadata.id, "files.retrieve(output_file_id).id") + assert metadata.id == output_file_id, ( + f"Output file ID mismatch: retrieve returned '{metadata.id}' but expected '{output_file_id}'" + ) + assert metadata.object == "file" + assert metadata.bytes > 0, "bytes not set" + assert metadata.filename, "filename not set" + assert metadata.purpose in ["batch_output", "batch"] + assert metadata.created_at > 0 + self.print_file_metadata(metadata, "Output file") + + print("\nFetching batch output file content...") + content = self.openai_client.files.content(output_file_id) + assert content.text, "No batch file content returned" + assert len(content.text) > 0, "Batch file content is empty" + print(f"Output file content ({len(content.text)} bytes):") + for line in content.text.strip().split("\n")[:3]: + print(f"\t{line}") + + return metadata + + def _delete_file(self, file_id, label, max_retries=6, retry_delay=10): + print(f"\nDeleting {label}: {self.shorten_id(file_id)}") + for attempt in range(max_retries): + try: + self.openai_client.files.delete(file_id) + return + except openai.BadRequestError as e: + if "batch_processed" in str(e) and attempt < max_retries - 1: + print( + f" File still referenced by unprocessed batch, " + f"retrying in {retry_delay}s ({attempt + 1}/{max_retries})" + ) + time.sleep(retry_delay) + else: + pytest.fail(f"files.delete({label}) failed: {e}") + except openai.APIError as e: + pytest.fail(f"files.delete({label}) failed: {e}") + + def _verify_file_deleted(self, file_id, label): + print(f"Verifying {label} is deleted...") + try: + self.openai_client.files.content(file_id) + assert False, f"{label} {file_id} still accessible after deletion" + except openai.NotFoundError: + print(f"{label} correctly not accessible after deletion") + + # ------------------------------------------------------------------ + # Tests + # ------------------------------------------------------------------ + + @pytest.mark.parametrize( + "model_name", + get_batch_model_names(), + ids=model_id, + ) + def test_e2e_managed_batch(self, tmp_path, model_name): + print( + f"\n\nStarting test with base_url={self.base_url} and model_name={model_name}\n", + ) + self.reset_mock_server() + tracker = self.create_state_tracker() + + batch_input_file = self._create_and_verify_batch_input_file( + tmp_path, + model_name, + ) + tracker.set_file_id(batch_input_file.id) + tracker.print_state("After creating batch input file") + + batch = self._create_and_verify_batch(batch_input_file.id) + tracker.set_batch_id(batch.id) + tracker.print_state("After creating batch") + + self._list_batches(batch.id, model_name) + + batch_response = self._wait_for_batch_completion(batch.id, tracker) + tracker.print_state("After batch completed") + + self._get_and_verify_batch_output(batch_response.output_file_id) + tracker.print_state("After retrieving output file") + + tracker.print_state("Final state after cleanup") + tracker.wait_and_print_s3_callbacks() + tracker.assert_batch_cost_callback() + + self._delete_file(batch_input_file.id, "input file") + self._delete_file(batch_response.output_file_id, "output file") + + self._verify_file_deleted(batch_input_file.id, "input file") + self._verify_file_deleted(batch_response.output_file_id, "output file") + + def cleanup_batches_in_database(self): + import psycopg2 + + print("Cleaning up stale batch records from database...") + try: + conn = psycopg2.connect( + host="localhost", + port=5432, + database="litellm", + user="llmproxy", + password="dbpassword9090", + ) + with conn.cursor() as cur: + cur.execute(""" + DELETE FROM "LiteLLM_ManagedObjectTable" + WHERE file_purpose = 'batch' AND status = 'validating' + """) + deleted = cur.rowcount + conn.commit() + if deleted > 0: + print(f"Deleted {deleted} stale batch records") + conn.close() + except Exception as e: + print(f"Warning: Could not clean up database: {e}") + + def clear_s3_callbacks(self): + clear_response = httpx.delete(f"{get_mock_server_base_url()}/mock-s3/callbacks") + assert clear_response.status_code == 200, ( + f"Failed to clear callbacks: {clear_response.text}" + ) + return clear_response.json() + + @pytest.mark.skipif( + True, + reason="Skipping managed files test till managed files feature is available", + ) + @pytest.mark.parametrize( + "model_name", + get_batch_model_names(), + ids=model_id, + ) + def test_error_files(self, tmp_path, model_name): + raise NotImplementedError( + "To implement. Fail a batch and retrieve the error file.", + ) \ No newline at end of file diff --git a/tests/proxy_e2e_azure_batches_tests/validate_e2e_setup.py b/tests/proxy_e2e_azure_batches_tests/validate_e2e_setup.py new file mode 100644 index 0000000000..e3991f2100 --- /dev/null +++ b/tests/proxy_e2e_azure_batches_tests/validate_e2e_setup.py @@ -0,0 +1,119 @@ +#!/usr/bin/env python +""" +Validation script for Azure Batch E2E test setup. +Run this before running the actual tests to verify all components are accessible. +""" + +import os +import sys +from pathlib import Path + +sys.path.insert(0, os.path.abspath("../..")) + +def check_imports(): + """Verify all required imports work.""" + print("Checking imports...") + try: + from base_integration_test import ( + get_mock_server_base_url, + get_litellm_base_url, + get_litellm_api_key, + ) + print(" ✓ base_integration_test imports OK") + + from test_managed_files_base import ManagedFilesBase, get_batch_model_names + print(" ✓ test_managed_files_base imports OK") + + from fixtures.mock_azure_batch_server import create_mock_azure_batch_server + print(" ✓ mock_azure_batch_server imports OK") + + import httpx + import openai + import psycopg2 + import uvicorn + print(" ✓ All external dependencies OK") + + return True + except ImportError as e: + print(f" ✗ Import error: {e}") + return False + + +def check_config_file(): + """Verify config file exists.""" + print("\nChecking config file...") + config_path = Path(__file__).parent / "fixtures" / "config.yml" + if config_path.exists(): + print(f" ✓ Config file found: {config_path}") + return True + else: + print(f" ✗ Config file not found: {config_path}") + return False + + +def check_database(): + """Verify database connection.""" + print("\nChecking database connection...") + try: + import psycopg2 + conn = psycopg2.connect( + host="localhost", + port=5432, + database="litellm", + user="llmproxy", + password="dbpassword9090", + ) + conn.close() + print(" ✓ Database connection OK") + return True + except Exception as e: + print(f" ✗ Database connection failed: {e}") + print(" Start PostgreSQL with:") + print(" docker run --name litellm-postgres -e POSTGRES_USER=llmproxy \\") + print(" -e POSTGRES_PASSWORD=dbpassword9090 -e POSTGRES_DB=litellm \\") + print(" -p 5432:5432 -d postgres:15") + return False + + +def check_ports(): + """Check if required ports are available.""" + print("\nChecking ports...") + import socket + + for port, name in [(4000, "LiteLLM Proxy"), (8090, "Mock Server")]: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + try: + s.bind(("localhost", port)) + print(f" ✓ Port {port} ({name}) is available") + except OSError: + print(f" ⚠ Port {port} ({name}) is in use (will reuse if healthy)") + return True + + +def main(): + print("=" * 70) + print("Azure Batch E2E Test Setup Validation") + print("=" * 70) + + checks = [ + check_imports(), + check_config_file(), + check_database(), + check_ports(), + ] + + print("\n" + "=" * 70) + if all(checks): + print("✓ All checks passed! Ready to run E2E tests.") + print("\nRun tests with:") + print(" cd litellm") + print(" export DATABASE_URL='postgresql://llmproxy:dbpassword9090@localhost:5432/litellm'") + print(" poetry run pytest tests/proxy_e2e_azure_batches_tests/test_proxy_e2e_azure_batches.py -vv") + return 0 + else: + print("✗ Some checks failed. Please fix the issues above.") + return 1 + + +if __name__ == "__main__": + sys.exit(main())