perf(ci): fan the budget checkers out across cores (#37784)

* perf(ci): fan the budget checkers out across cores

check_type_discipline.py and check_test_quality.py each walk a few thousand
files and parse every one, single-threaded. In the lint job those two steps
measure 2.3 and 1.6 minutes, second and third behind dependency install, and
lint is the slowest required check on 9 of the last 10 merged staging PRs.

check_file is already pure per-file work, so the walk fans out over a process
pool with no change to what either rule reports. Callers sort, which is what
keeps output order stable when results land out of order. Runs below
PARALLEL_MIN_PATHS stay serial rather than pay for process startup, and the
worker count is capped so a large runner does not oversubscribe.

Measured locally over the same trees, output byte-identical both times:
type-discipline 17.8s -> 3.0s over litellm/ (78,768 report lines), test-quality
14.4s -> 2.3s over tests/ (6,321 report lines), per-rule counts unchanged.

* test(ci): type the fan-out helpers and skip the comparison on one core

scripts/check_test_quality.py

@@ -102,11 +102,13 @@ from __future__ import annotations
 
 import ast
 import io
+import os
 import re
 import sys
 import tokenize
 from collections.abc import Iterable, Iterator, Mapping, Sequence
 from dataclasses import dataclass
+from multiprocessing import Pool
 from pathlib import Path
 from types import MappingProxyType
 from typing import Final, NamedTuple
@@ -692,13 +694,36 @@ def collect_paths(raw: Iterable[str]) -> Iterator[Path]:
             yield candidate
 
 
+PARALLEL_MIN_PATHS = 200
+MAX_WORKERS = 8
+
+
+def _worker_count(path_count: int) -> int:
+    """1 when the run is too small to repay process startup, else one worker per
+    core up to MAX_WORKERS."""
+    if path_count < PARALLEL_MIN_PATHS:
+        return 1
+    return max(1, min(os.cpu_count() or 1, MAX_WORKERS))
+
+
+def scan_paths(paths: Sequence[Path]) -> tuple[Violation, ...]:
+    """check_file over every path. Pure per-file work, so it fans out across
+    processes; callers sort, which is what keeps output order stable."""
+    workers = _worker_count(len(paths))
+    if workers == 1:
+        return tuple(v for path in paths for v in check_file(path))
+    with Pool(workers) as pool:
+        return tuple(v for found in pool.imap_unordered(check_file, paths, chunksize=32) for v in found)
+
+
 def main(argv: Sequence[str]) -> int:
     paths: Final = tuple(a for a in argv if not a.startswith("-"))
     if not paths:
         print("usage: check_test_quality.py <files-or-dirs>...", file=sys.stderr)
         return 2
 
-    violations: Final = sorted(v for path in collect_paths(paths) for v in check_file(path))
+    targets: Final = tuple(collect_paths(paths))
+    violations: Final = sorted(scan_paths(targets))
     for violation in violations:
         print(violation.render())
 

scripts/check_type_discipline.py

@@ -114,10 +114,12 @@ from __future__ import annotations
  
 import ast
 import io
+import os
 import re
 import sys
 import tokenize
 from dataclasses import dataclass
+from multiprocessing import Pool
 from pathlib import Path
 from collections.abc import Iterable, Iterator, Mapping, Sequence
 from typing import NamedTuple
@@ -1070,13 +1072,36 @@ def collect_paths(raw: Iterable[str]) -> Iterator[Path]:
             yield p
  
  
+PARALLEL_MIN_PATHS = 200
+MAX_WORKERS = 8
+
+
+def _worker_count(path_count: int) -> int:
+    """1 when the run is too small to repay process startup, else one worker per
+    core up to MAX_WORKERS."""
+    if path_count < PARALLEL_MIN_PATHS:
+        return 1
+    return max(1, min(os.cpu_count() or 1, MAX_WORKERS))
+
+
+def scan_paths(paths: Sequence[Path]) -> tuple[Violation, ...]:
+    """check_file over every path. Pure per-file work, so it fans out across
+    processes; callers sort, which is what keeps output order stable."""
+    workers = _worker_count(len(paths))
+    if workers == 1:
+        return tuple(v for path in paths for v in check_file(path))
+    with Pool(workers) as pool:
+        return tuple(v for found in pool.imap_unordered(check_file, paths, chunksize=32) for v in found)
+
+
 def main(argv: Sequence[str]) -> int:
     paths = tuple(a for a in argv if not a.startswith("-"))
     if not paths:
         print("usage: check_type_discipline.py <files-or-dirs>...", file=sys.stderr)
         return 2
  
-    violations = sorted(v for path in collect_paths(paths) for v in check_file(path))
+    targets = tuple(collect_paths(paths))
+    violations = sorted(scan_paths(targets))
     for v in violations:
         print(v.render())
  

tests/test_litellm/test_check_test_quality.py

@@ -8,9 +8,13 @@ in the test body.
 """
 
 import importlib.util
+import os
+import subprocess
 import sys
 from pathlib import Path
 
+import pytest
+
 _REPO_ROOT = Path(__file__).resolve().parents[2]
 _MODULE_PATH = _REPO_ROOT / "scripts" / "check_test_quality.py"
 _spec = importlib.util.spec_from_file_location("check_test_quality", _MODULE_PATH)
@@ -548,3 +552,61 @@ def test_the_read_may_sit_a_statement_above_the_store(tmp_path):
 def test_a_loop_storing_under_a_key_that_is_not_the_loop_variable_is_not_an_inventory(tmp_path):
     source = _HELPER_DICT_CONFTEST.replace("state[attr] =", 'state["fixed"] =')
     assert [v.code for v in checker.check_file(_written(tmp_path, source))] == []
+
+
+_FANS_OUT = checker._worker_count(checker.PARALLEL_MIN_PATHS) > 1
+_SERIAL_ONLY = "one usable core, so scan_paths stays serial and there is no fan-out to compare"
+
+
+def _corpus(tmp_path: Path, count: int) -> tuple[Path, ...]:
+    for index in range(count):
+        (tmp_path / f"test_gen_{index}.py").write_text(
+            f"def test_flagged_{index}():\n    compute()\n\n\ndef test_clean_{index}():\n    assert compute() == {index}\n",
+            encoding="utf-8",
+        )
+    return tuple(sorted(tmp_path.rglob("*.py")))
+
+
+def _run_checker(target: Path) -> list[str]:
+    completed = subprocess.run(
+        [sys.executable, str(_MODULE_PATH), str(target)],
+        capture_output=True, text=True, timeout=300,
+    )
+    return completed.stdout.splitlines()
+
+
+def test_worker_count_stays_serial_below_the_threshold():
+    assert checker._worker_count(checker.PARALLEL_MIN_PATHS - 1) == 1
+
+
+def test_worker_count_fans_out_at_the_threshold():
+    assert checker._worker_count(checker.PARALLEL_MIN_PATHS) == max(
+        1, min(os.cpu_count() or 1, checker.MAX_WORKERS)
+    )
+
+
+def test_worker_count_never_exceeds_the_cap():
+    assert checker._worker_count(100_000) <= checker.MAX_WORKERS
+
+
+def test_scan_paths_below_the_threshold_returns_every_violation(tmp_path):
+    paths = _corpus(tmp_path, 3)
+    assert checker._worker_count(len(paths)) == 1
+    assert [v.code for v in checker.scan_paths(paths)] == ["TQ001"] * 3
+
+
+@pytest.mark.skipif(not _FANS_OUT, reason=_SERIAL_ONLY)
+def test_a_fanned_out_run_reports_exactly_what_a_serial_run_reports(tmp_path):
+    paths = _corpus(tmp_path, checker.PARALLEL_MIN_PATHS + 5)
+    serial = [v.render() for v in sorted(v for path in paths for v in checker.check_file(path))]
+    assert serial, "corpus must produce violations or the comparison proves nothing"
+    assert _run_checker(tmp_path) == serial
+
+
+@pytest.mark.skipif(not _FANS_OUT, reason=_SERIAL_ONLY)
+def test_a_fanned_out_run_reports_each_generated_file_exactly_once(tmp_path):
+    paths = _corpus(tmp_path, checker.PARALLEL_MIN_PATHS + 5)
+    reported = _run_checker(tmp_path)
+    assert len(reported) == len(paths)
+    assert len({line.split(":")[0] for line in reported}) == len(paths)
+    assert all(" TQ001 " in line for line in reported)

tests/test_litellm/test_check_type_discipline.py

@@ -8,10 +8,14 @@ a test fail. The comment-scanner cases are the regression for the readline path:
 
 import importlib.util
 import json
+import os
 import re
+import subprocess
 import sys
 from pathlib import Path
 
+import pytest
+
 _REPO_ROOT = Path(__file__).resolve().parents[2]
 _MODULE_PATH = _REPO_ROOT / "scripts" / "check_type_discipline.py"
 _spec = importlib.util.spec_from_file_location("check_type_discipline", _MODULE_PATH)
@@ -695,3 +699,61 @@ def test_budget_covers_exactly_the_checker_rules():
     for spec in budget.values():
         assert isinstance(spec["limit"], int)
         assert spec["limit"] >= 0
+
+
+_FANS_OUT = checker._worker_count(checker.PARALLEL_MIN_PATHS) > 1
+_SERIAL_ONLY = "one usable core, so scan_paths stays serial and there is no fan-out to compare"
+
+
+def _corpus(tmp_path: Path, count: int) -> tuple[Path, ...]:
+    for index in range(count):
+        (tmp_path / f"mod_{index}.py").write_text(
+            f"def build_{index}(items: list[int]) -> None:\n    return None\n",
+            encoding="utf-8",
+        )
+    return tuple(sorted(tmp_path.rglob("*.py")))
+
+
+def _run_checker(target: Path) -> list[str]:
+    completed = subprocess.run(
+        [sys.executable, str(_MODULE_PATH), str(target)],
+        capture_output=True, text=True, timeout=300,
+    )
+    return completed.stdout.splitlines()
+
+
+def test_worker_count_stays_serial_below_the_threshold():
+    assert checker._worker_count(checker.PARALLEL_MIN_PATHS - 1) == 1
+
+
+def test_worker_count_fans_out_at_the_threshold():
+    assert checker._worker_count(checker.PARALLEL_MIN_PATHS) == max(
+        1, min(os.cpu_count() or 1, checker.MAX_WORKERS)
+    )
+
+
+def test_worker_count_never_exceeds_the_cap():
+    assert checker._worker_count(100_000) <= checker.MAX_WORKERS
+
+
+def test_scan_paths_below_the_threshold_returns_every_violation(tmp_path):
+    paths = _corpus(tmp_path, 3)
+    assert checker._worker_count(len(paths)) == 1
+    found = checker.scan_paths(paths)
+    assert found and len({v.path for v in found}) == 3
+
+
+@pytest.mark.skipif(not _FANS_OUT, reason=_SERIAL_ONLY)
+def test_a_fanned_out_run_reports_exactly_what_a_serial_run_reports(tmp_path):
+    paths = _corpus(tmp_path, checker.PARALLEL_MIN_PATHS + 5)
+    serial = [v.render() for v in sorted(v for path in paths for v in checker.check_file(path))]
+    assert serial, "corpus must produce violations or the comparison proves nothing"
+    assert _run_checker(tmp_path) == serial
+
+
+@pytest.mark.skipif(not _FANS_OUT, reason=_SERIAL_ONLY)
+def test_a_fanned_out_run_reports_each_generated_file_exactly_once(tmp_path):
+    paths = _corpus(tmp_path, checker.PARALLEL_MIN_PATHS + 5)
+    reported = _run_checker(tmp_path)
+    assert reported
+    assert len({line.split(":")[0] for line in reported}) == len(paths)