litellm/.github/workflows/test-unit-proxy-db.yml
Yuneng Jiang e0201ece1e
[Infra] Split slow proxy-db shards to hit 7m wall-clock target
Previous run (13.8m total) was bottlenecked by shards with 9-12m wall-clock.
Setup + xdist spawn + coverage teardown is ~3m per shard, so each shard's
pytest runtime must stay under ~4m to fit inside 7m total.

Observed per-shard pytest times (before split):
  db-and-spend            9:08   (170s outlier: test_aaaasschema_migration_check)
  proxy-server            7:15
  logging-and-callbacks   6:45
  guardrails-budget-hooks 6:37
  proxy-utils             6:23
  auth-and-jwt            6:54

Split 6 shards into 12, keeping key-generation and endpoints-and-responses
(already <7m). Adds a `keyword` input to _test-unit-services-base.yml so
test_proxy_utils.py can be split by -k expression (same file, two runners).
New matrix entries:

  auth-and-jwt           -> auth-checks + jwt-and-keys
  proxy-server           -> proxy-server-core + proxy-runtime
  logging-and-callbacks  -> custom-logging + logging-misc
  db-and-spend           -> schema-migration (isolated 170s test) + db-and-spend
  guardrails-budget-hooks-> guardrails-hooks + budgets
  proxy-utils            -> proxy-utils-a-h + proxy-utils-i-z (-k split)

The -k expression split is verified to cover every one of the 64 test
functions in test_proxy_utils.py exactly once. The assert-shard-coverage
guard still catches any file not in any shard.
2026-04-23 15:25:37 -07:00

270 lines
12 KiB
YAML

