From 7b43f5981fa60a0eae3218ffe9b2ce907d64900e Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 20 Apr 2026 16:36:05 -0700 Subject: [PATCH 1/4] [Fix] CI: split test_proxy_utils.py into its own proxy-db matrix entry The "remaining" proxy-db job was consistently timing out at ~98% because --dist=loadscope pins every test in test_proxy_utils.py (168+ parametrized tests) to a single xdist worker. 7 workers finished their files in ~15 minutes, then one worker ran alone for another 8+ minutes and hit the 30-minute job cap. Give test_proxy_utils.py its own matrix entry so its tests spread across all 8 workers, and add it to the "remaining" ignore list. --- .github/workflows/test-unit-proxy-db.yml | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test-unit-proxy-db.yml b/.github/workflows/test-unit-proxy-db.yml index 87e7e17feb7..a631a7c3005 100644 --- a/.github/workflows/test-unit-proxy-db.yml +++ b/.github/workflows/test-unit-proxy-db.yml @@ -31,8 +31,15 @@ jobs: test-path: "tests/proxy_unit_tests/test_auth_checks.py tests/proxy_unit_tests/test_user_api_key_auth.py" workers: 8 timeout: 20 + # test_proxy_utils.py is large (168+ parametrized tests) — run it on its + # own matrix so --dist=loadscope doesn't pin all of it to a single xdist + # worker and push the "remaining" group past the job timeout. + - test-group: proxy-utils + test-path: "tests/proxy_unit_tests/test_proxy_utils.py" + workers: 8 + timeout: 20 - test-group: remaining - test-path: "tests/proxy_unit_tests --ignore=tests/proxy_unit_tests/test_key_generate_prisma.py --ignore=tests/proxy_unit_tests/test_auth_checks.py --ignore=tests/proxy_unit_tests/test_user_api_key_auth.py" + test-path: "tests/proxy_unit_tests --ignore=tests/proxy_unit_tests/test_key_generate_prisma.py --ignore=tests/proxy_unit_tests/test_auth_checks.py --ignore=tests/proxy_unit_tests/test_user_api_key_auth.py --ignore=tests/proxy_unit_tests/test_proxy_utils.py" workers: 8 timeout: 30 uses: ./.github/workflows/_test-unit-services-base.yml From ccf928361be6c36d8f5e6ac197775067cff442e4 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 20 Apr 2026 17:59:05 -0700 Subject: [PATCH 2/4] [Infra] Speed up proxy unit tests by replacing litellm reload with state snapshot MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit tests/proxy_unit_tests/conftest.py was calling importlib.reload(litellm) in an autouse function-scoped fixture, which cost ~17s per test because it re-ran the full litellm __init__ import chain. With 400+ proxy unit tests, this was the single biggest driver of CI wall time — 18 of the top 20 slowest durations in a typical run were just the 17s fixture setup. Replace the reload with a snapshot-and-restore approach: snapshot the mutable lists/dicts/sets on litellm and litellm.proxy.proxy_server once at conftest import, then deep-copy that snapshot back before each test. Callback lists, caches, router state, etc. still get reset between tests, but the expensive import chain only runs once per worker. Local measurement on test_proxy_utils.py: 188 tests in 3.50s (previously took ~15 minutes of CI wall time on a single worker). --- tests/proxy_unit_tests/conftest.py | 81 +++++++++++++++++++++--------- 1 file changed, 56 insertions(+), 25 deletions(-) diff --git a/tests/proxy_unit_tests/conftest.py b/tests/proxy_unit_tests/conftest.py index 1421700c9a8..0cde5bdf28b 100644 --- a/tests/proxy_unit_tests/conftest.py +++ b/tests/proxy_unit_tests/conftest.py @@ -1,6 +1,7 @@ # conftest.py -import importlib +import asyncio +import copy import os import sys @@ -9,40 +10,70 @@ import pytest sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path + import litellm +import litellm.proxy.proxy_server + + +def _snapshot_mutable_state(module): + """Deep-copy every list/dict/set module attribute for later restore. + + Classes, functions, submodules and primitives are skipped — only the + collections that tests mutate (callbacks, caches, routers, etc.) need + per-test isolation. + """ + snapshot = {} + for attr in list(vars(module)): + if attr.startswith("_"): + continue + try: + value = getattr(module, attr) + except Exception: + continue + if isinstance(value, (list, dict, set)): + try: + snapshot[attr] = copy.deepcopy(value) + except Exception: + # Unpickleable collections (e.g. holding open clients) can't + # round-trip through deepcopy; skip them rather than crash. + pass + return snapshot + + +def _restore_mutable_state(module, snapshot): + for attr, default in snapshot.items(): + try: + setattr(module, attr, copy.deepcopy(default)) + except Exception: + pass + + +# Snapshot once at conftest import — these are the "clean" module states. +_LITELLM_STATE = _snapshot_mutable_state(litellm) +_PROXY_SERVER_STATE = _snapshot_mutable_state(litellm.proxy.proxy_server) @pytest.fixture(scope="function", autouse=True) def setup_and_teardown(): """ - This fixture reloads litellm before every function. To speed up testing by removing callbacks being chained. + Reset mutable module state on litellm and proxy_server before every test. + + Replaces a previous importlib.reload(litellm) approach that cost ~17s + per test (re-executing the full litellm __init__ import chain). The + snapshot-and-restore below only touches collections that actually leak + across tests — callbacks, caches, router, etc. — and is effectively + instantaneous. """ - curr_dir = os.getcwd() # Get the current working directory - sys.path.insert( - 0, os.path.abspath("../..") - ) # Adds the project directory to the system path - - import litellm - from litellm import Router - - importlib.reload(litellm) - try: - if hasattr(litellm, "proxy") and hasattr(litellm.proxy, "proxy_server"): - importlib.reload(litellm.proxy.proxy_server) - except Exception as e: - print(f"Error reloading litellm.proxy.proxy_server: {e}") - - import asyncio + _restore_mutable_state(litellm, _LITELLM_STATE) + _restore_mutable_state(litellm.proxy.proxy_server, _PROXY_SERVER_STATE) loop = asyncio.get_event_loop_policy().new_event_loop() asyncio.set_event_loop(loop) - print(litellm) - # from litellm import Router, completion, aembedding, acompletion, embedding - yield - - # Teardown code (executes after the yield point) - loop.close() # Close the loop created earlier - asyncio.set_event_loop(None) # Remove the reference to the loop + try: + yield + finally: + loop.close() + asyncio.set_event_loop(None) def pytest_collection_modifyitems(config, items): From 5411ebedae0f77ed0832289ae877a75a1cca836f Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 20 Apr 2026 21:03:07 -0700 Subject: [PATCH 3/4] [Fix] conftest snapshot: also reset scalar module attributes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous snapshot only tracked list/dict/set values. Tests mutate scalar module attrs too — master_key, premium_user, prisma_client — and importlib.reload used to reset those implicitly. Under the snapshot approach they were leaking between tests, so test_active_callbacks failed in CI with "No api key passed in." once an earlier test left master_key set to sk-1234. Expand the snapshot to cover primitives (str/int/float/bool/bytes/tuple) and None-valued attributes. Complex object instances are still skipped to avoid deepcopy issues. --- tests/proxy_unit_tests/conftest.py | 21 +++++++++++++-------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/tests/proxy_unit_tests/conftest.py b/tests/proxy_unit_tests/conftest.py index 0cde5bdf28b..544b6a0b421 100644 --- a/tests/proxy_unit_tests/conftest.py +++ b/tests/proxy_unit_tests/conftest.py @@ -15,12 +15,17 @@ import litellm import litellm.proxy.proxy_server -def _snapshot_mutable_state(module): - """Deep-copy every list/dict/set module attribute for later restore. +_SNAPSHOT_TYPES = (list, dict, set, tuple, str, int, float, bool, bytes) - Classes, functions, submodules and primitives are skipped — only the - collections that tests mutate (callbacks, caches, routers, etc.) need - per-test isolation. + +def _snapshot_mutable_state(module): + """Snapshot every module attribute that importlib.reload would have reset. + + Covers the top-level assignments that tests mutate — collections + (callbacks, caches, general_settings) plus scalar flags (master_key, + premium_user, etc.) that gate auth and feature behavior. Classes, + functions, submodules and complex object instances are skipped: those + either aren't meant to be reset or can't round-trip through deepcopy. """ snapshot = {} for attr in list(vars(module)): @@ -30,12 +35,12 @@ def _snapshot_mutable_state(module): value = getattr(module, attr) except Exception: continue - if isinstance(value, (list, dict, set)): + if value is None or isinstance(value, _SNAPSHOT_TYPES): try: snapshot[attr] = copy.deepcopy(value) except Exception: - # Unpickleable collections (e.g. holding open clients) can't - # round-trip through deepcopy; skip them rather than crash. + # Skip anything that can't round-trip through deepcopy + # rather than crash collection. pass return snapshot From 4b3f5d7f81d38e2019882a63bbe91411b2e31065 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 20 Apr 2026 22:19:36 -0700 Subject: [PATCH 4/4] [Fix] conftest: flush cache instances and warn on silent skips MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses review feedback on the snapshot approach: 1. Class-instance mutable state The snapshot only covers primitives + collections + None. Class instances (DualCache, LLMClientCache) weren't reset between tests, so in-place cache mutations could leak. Can't deepcopy these — they hold thread locks — but they expose flush_cache(). Collect every module attribute whose value implements flush_cache() at conftest import, and invoke it per-test alongside the snapshot restore. 2. Silent skips are now warnings _snapshot_mutable_state and _restore_mutable_state previously swallowed exceptions, so if a future attr gained a property without a setter (or other non-round-trippable state), an isolation gap would have no signal. Emit warnings.warn on each failure path. 3. Docstring Explicitly documents what IS and IS NOT reset, and tells authors to use monkeypatch.setattr() for in-place mutations of instances without flush_cache() (ProxyLogging, JWTHandler, etc.). --- tests/proxy_unit_tests/conftest.py | 99 ++++++++++++++++++++++++------ 1 file changed, 79 insertions(+), 20 deletions(-) diff --git a/tests/proxy_unit_tests/conftest.py b/tests/proxy_unit_tests/conftest.py index 544b6a0b421..a0326f64ed7 100644 --- a/tests/proxy_unit_tests/conftest.py +++ b/tests/proxy_unit_tests/conftest.py @@ -2,8 +2,10 @@ import asyncio import copy +import inspect import os import sys +import warnings import pytest @@ -15,33 +17,34 @@ import litellm import litellm.proxy.proxy_server +# Top-level assignments of these types are the ones importlib.reload(litellm) +# would have effectively reset. We snapshot them at conftest import time and +# deep-copy the snapshot back before every test. _SNAPSHOT_TYPES = (list, dict, set, tuple, str, int, float, bool, bytes) def _snapshot_mutable_state(module): - """Snapshot every module attribute that importlib.reload would have reset. - - Covers the top-level assignments that tests mutate — collections - (callbacks, caches, general_settings) plus scalar flags (master_key, - premium_user, etc.) that gate auth and feature behavior. Classes, - functions, submodules and complex object instances are skipped: those - either aren't meant to be reset or can't round-trip through deepcopy. - """ + """Capture a per-module snapshot of primitive and collection attributes.""" snapshot = {} for attr in list(vars(module)): if attr.startswith("_"): continue try: value = getattr(module, attr) - except Exception: + except Exception as exc: + warnings.warn( + f"conftest: could not read {module.__name__}.{attr} during snapshot: {exc}", + stacklevel=2, + ) continue if value is None or isinstance(value, _SNAPSHOT_TYPES): try: snapshot[attr] = copy.deepcopy(value) - except Exception: - # Skip anything that can't round-trip through deepcopy - # rather than crash collection. - pass + except Exception as exc: + warnings.warn( + f"conftest: could not snapshot {module.__name__}.{attr}: {exc}", + stacklevel=2, + ) return snapshot @@ -49,28 +52,84 @@ def _restore_mutable_state(module, snapshot): for attr, default in snapshot.items(): try: setattr(module, attr, copy.deepcopy(default)) + except Exception as exc: + warnings.warn( + f"conftest: could not restore {module.__name__}.{attr}: {exc}", + stacklevel=2, + ) + + +def _collect_flushable_caches(): + """Return (module, attr) pairs whose values expose flush_cache().""" + targets = [] + for module in (litellm, litellm.proxy.proxy_server): + for attr in list(vars(module)): + if attr.startswith("_"): + continue + try: + value = getattr(module, attr) + except Exception: + continue + # Only instances — a class reference has an unbound flush_cache + # that can't be called without a self argument. + if inspect.isclass(value) or inspect.ismodule(value): + continue + if callable(getattr(value, "flush_cache", None)): + targets.append((module, attr)) + return targets + + +def _flush_caches(targets): + for module, attr in targets: + try: + value = getattr(module, attr) except Exception: - pass + continue + flush = getattr(value, "flush_cache", None) + if callable(flush): + try: + flush() + except Exception as exc: + warnings.warn( + f"conftest: flush_cache failed on {module.__name__}.{attr}: {exc}", + stacklevel=2, + ) # Snapshot once at conftest import — these are the "clean" module states. _LITELLM_STATE = _snapshot_mutable_state(litellm) _PROXY_SERVER_STATE = _snapshot_mutable_state(litellm.proxy.proxy_server) +_FLUSHABLE_CACHES = _collect_flushable_caches() @pytest.fixture(scope="function", autouse=True) def setup_and_teardown(): - """ - Reset mutable module state on litellm and proxy_server before every test. + """Reset mutable module state on litellm and proxy_server before each test. Replaces a previous importlib.reload(litellm) approach that cost ~17s - per test (re-executing the full litellm __init__ import chain). The - snapshot-and-restore below only touches collections that actually leak - across tests — callbacks, caches, router, etc. — and is effectively - instantaneous. + per test (re-executing the full litellm __init__ import chain). + + What IS reset: + - Top-level module attributes of type list / dict / set / tuple + / str / int / float / bool / bytes, and None-valued attributes. + These cover callback lists, general_settings, master_key, + premium_user, prisma_client, etc. — anything the old reload() reset + by re-executing the module body. + - Any module-level object instance that exposes flush_cache() (the + DualCache and LLMClientCache family), which handles cache state + that can't round-trip through deepcopy because of internal locks. + + What is NOT reset: + - Class instances without flush_cache() (e.g. ProxyLogging, + JWTHandler, FastAPI routers, loggers). If a test mutates such an + instance in-place (setattr on the instance, appending to one of + its internal lists, etc.), the mutation will leak into later tests. + Use pytest's monkeypatch.setattr() or a local fixture for those + cases — don't rely on this autouse fixture to undo them. """ _restore_mutable_state(litellm, _LITELLM_STATE) _restore_mutable_state(litellm.proxy.proxy_server, _PROXY_SERVER_STATE) + _flush_caches(_FLUSHABLE_CACHES) loop = asyncio.get_event_loop_policy().new_event_loop() asyncio.set_event_loop(loop)