feat(ci): guard shard assignment across every sharded test tree (#37593)

tests/proxy_unit_tests had a 30-line YAML parser inlined in its workflow that
failed the run when a test file there belonged to no shard. tests/test_litellm
is sharded the same way, with no catch-all bucket, and had no such guard: a new
directory under it (or under its proxy subtree) is collected by nothing and runs
nowhere, and the coverage census cannot see it because a token like
tests/test_litellm/test_*.py already answers 'yes, that tree runs'.

The two questions differ. The census asks whether a file runs at all, so an
ancestor path standing in for everything beneath it is a fine answer. Shard
assignment asks which shard owns a child, and there that same ancestor path is
precisely the bug. _token_covers keeps the first meaning; _token_names adds the
second, and the guard now walks a list of sharded trees rather than one hardcoded
directory. Both read the same test-path keys, so there is one workflow parser.

A directory needs a shard when it holds a test file, not when it is named test_*.
That drops the hardcoded test_configs exception and keeps fixture directories
like expected_fine_tuning_api out on their own merits.

The job keeps its name and its workflow, since assert-shard-coverage is a
required status check on litellm_internal_staging.

Verified red-first: a planted directory under tests/test_litellm, a planted
directory under tests/test_litellm/proxy, and a planted file under
tests/proxy_unit_tests each fail the guard, while a fixture-only directory does
not. 327 children across the three trees are assigned today.

.github/scripts/assert_ci_coverage.py

@@ -25,6 +25,15 @@ DOCKERFILE_TOKEN_RE = re.compile(r"[A-Za-z0-9_./-]*Dockerfile[A-Za-z0-9_.-]*")
 COMMENT_RE = re.compile(r"^\s*#.*$", re.MULTILINE)
 GLOB_CHARS = frozenset("*?")
 
+# Trees whose jobs are sharded with no catch-all bucket, so every child that holds
+# tests has to be named by some shard or it runs nowhere. A child listed here is
+# itself decomposed one level deeper and is checked through its own entry.
+SHARDED_ROOTS: tuple[str, ...] = (
+    "tests/proxy_unit_tests",
+    "tests/test_litellm",
+    "tests/test_litellm/proxy",
+)
+
 
 @dataclass(frozen=True, slots=True)
 class AllowEntry:
@@ -107,20 +116,32 @@ def _built_dockerfile_tokens(scalars: Iterable[Scalar]) -> frozenset[str]:
     )
 
 
-def _glob_to_regex(token: str) -> re.Pattern[str]:
+def _glob_to_regex(token: str, *, subtree: bool) -> re.Pattern[str]:
     parts = re.split(r"(\*\*/|\*\*|\*|\?)", token)
     translated = "".join(
         {"**/": r"(?:.*/)?", "**": r".*", "*": r"[^/]*", "?": r"[^/]"}.get(part, re.escape(part)) for part in parts
     )
-    return re.compile(rf"{translated}(?:/.*)?$")
+    return re.compile(rf"{translated}(?:/.*)?$" if subtree else rf"{translated}$")
 
 
 def _token_covers(token: str, relative_path: str) -> bool:
     if GLOB_CHARS & set(token):
-        return _glob_to_regex(token).match(relative_path) is not None
+        return _glob_to_regex(token, subtree=True).match(relative_path) is not None
     return relative_path == token or relative_path.startswith(f"{token}/")
 
 