name: "Unit Tests: Proxy DB Operations"
# Uses DATABASE_URL secret — only runs on trusted branches, not PRs.
on:
push:
branches: [main, "litellm_**"]
permissions:
contents: read
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: true
# Semantic matrix: each shard groups tests by concern (auth, server, logging, …)
# rather than alphabetical letter ranges. Adding a new test file means adding it
# to whichever group it belongs to, not reshuffling slices.
#
# Design targets:
# * Every shard runs in <= 7 minutes of wall-clock on the default runner.
# Setup + xdist worker spawn + coverage teardown is ~3 minutes per shard,
# so each shard's pytest runtime must stay under ~4 minutes. That drives
# the split granularity: shards get subdivided when pytest call time
# exceeds ~4m or any single test exceeds ~3m (it pins one xdist worker).
# * test_key_generate_prisma.py stays serial (workers=0) — it has event-loop
# conflicts with the logging worker when run in parallel.
# * test_proxy_utils.py is split into two -k-filtered shards (by first
# character of the test function name) so its 188 parametrized cases
# fan out across two runners rather than one. --dist=worksteal within
# each shard balances parametrized cases across xdist workers.
# * test_db_schema_migration.py is isolated because one test in it
# (test_aaaasschema_migration_check) takes ~170s — by itself it
# determines the shard's wall-clock floor.
jobs:
# Fast guard — fails the workflow if a test_*.py file under
# tests/proxy_unit_tests/ is not referenced by any matrix entry below.
# The semantic-shard design (no catch-all "remaining" bucket) relies on
# every test file being explicitly assigned; this guard prevents a new
# file from silently dropping out of CI.
assert-shard-coverage:
runs-on: ubuntu-latest
timeout-minutes: 2
permissions:
contents: read
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Assert every test_*.py is in a matrix shard
run: |
python3 - <<'PY'
import pathlib, sys, yaml
wf = yaml.safe_load(open(".github/workflows/test-unit-proxy-db.yml"))
matrix = wf["jobs"]["proxy-db"]["strategy"]["matrix"]["include"]
referenced = set()
for entry in matrix:
for token in entry["test-path"].split():
if token.startswith("tests/proxy_unit_tests/"):
referenced.add(pathlib.PurePosixPath(token).name)
actual = {p.name for p in pathlib.Path("tests/proxy_unit_tests").iterdir()
if p.name.startswith("test_") and (p.suffix == ".py" or p.is_dir())
and p.name != "test_configs"}
orphans = sorted(actual - referenced)
if orphans:
print("ERROR: the following files/dirs under tests/proxy_unit_tests/")
print(" are not assigned to any shard in test-unit-proxy-db.yml:")
for o in orphans:
print(f" - {o}")
print()
print("Add each to whichever semantic shard it belongs to.")
sys.exit(1)
print(f"OK: all {len(actual)} files assigned to a shard.")
PY
proxy-db:
needs: assert-shard-coverage
permissions:
contents: read
id-token: write
pull-requests: write
strategy:
fail-fast: false
matrix:
include:
# Must run serially — event-loop conflict with the logging worker.
- test-group: key-generation
test-path: "tests/proxy_unit_tests/test_key_generate_prisma.py"
workers: 0
dist: loadscope
keyword: ""
timeout: 20
# ---- auth: split into 2 shards (was 1 at ~10.4m wall-clock) ----
- test-group: auth-checks
test-path: >-
tests/proxy_unit_tests/test_auth_checks.py
tests/proxy_unit_tests/test_user_api_key_auth.py
workers: 8
dist: loadscope
keyword: ""
timeout: 15
- test-group: jwt-and-keys
test-path: >-
tests/proxy_unit_tests/test_jwt.py
tests/proxy_unit_tests/test_jwt_key_mapping.py
tests/proxy_unit_tests/test_proxy_custom_auth.py
tests/proxy_unit_tests/test_key_generate_dynamodb.py
tests/proxy_unit_tests/test_deployed_proxy_keygen.py
workers: 8
dist: loadscope
keyword: ""
timeout: 15
# ---- test_proxy_utils.py split into 2 by -k (was 1 at ~9.7m) ----
# Same file, same --dist=worksteal, filtered by first char of test
# function name. Keywords below cover all 63 test functions in the
# file. If new functions are added, balance between the two shards.
- test-group: proxy-utils-a-h
test-path: "tests/proxy_unit_tests/test_proxy_utils.py"
workers: 8
dist: worksteal
keyword: >-
test_add or test_check or test_custom or test_during or
test_dynamic or test_end_user or test_enforced or test_foward or
test_get_admin or test_get_complete or test_get_docs or
test_get_known or test_get_model_group or test_get_openapi or
test_get_redoc or test_get_temp or test_get_user_info or
test_handle or test_health
timeout: 15
- test-group: proxy-utils-i-z
test-path: "tests/proxy_unit_tests/test_proxy_utils.py"
workers: 8
dist: worksteal
keyword: >-
test_is or test_litellm or test_merge or test_post_call or
test_prepare or test_provider or test_proxy_config or
test_reading or test_spend or test_team or test_traceparent or
test_update or test_get_key or test_get_team
timeout: 15
# ---- proxy server: split into 2 shards (was 1 at ~11.1m) ----
- test-group: proxy-server-core
test-path: >-
tests/proxy_unit_tests/test_proxy_server.py
tests/proxy_unit_tests/test_proxy_server_keys.py
tests/proxy_unit_tests/test_proxy_server_caching.py
tests/proxy_unit_tests/test_proxy_server_langfuse.py
tests/proxy_unit_tests/test_proxy_server_spend.py
tests/proxy_unit_tests/test_aproxy_startup.py
workers: 8
dist: loadscope
keyword: ""
timeout: 15
- test-group: proxy-runtime
test-path: >-
tests/proxy_unit_tests/test_proxy_config_unit_test.py
tests/proxy_unit_tests/test_proxy_routes.py
tests/proxy_unit_tests/test_proxy_gunicorn.py
tests/proxy_unit_tests/test_server_root_path.py
tests/proxy_unit_tests/test_proxy_pass_user_config.py
tests/proxy_unit_tests/test_proxy_token_counter.py
workers: 8
dist: loadscope
keyword: ""
timeout: 15
# ---- logging: split into 2 shards (was 1 at ~10.1m) ----
- test-group: custom-logging
test-path: >-
tests/proxy_unit_tests/test_custom_callback_input.py
tests/proxy_unit_tests/test_custom_logger_s3_gcs.py
tests/proxy_unit_tests/test_proxy_custom_logger.py
workers: 8
dist: loadscope
keyword: ""
timeout: 15
- test-group: logging-misc
test-path: >-
tests/proxy_unit_tests/test_proxy_reject_logging.py
tests/proxy_unit_tests/test_audit_logs_proxy.py
tests/proxy_unit_tests/test_search_api_logging.py
workers: 8
dist: loadscope
keyword: ""
timeout: 15
# ---- db-and-spend: split out the 170s schema-migration test ----
# test_db_schema_migration.py has one test that runs ~170s; it
# single-handedly pins one xdist worker and determined the whole
# shard's 12.3m wall-clock. Isolated here so the other 45 tests
# finish faster.
- test-group: schema-migration
test-path: "tests/proxy_unit_tests/test_db_schema_migration.py"
workers: 8
dist: loadscope
keyword: ""
timeout: 15
- test-group: db-and-spend
test-path: >-
tests/proxy_unit_tests/test_prisma_client_backoff_retry.py
tests/proxy_unit_tests/test_db_schema_changes.py
tests/proxy_unit_tests/test_e2e_pod_lock_manager.py
tests/proxy_unit_tests/test_skills_db.py
tests/proxy_unit_tests/test_update_daily_tag_spend.py
tests/proxy_unit_tests/test_update_spend.py
tests/proxy_unit_tests/test_project_endpoints_prisma.py
tests/proxy_unit_tests/test_proxy_encrypt_decrypt.py
workers: 8
dist: loadscope
keyword: ""
timeout: 15
# ---- guardrails + budget + hooks: split into 2 (was 1 at ~10.1m) ----
- test-group: guardrails-hooks
test-path: >-
tests/proxy_unit_tests/test_proxy_setting_guardrails.py
tests/proxy_unit_tests/test_banned_keyword_list.py
tests/proxy_unit_tests/test_unit_test_proxy_hooks.py
workers: 8
dist: loadscope
keyword: ""
timeout: 15
- test-group: budgets
test-path: >-
tests/proxy_unit_tests/test_default_end_user_budget_simple.py
tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py
tests/proxy_unit_tests/test_zero_cost_model_budget_bypass.py
workers: 8
dist: loadscope
keyword: ""
timeout: 15
# Already under 7m; left as a single shard.
- test-group: endpoints-and-responses
test-path: >-
tests/proxy_unit_tests/test_blog_posts_endpoint.py
tests/proxy_unit_tests/test_models_fallback_endpoint.py
tests/proxy_unit_tests/test_google_endpoint_routing.py
tests/proxy_unit_tests/test_google_gemini_proxy_request.py
tests/proxy_unit_tests/test_get_favicon.py
tests/proxy_unit_tests/test_get_image.py
tests/proxy_unit_tests/test_ui_path_detection.py
tests/proxy_unit_tests/test_prompt_test_endpoint.py
tests/proxy_unit_tests/test_check_batch_cost.py
tests/proxy_unit_tests/test_check_responses_cost.py
tests/proxy_unit_tests/test_response_polling_handler.py
tests/proxy_unit_tests/test_response_polling_pre_call_checks.py
tests/proxy_unit_tests/test_realtime_cache.py
tests/proxy_unit_tests/test_proxy_exception_mapping.py
tests/proxy_unit_tests/test_custom_tokenizer_bug.py
tests/proxy_unit_tests/test_model_response_typing
workers: 8
dist: loadscope
keyword: ""
timeout: 15
uses: ./.github/workflows/_test-unit-services-base.yml
with:
test-path: ${{ matrix.test-path }}
workers: ${{ matrix.workers }}
reruns: 2
timeout-minutes: ${{ matrix.timeout }}
enable-postgres: true
dist: ${{ matrix.dist }}
keyword: ${{ matrix.keyword }}
artifact-name: proxy-db-${{ matrix.test-group }}
secrets:
DATABASE_URL: ${{ secrets.DATABASE_URL }}
POSTGRES_USER: ${{ secrets.POSTGRES_USER }}
POSTGRES_PASSWORD: ${{ secrets.POSTGRES_PASSWORD }}