Merge pull request #41948 from BerriAI/litellm_unit_shard_per_test_timeout

ci(unit): fail a hung test in 120s with a traceback instead of idling the shard to its step timeout

.github/workflows/_test-unit-base.yml

@@ -37,6 +37,18 @@ on:
         required: false
         type: number
         default: 60
+      test-timeout-seconds:
+        description: >-
+          Per-test ceiling enforced by pytest-timeout, covering fixture setup and
+          teardown as well as the test body. A test that hangs fails with a
+          traceback of where it was stuck instead of idling the shard until
+          `timeout-minutes` cancels it. Timed-out tests are excluded from reruns
+          because pytest-timeout arms its timer once per test and
+          pytest-rerunfailures reruns inside that same window, so a rerun of a
+          timed-out test would run with no timer at all.
+        required: false
+        type: number
+        default: 120
       max-failures:
         description: "Stop after this many failures"
         required: false
@@ -137,6 +149,7 @@ jobs:
           MAX_FAILURES: ${{ inputs.max-failures }}
           WORKERS: ${{ inputs.workers }}
           RERUNS: ${{ inputs.reruns }}
+          TEST_TIMEOUT_SECONDS: ${{ inputs.test-timeout-seconds }}
           DIST: ${{ inputs.dist }}
           COVERAGE_CORE: sysmon
         run: |
@@ -146,6 +159,8 @@ jobs:
               --maxfail="${MAX_FAILURES}" \
               --reruns "${RERUNS}" \
               --reruns-delay 1 \
+              --timeout="${TEST_TIMEOUT_SECONDS}" \
+              --rerun-except "from pytest-timeout" \
               --durations=20 \
               --cov=./litellm --cov=./enterprise/litellm_enterprise \
               --cov-report=xml:coverage.xml \
@@ -157,6 +172,8 @@ jobs:
               -n "${WORKERS}" \
               --reruns "${RERUNS}" \
               --reruns-delay 1 \
+              --timeout="${TEST_TIMEOUT_SECONDS}" \
+              --rerun-except "from pytest-timeout" \
               --dist="${DIST}" \
               --durations=20 \
               --cov=./litellm --cov=./enterprise/litellm_enterprise \

tests/test_litellm/rerank_api/test_main.py

@@ -239,6 +239,7 @@ async def test_arerank_error_is_mapped_to_litellm_exception(respx_mock: respx.Mo
 
 
 @pytest.mark.asyncio
+@pytest.mark.timeout(300)
 async def test_arerank_declared_authenticating_provider_skips_resolution(monkeypatch):
     """Regression for the event-loop hazard in arerank's provider pre-resolution:
     get_llm_provider runs the blocking OAuth device flow for github_copilot/chatgpt,

tests/test_litellm/test_unit_shard_per_test_timeout.py

@@ -0,0 +1,83 @@
+import shlex
+import subprocess
+import sys
+from pathlib import Path
+from string import Template
+from types import MappingProxyType
+from typing import Final
+
+import pytest
+import yaml
+
+_REPO_ROOT: Final = Path(__file__).resolve().parents[2]
+_BASE_WORKFLOW: Final = _REPO_ROOT / ".github" / "workflows" / "_test-unit-base.yml"
+_SHARD_ENV: Final = MappingProxyType({"WORKERS": "2", "RERUNS": "2", "DIST": "loadscope", "TEST_TIMEOUT_SECONDS": "1"})
+_HANG_GUARD_FLAGS: Final = frozenset(("-n", "--dist", "--reruns", "--reruns-delay", "--timeout", "--rerun-except"))
+_HUNG_TEST_MODULE: Final = """
+import threading
+
+import pytest
+
+
+@pytest.fixture
+def hangs_on_teardown():
+    yield
+    threading.Event().wait()
+
+
+def test_body_waits_forever():
+    threading.Event().wait()
+
+
+def test_fixture_teardown_waits_forever(hangs_on_teardown):
+    assert True
+
+
+def test_passes():
+    assert True
+"""
+
+
+def _run_tests_script() -> str:
+    workflow: Final = yaml.safe_load(_BASE_WORKFLOW.read_text())
+    return next(step["run"] for step in workflow["jobs"]["run"]["steps"] if step.get("name") == "Run tests")
+
+
+def _pytest_invocations(script: str) -> tuple[tuple[str, ...], ...]:
+    return tuple(tuple(shlex.split(line)) for line in script.replace("\\\n", " ").splitlines() if " pytest " in line)
+
+
+def _hang_guard_args(invocation: tuple[str, ...]) -> tuple[str, ...]:
+    return tuple(
+        Template(token).safe_substitute(_SHARD_ENV)
+        for previous, token in zip(("", *invocation), invocation)
+        if token.split("=", 1)[0] in _HANG_GUARD_FLAGS or previous in _HANG_GUARD_FLAGS
+    )
+
+
+_INVOCATIONS: Final = _pytest_invocations(_run_tests_script())
+
+
+@pytest.mark.parametrize(
+    "invocation", _INVOCATIONS, ids=tuple("xdist" if "-n" in invocation else "serial" for invocation in _INVOCATIONS)
+)
+def test_a_hung_test_fails_fast_and_names_itself_under_the_shard_flags(
+    invocation: tuple[str, ...], tmp_path: Path
+) -> None:
+    hung_module: Final = tmp_path / "test_hung.py"
+    hung_module.write_text(_HUNG_TEST_MODULE)
+
+    result: Final = subprocess.run(
+        (sys.executable, "-m", "pytest", str(hung_module), "-p", "no:cacheprovider", *_hang_guard_args(invocation)),
+        cwd=tmp_path,
+        capture_output=True,
+        text=True,
+        timeout=90,
+        check=False,
+    )
+
+    assert result.returncode == 1, result.stdout
+    assert "FAILED test_hung.py::test_body_waits_forever" in result.stdout
+    assert "ERROR test_hung.py::test_fixture_teardown_waits_forever" in result.stdout
+    assert "Timeout (>1.0s) from pytest-timeout" in result.stdout
+    assert "1 failed, 2 passed, 1 error" in result.stdout