+def _token_names(token: str, relative_path: str) -> bool:
+    """Whether the token names this path itself, rather than merely containing it.
+
+    A sharded tree has no catch-all bucket, so the ancestor token the census is happy
+    with (`tests/x` standing in for everything below it) is exactly what would let a
+    newly added child ride along without a shard.
+    """
+    if GLOB_CHARS & set(token):
+        return _glob_to_regex(token, subtree=False).match(relative_path) is not None
+    return token == relative_path
+
+
 def _test_files() -> tuple[str, ...]:
     return tuple(
         sorted(
@@ -166,6 +187,45 @@ def _describe(paths: tuple[str, ...]) -> str:
     return f"{len(paths)} test file(s) invoked by no job: {names}{suffix}"
 
 
+def _holds_tests(directory: pathlib.Path) -> bool:
+    return any(directory.rglob("test_*.py"))
+
+
+def _shard_children(root: str, repo_root: pathlib.Path = REPO_ROOT) -> tuple[str, ...]:
+    """Children of a sharded root that carry tests, so each one needs its own shard.
+
+    A directory earns an entry by containing a test file rather than by being named
+    `test_*`, which is what keeps fixture directories (`test_configs`, `expected_*`)
+    out without a hand-maintained list of exceptions.
+    """
+    return tuple(
+        sorted(
+            child.relative_to(repo_root).as_posix()
+            for child in (repo_root / root).iterdir()
+            if not child.name.startswith(".")
+            and (
+                _holds_tests(child)
+                if child.is_dir()
+                else child.name.startswith("test_") and child.suffix == ".py"
+            )
+        )
+    )
+
+
+def _unassigned_shard_children(
+    tokens: frozenset[str],
+    roots: tuple[str, ...] = SHARDED_ROOTS,
+    repo_root: pathlib.Path = REPO_ROOT,
+) -> tuple[Finding, ...]:
+    return tuple(
+        Finding(subject=child, detail=f"holds tests but no shard of {root} names it")
+        for root in roots
+        if (repo_root / root).is_dir()
+        for child in _shard_children(root, repo_root)
+        if child not in roots and not any(_token_names(token, child) for token in tokens)
+    )
+
+
 def _uncovered_dockerfiles(allowlist: Allowlist, tokens: frozenset[str]) -> tuple[Finding, ...]:
     return tuple(
         Finding(subject=relative_path, detail="built by no job")
@@ -229,7 +289,26 @@ def _report(title: str, findings: tuple[Finding, ...], remedy: str) -> None:
     _write("")
 
 
+def _check_shards() -> int:
+    findings = _unassigned_shard_children(_invoked_test_tokens(_all_scalars()))
+    if findings:
+        _report(
+            "test directories and files that no shard claims",
+            findings,
+            "Add each to the shard it belongs to. A directory that is itself split across "
+            "several shards belongs in SHARDED_ROOTS instead, so its own children get checked.",
+        )
+        return 1
+
+    counted = sum(len(_shard_children(root)) for root in SHARDED_ROOTS if (REPO_ROOT / root).is_dir())
+    _write(f"OK: all {counted} test children across {len(SHARDED_ROOTS)} sharded trees are assigned to a shard.")
+    return 0
+
+
 def main() -> int:
+    if "--shards" in sys.argv[1:]:
+        return _check_shards()
+
     allowlist = _load_allowlist()
     scalars = _all_scalars()
 

.github/workflows/test-unit-proxy-db.yml

@@ -42,11 +42,10 @@ concurrency:
 #     pinning the whole file to one worker (the default --dist=loadscope
 #     behavior for single-file targets).
 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.
+  # Fast guard — fails the workflow when a test directory or file inside a sharded
+  # tree is claimed by no shard. The semantic-shard design has no catch-all bucket,
+  # so an unassigned child runs nowhere; assert_ci_coverage.py holds the tree list
+  # and reads the same test-path keys the coverage census does.
   assert-shard-coverage:
     runs-on: ubuntu-latest
     timeout-minutes: 2
@@ -56,31 +55,8 @@ jobs:
       - 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
+      - name: Assert every test directory and file is claimed by a shard
+        run: python3 .github/scripts/assert_ci_coverage.py --shards
 
   proxy-db:
     needs: assert-shard-coverage

tests/test_litellm/test_assert_ci_coverage.py

@@ -0,0 +1,119 @@
+"""Tests for .github/scripts/assert_ci_coverage.py.
+
+Two guards share one workflow parser. The census asks whether a test file is run at
+all, so an ancestor path standing in for everything below it is a valid answer. The
+shard guard asks whether a sharded tree, which has no catch-all bucket, names each
+child outright, so that same ancestor path must NOT be an answer. The pair of
+matchers that splits those two questions is what these tests pin.
+"""
+
+import importlib.util
+import sys
+from pathlib import Path
+
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+_MODULE_PATH = _REPO_ROOT / ".github" / "scripts" / "assert_ci_coverage.py"
+_spec = importlib.util.spec_from_file_location("assert_ci_coverage", _MODULE_PATH)
+coverage = importlib.util.module_from_spec(_spec)
+sys.modules[_spec.name] = coverage  # @dataclass(slots=True) rebuilds via sys.modules
+_spec.loader.exec_module(coverage)
+
+
+def test_an_ancestor_directory_covers_a_file_but_does_not_name_it():
+    # The whole point of the split: `tests/x` answers "does it run?" but not
+    # "which shard owns it?" — accepting it for the latter is how a new child
+    # silently drops out of a tree that has no catch-all bucket.
+    assert coverage._token_covers("tests/test_litellm", "tests/test_litellm/caching") is True
+    assert coverage._token_names("tests/test_litellm", "tests/test_litellm/caching") is False
+
+
+def test_an_exact_token_both_covers_and_names():
+    assert coverage._token_covers("tests/test_litellm/caching", "tests/test_litellm/caching") is True
+    assert coverage._token_names("tests/test_litellm/caching", "tests/test_litellm/caching") is True
+
+
+def test_a_glob_names_only_what_it_matches_not_what_sits_below_it():
+    glob = "tests/test_litellm/test_*.py"
+    assert coverage._token_names(glob, "tests/test_litellm/test_router.py") is True
+    assert coverage._token_names(glob, "tests/test_litellm/test_router.py/nested.py") is False
+    assert coverage._token_names(glob, "tests/test_litellm/proxy/test_router.py") is False
+
+
+def test_a_glob_still_covers_the_subtree_for_the_census():
+    assert coverage._token_covers("tests/llm_translation/**", "tests/llm_translation/a/b.py") is True
+
+
+def test_a_directory_earns_a_shard_by_holding_tests_not_by_its_name(tmp_path):
+    fixtures = tmp_path / "expected_payloads"
+    fixtures.mkdir()
+    (fixtures / "body.json").write_text("{}")
+    tests = tmp_path / "some_area"
+    tests.mkdir()
+    (tests / "test_thing.py").write_text("def test_thing(): assert True\n")
+
+    assert coverage._holds_tests(fixtures) is False
+    assert coverage._holds_tests(tests) is True
+
+
+def test_shard_children_lists_test_dirs_and_test_files_and_skips_fixture_dirs(tmp_path):
+    root = tmp_path / "tests" / "tree"
+    (root / "billing").mkdir(parents=True)
+    (root / "billing" / "test_billing.py").write_text("def test_b(): assert True\n")
+    (root / "test_configs").mkdir()
+    (root / "test_configs" / "config.yaml").write_text("model_list: []\n")
+    (root / "test_top_level.py").write_text("def test_t(): assert True\n")
+    (root / "helpers.py").write_text("VALUE = 1\n")
+
+    assert coverage._shard_children("tests/tree", tmp_path) == (
+        "tests/tree/billing",
+        "tests/tree/test_top_level.py",
+    )
+
+
+def test_an_unnamed_child_is_reported_and_a_named_one_is_not(tmp_path):
+    root = tmp_path / "tests" / "tree"
+    (root / "claimed").mkdir(parents=True)
+    (root / "claimed" / "test_a.py").write_text("def test_a(): assert True\n")
+    (root / "orphan").mkdir()
+    (root / "orphan" / "test_b.py").write_text("def test_b(): assert True\n")
+
+    findings = coverage._unassigned_shard_children(
+        frozenset({"tests/tree/claimed"}), roots=("tests/tree",), repo_root=tmp_path
+    )
+
+    assert tuple(f.subject for f in findings) == ("tests/tree/orphan",)
+
+
+def test_the_parent_token_alone_does_not_satisfy_any_child(tmp_path):
+    root = tmp_path / "tests" / "tree"
+    (root / "billing").mkdir(parents=True)
+    (root / "billing" / "test_a.py").write_text("def test_a(): assert True\n")
+
+    findings = coverage._unassigned_shard_children(
+        frozenset({"tests/tree"}), roots=("tests/tree",), repo_root=tmp_path
+    )
+
+    assert tuple(f.subject for f in findings) == ("tests/tree/billing",)
+
+
+def test_a_child_that_is_itself_a_sharded_root_is_checked_there_not_here(tmp_path):
+    root = tmp_path / "tests" / "tree"
+    (root / "proxy" / "endpoints").mkdir(parents=True)
+    (root / "proxy" / "endpoints" / "test_a.py").write_text("def test_a(): assert True\n")
+
+    findings = coverage._unassigned_shard_children(
+        frozenset(), roots=("tests/tree", "tests/tree/proxy"), repo_root=tmp_path
+    )
+
+    assert tuple(f.subject for f in findings) == ("tests/tree/proxy/endpoints",)
+
+
+def test_every_sharded_root_named_in_the_script_exists_on_disk():
+    # A stale root would make the guard pass by checking nothing.
+    missing = [root for root in coverage.SHARDED_ROOTS if not (_REPO_ROOT / root).is_dir()]
+    assert missing == []
+
+
+def test_the_repo_as_it_stands_has_every_shard_child_assigned():
+    findings = coverage._unassigned_shard_children(coverage._invoked_test_tokens(coverage._all_scalars()))
+    assert [f.subject for f in findings] == []