From bf189daea9b632ba37a43ce06cd3f7d46c768e3a Mon Sep 17 00:00:00 2001 From: Stephen Zhou Date: Sun, 20 Sep 2026 13:42:38 +0000 Subject: [PATCH] ci(api): balance unit shards with shared test timings (#42593) --- .github/workflows/api-tests.yml | 63 +++++++++++ api/conftest.py | 28 ++++- api/dev/pytest_sharding.md | 59 ++++++++++ api/dev/pytest_sharding.py | 80 +++++++++++++- api/tests/pytest_timing.py | 26 +++++ api/tests/unit_tests/test_pytest_sharding.py | 109 +++++++++++++++++++ api/tests/unit_tests/test_pytest_timing.py | 36 ++++++ 7 files changed, 397 insertions(+), 4 deletions(-) create mode 100644 api/dev/pytest_sharding.md create mode 100644 api/tests/pytest_timing.py create mode 100644 api/tests/unit_tests/test_pytest_timing.py diff --git a/.github/workflows/api-tests.yml b/.github/workflows/api-tests.yml index ce3ec6c10ce..5544ef470fe 100644 --- a/.github/workflows/api-tests.yml +++ b/.github/workflows/api-tests.yml @@ -14,7 +14,52 @@ concurrency: cancel-in-progress: true jobs: + api-unit-plan: + name: Plan API unit shards + runs-on: depot-ubuntu-24.04 + steps: + - name: Checkout code + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + persist-credentials: false + - name: Setup UV and Python + uses: astral-sh/setup-uv@bec219d24cd3e171d82865faccec33120bb574f4 # v10.1.0 + with: + python-version: '3.12' + enable-cache: false + - name: Download shared timing baseline + # Public statistics are independent of PR refs and require no credentials. + run: | + if ! curl --fail --silent --show-error --location \ + --connect-timeout 5 --max-time 15 --retry 2 --retry-max-time 30 \ + https://raw.githubusercontent.com/langgenius/dify-test-infra/generated-stats/stats/test-times.json \ + --output unit-test-durations.json; then + rm -f unit-test-durations.json + echo '::warning::Timing baseline unavailable; using round-robin file allocation.' + fi + - name: Freeze file assignments for every shard + run: | + uv run --project api --no-sync python api/dev/pytest_sharding.py \ + --shard-index 1 --shard-total 2 \ + --durations unit-test-durations.json \ + --write-plan unit-test-plan.json \ + --split-threshold 60 \ + --ignore api/tests/unit_tests/controllers \ + api/tests/unit_tests \ + api/providers/vdb/*/tests/unit_tests \ + api/providers/trace/*/tests/unit_tests + - name: Upload shared plan + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: api-unit-plan + path: | + unit-test-plan.json + unit-test-durations.json + retention-days: 7 + if-no-files-found: error + api-unit: + needs: api-unit-plan name: API Unit Tests (${{ matrix.python-version }}, shard ${{ matrix.shardIndex }}/${{ matrix.shardTotal }}) runs-on: depot-ubuntu-24.04 env: @@ -53,9 +98,15 @@ jobs: - name: Install dependencies run: uv sync --project api --dev + - name: Download shared plan + uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 + with: + name: api-unit-plan + - name: Run Unit Tests run: | uv run --project api python api/dev/pytest_sharding.py \ + --plan unit-test-plan.json \ --shard-index "${{ matrix.shardIndex }}" \ --shard-total "${{ matrix.shardTotal }}" \ --ignore api/tests/unit_tests/controllers \ @@ -68,6 +119,10 @@ jobs: -n auto \ -o 'addopts=--cov=./api --cov-branch --import-mode=importlib' \ --cov-report= \ + --file-shard-plan unit-test-plan.json \ + --shard-index "${{ matrix.shardIndex }}" \ + --shard-total "${{ matrix.shardTotal }}" \ + --write-test-durations "test-results/durations-${{ matrix.shardIndex }}.json" \ --junitxml=test-results/unit.xml \ "@$RUNNER_TEMP/unit-test-files.txt" @@ -80,6 +135,14 @@ jobs: if-no-files-found: warn retention-days: 7 + - name: Upload successful timing observations + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: api-unit-durations-${{ matrix.shardIndex }} + path: test-results/durations-*.json + retention-days: 7 + if-no-files-found: error + - name: Upload unit coverage data uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: diff --git a/api/conftest.py b/api/conftest.py index a008a020771..cf50a505e6f 100644 --- a/api/conftest.py +++ b/api/conftest.py @@ -34,6 +34,8 @@ ensure_backend_test_environment(_REPO_ROOT) def pytest_addoption(parser: pytest.Parser) -> None: group = parser.getgroup("dify") + group.addoption("--file-shard-plan", type=Path, help="Shared pre-collection file/partition assignment.") + group.addoption("--write-test-durations", type=Path, help="Write successful-run file durations for CI history.") group.addoption( "--middleware-stop-timeout", type=int, @@ -87,13 +89,37 @@ def pytest_configure(config: pytest.Config) -> None: if stop_timeout is not None and stop_timeout < 0: raise pytest.UsageError("--middleware-stop-timeout must be nonnegative") + durations_path = config.getoption("write_test_durations") + if durations_path is not None and not hasattr(config, "workerinput"): + from tests.pytest_timing import TestDurationsPlugin + + config.pluginmanager.register(TestDurationsPlugin(durations_path), "dify-test-durations") + config.stash[_DIFY_COMPOSE_STACKS_KEY] = [] def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> None: - """Select a deterministic, evenly sized slice of collected tests for this shard.""" + """Apply the shared file plan, or the legacy round-robin testcase shard.""" shard_index = config.getoption("shard_index") shard_total = config.getoption("shard_total") + if plan_path := config.getoption("file_shard_plan"): + import json + + from dev.pytest_sharding import case_shard + + plan = json.loads(plan_path.read_text()) + selected_items: list[pytest.Item] = [] + deselected_items: list[pytest.Item] = [] + for item in items: + filename = item.path.relative_to(_REPO_ROOT).as_posix() + # Missing files are a broken plan, not an excuse to silently drop tests. + targets = plan[filename] + case = item.nodeid.split("::", 1)[1] + destination = selected_items if case_shard(case, targets) == shard_index else deselected_items + destination.append(item) + config.hook.pytest_deselected(items=deselected_items) + items[:] = selected_items + return if shard_total == 1: return diff --git a/api/dev/pytest_sharding.md b/api/dev/pytest_sharding.md new file mode 100644 index 00000000000..d75751aacd8 --- /dev/null +++ b/api/dev/pytest_sharding.md @@ -0,0 +1,59 @@ +# API unit-test sharding + +The workflow plans shards before pytest collection. Files with recorded durations +are sorted longest first and assigned to the least-loaded shard. Unseen files are +assigned round robin. File discovery always comes from the current checkout, so +stale history cannot omit new tests or reintroduce deleted files. + +Like PyTorch's `tools/testing/test_selections.py`, oversized files become logical +parts. Our default is 60 summed testcase seconds per part, capped at the number of +CI shards. This is worker work, including setup/call/teardown, not wall time; the +threshold corresponds roughly to 30 seconds of work on two workers. Each case's +SHA-256 hash assigns it to one part, including new parametrized cases. Only files +assigned to a runner are collected; shared oversized files are collected on each +assigned runner. All workers use the same frozen plan. Ordinary xdist scheduling +continues inside each runner. + +Successful coordinator reports are uploaded as two small timing artifacts. The +public [dify-test-infra](https://github.com/langgenius/dify-test-infra) repository +collects complete observations from successful CI on merged Dify PRs daily at +03:05 UTC. It averages up to five recent samples and publishes `stats/test-times.json` +on its `generated-stats` branch. Metadata records the source runs. No timing data +or generated history is committed to Dify itself. + +Every PR, including fork PRs, downloads that same public baseline without a token. +Unmerged PR observations do not update the shared baseline automatically. The +initial baseline is explicitly marked as a bootstrap from validated CI on #42593; +automatic updates use merged PRs only. Keeping the producer in the statistics +repository avoids needing a cross-repository write token in Dify or additional +main-branch test runs. + +A single planning job downloads the baseline once and publishes both the input +and assignment plan as run artifacts. All shards therefore use the same snapshot, +even if statistics are updated during a run. Downloads have bounded timeouts and +retries; missing or invalid JSON falls back to round-robin allocation. Successful +unit runs continue publishing observations for the external updater; failed or +incomplete runs are not accepted as baseline samples. + +This trades an extra planning job for consistent assignments. Timing noise, +shared fixtures, collection, and an indivisible slow case can still produce +imbalance. Logical splitting can duplicate expensive module/session setup. Compare both +shard test steps and the end-to-end API completion time when evaluating a change. + +Local examples (from repository root): + +```sh +uv run --project api python api/dev/pytest_sharding.py \ + --shard-index 1 --shard-total 2 --durations /tmp/durations.json \ + --write-plan /tmp/plan.json --ignore api/tests/unit_tests/controllers \ + api/tests/unit_tests api/providers/vdb/*/tests/unit_tests api/providers/trace/*/tests/unit_tests +uv run --project api python api/dev/pytest_sharding.py \ + --shard-index 1 --shard-total 2 --plan /tmp/plan.json api > /tmp/files.txt +uv run --project api pytest --file-shard-plan /tmp/plan.json \ + --shard-index 1 --shard-total 2 --write-test-durations /tmp/durations-1.json \ + @/tmp/files.txt +``` + +Repeat the last two commands for shard 2. Merge their file durations by addition +before using them as the next history snapshot. Keep generated files outside the +checkout locally. diff --git a/api/dev/pytest_sharding.py b/api/dev/pytest_sharding.py index df6ebf6eda0..2b33f2e8e85 100644 --- a/api/dev/pytest_sharding.py +++ b/api/dev/pytest_sharding.py @@ -1,11 +1,68 @@ """Select pytest files before collection so workers do not import other shards.""" import argparse +import hashlib +import json +import math import sys from collections.abc import Sequence from pathlib import Path +def load_durations(path: Path) -> dict[str, float]: + """History is optional and never controls which files are discovered.""" + try: + data = json.loads(path.read_text()) + except (OSError, ValueError): + return {} + if not isinstance(data, dict): + return {} + return { + name: float(value) + for name, value in data.items() + if isinstance(value, (int, float)) and not isinstance(value, bool) and math.isfinite(value) and value >= 0 + } + + +def build_plan( + files: Sequence[Path], durations: dict[str, float], total: int, threshold: float +) -> dict[str, list[int]]: + """Longest estimated work first; split oversized files into logical parts. + + Each list maps a file's hash partitions to one-based CI shard numbers. + All CI shards must consume the same plan. Unknown files use round robin. + Estimates are summed worker seconds, not predicted job wall time. + """ + if total < 1 or not math.isfinite(threshold) or threshold <= 0: + raise ValueError("shard-total and split threshold must be positive") + plan: dict[str, list[int]] = {} + tasks: list[tuple[float, str, int]] = [] + unknown: list[str] = [] + for path in sorted(set(files)): + name = path.as_posix() + duration = durations.get(name) + if duration is None: + unknown.append(name) + continue + parts = min(total, max(1, math.ceil(duration / threshold))) + plan[name] = [0] * parts + tasks.extend((duration / parts, name, part) for part in range(parts)) + loads = [0.0] * total + for duration, name, part in sorted(tasks, key=lambda task: (-task[0], task[1], task[2])): + target = min(range(total), key=lambda index: (loads[index], index)) + plan[name][part] = target + 1 + loads[target] += duration + for index, name in enumerate(unknown): + plan[name] = [index % total + 1] + return plan + + +def case_shard(nodeid: str, targets: Sequence[int]) -> int: + """Stable across processes, including newly added parametrized cases.""" + digest = hashlib.sha256(nodeid.encode()).digest() + return targets[int.from_bytes(digest[:8], "big") % len(targets)] + + def select_test_files( roots: Sequence[Path], *, shard_index: int, shard_total: int, ignored: Sequence[Path] = () ) -> list[Path]: @@ -29,11 +86,28 @@ def main() -> None: parser.add_argument("--shard-index", type=int, required=True) parser.add_argument("--shard-total", type=int, required=True) parser.add_argument("--ignore", type=Path, action="append", default=[]) + parser.add_argument("--durations", type=Path) + parser.add_argument("--write-plan", type=Path) + parser.add_argument("--plan", type=Path) + parser.add_argument("--split-threshold", type=float, default=60, help="Summed testcase seconds per logical part.") parser.add_argument("roots", type=Path, nargs="+") args = parser.parse_args() - files = select_test_files( - args.roots, shard_index=args.shard_index, shard_total=args.shard_total, ignored=args.ignore - ) + if args.write_plan: + files = select_test_files(args.roots, shard_index=1, shard_total=1, ignored=args.ignore) + durations = load_durations(args.durations) if args.durations else {} + plan = build_plan(files, durations, args.shard_total, args.split_threshold) + args.write_plan.write_text(json.dumps(plan, sort_keys=True) + "\n") + sys.stderr.write( + f"Planned {len(files)} files; {sum(name in durations for name in plan)} have historical durations.\n" + ) + return + if args.plan: + plan = json.loads(args.plan.read_text()) + files = [Path(name) for name, targets in sorted(plan.items()) if args.shard_index in targets] + else: + files = select_test_files( + args.roots, shard_index=args.shard_index, shard_total=args.shard_total, ignored=args.ignore + ) if not files: parser.error("No test files selected; refusing to fall back to pytest's default discovery") sys.stdout.write("\n".join(path.as_posix() for path in files) + "\n") diff --git a/api/tests/pytest_timing.py b/api/tests/pytest_timing.py new file mode 100644 index 00000000000..6740aeecb63 --- /dev/null +++ b/api/tests/pytest_timing.py @@ -0,0 +1,26 @@ +"""Successful-run timing observations for shared CI shard planning.""" + +from pathlib import Path + +import pytest + + +class TestDurationsPlugin: + """Persist coordinator reports only after a complete successful test run.""" + + def __init__(self, path: Path) -> None: + self.path = path + self.files: dict[str, float] = {} + + def pytest_runtest_logreport(self, report: pytest.TestReport) -> None: + filename = report.nodeid.split("::", 1)[0] + if not filename.startswith("api/"): + filename = f"api/{filename}" + self.files[filename] = self.files.get(filename, 0.0) + report.duration + + def pytest_sessionfinish(self, session: pytest.Session, exitstatus: int) -> None: + import json + + if exitstatus == 0: + self.path.parent.mkdir(parents=True, exist_ok=True) + self.path.write_text(json.dumps(self.files, sort_keys=True) + "\n") diff --git a/api/tests/unit_tests/test_pytest_sharding.py b/api/tests/unit_tests/test_pytest_sharding.py index e81d57e07a3..03950923734 100644 --- a/api/tests/unit_tests/test_pytest_sharding.py +++ b/api/tests/unit_tests/test_pytest_sharding.py @@ -49,3 +49,112 @@ def test_empty_shard_fails_instead_of_running_default_discovery( captured = capsys.readouterr() assert captured.out == "" assert "No test files selected" in captured.err + + +def test_duration_plan_balances_work_and_ignores_stale_history() -> None: + from dev.pytest_sharding import build_plan + + files = [Path(f"test_{name}.py") for name in "abcd"] + durations = dict(zip((str(path) for path in files), [90.0, 80.0, 50.0, 40.0], strict=True)) + durations["deleted.py"] = 999.0 + plan = build_plan(files, durations, total=2, threshold=120) + assert set(plan) == {str(path) for path in files} + loads = [sum(durations[name] for name, targets in plan.items() if targets == [i]) for i in (1, 2)] + assert loads == [130, 130] + assert plan == build_plan(list(reversed(files)), durations, total=2, threshold=120) + + +def test_logical_shards_cover_existing_and_new_cases_exactly_once() -> None: + from dev.pytest_sharding import build_plan, case_shard + + plan = build_plan([Path("test_large.py")], {"test_large.py": 250.0}, total=3, threshold=100) + targets = plan["test_large.py"] + assert sorted(targets) == [1, 2, 3] + cases = [f"TestLarge::test_case[param-{i}]" for i in range(100)] + ["test_new"] + partitions = [{case for case in cases if case_shard(case, targets) == shard} for shard in (1, 2, 3)] + assert set.union(*partitions) == set(cases) + assert sum(map(len, partitions)) == len(cases) + assert all(partitions) + + +def test_missing_or_corrupt_history_keeps_all_files(tmp_path: Path) -> None: + from dev.pytest_sharding import build_plan, load_durations + + path = tmp_path / "durations.json" + assert load_durations(path) == {} + for invalid in ["broken", "[]", "null"]: + path.write_text(invalid) + assert load_durations(path) == {} + path.write_text('{"good": 2, "negative": -1, "nan": NaN, "bool": true, "string": "2"}') + assert load_durations(path) == {"good": 2.0} + files = [Path(f"test_{i}.py") for i in range(5)] + assert list(build_plan(files, {}, total=2, threshold=60).values()) == [[1], [2], [1], [2], [1]] + + +def test_plan_cli_selects_only_files_assigned_to_shard( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] +) -> None: + import json + + plan = tmp_path / "plan.json" + plan.write_text(json.dumps({"api/test_a.py": [1], "api/test_b.py": [2], "api/test_big.py": [1, 2]})) + monkeypatch.setattr( + "sys.argv", ["pytest_sharding.py", "--shard-index", "2", "--shard-total", "2", "--plan", str(plan), "api"] + ) + main() + assert capsys.readouterr().out.splitlines() == ["api/test_b.py", "api/test_big.py"] + + +@pytest.mark.parametrize("with_history", [False, True]) +def test_cli_freezes_current_files_and_every_shard_consumes_the_same_plan( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str], with_history: bool +) -> None: + import json + + root = tmp_path / "tests" + root.mkdir() + slow = root / "test_slow.py" + new = root / "test_new.py" + for path in (slow, new, root / "helper.py"): + path.touch() + history = tmp_path / "history.json" + history.write_text(json.dumps({str(slow): 100, str(root / "test_deleted.py"): 1000})) + plan_path = tmp_path / "plan.json" + arguments = ["pytest_sharding.py", "--shard-index", "1", "--shard-total", "2", "--write-plan", str(plan_path)] + if with_history: + arguments.extend(["--durations", str(history)]) + monkeypatch.setattr("sys.argv", [*arguments, str(root)]) + main() + assert capsys.readouterr().out == "" + plan = json.loads(plan_path.read_text()) + assert set(plan) == {str(slow), str(new)} + if with_history: + assert sorted(plan[str(slow)]) == [1, 2] + selected: set[str] = set() + for shard in (1, 2): + monkeypatch.setattr( + "sys.argv", + [ + "pytest_sharding.py", + "--shard-index", + str(shard), + "--shard-total", + "2", + "--plan", + str(plan_path), + str(root), + ], + ) + main() + files = capsys.readouterr().out.splitlines() + assert set(files) == {name for name, targets in plan.items() if shard in targets} + selected.update(files) + assert selected == {str(slow), str(new)} + + +@pytest.mark.parametrize("threshold", [0, -1, float("nan")]) +def test_invalid_logical_split_threshold_is_rejected(threshold: float) -> None: + from dev.pytest_sharding import build_plan + + with pytest.raises(ValueError, match="positive"): + build_plan([Path("test_a.py")], {"test_a.py": 10.0}, total=2, threshold=threshold) diff --git a/api/tests/unit_tests/test_pytest_timing.py b/api/tests/unit_tests/test_pytest_timing.py new file mode 100644 index 00000000000..a0872f9d2e9 --- /dev/null +++ b/api/tests/unit_tests/test_pytest_timing.py @@ -0,0 +1,36 @@ +"""Persist complete successful observations for shared test timing history.""" + +from pathlib import Path +from typing import Literal +from unittest.mock import MagicMock + +import pytest + + +def report(nodeid: str, when: Literal["setup", "call", "teardown"], duration: float) -> pytest.TestReport: + return pytest.TestReport( + nodeid=nodeid, + location=(nodeid.split("::")[0], 0, "test_case"), + keywords={}, + outcome="passed", + longrepr=None, + when=when, + duration=duration, + ) + + +@pytest.mark.parametrize("exitstatus", [0, 1, 2, 3, 4, 5]) +def test_history_is_published_only_for_complete_success(tmp_path: Path, exitstatus: int) -> None: + import json + + from tests.pytest_timing import TestDurationsPlugin + + destination = tmp_path / "history.json" + plugin = TestDurationsPlugin(destination) + for phase in ("setup", "call", "teardown"): + plugin.pytest_runtest_logreport(report("tests/test_file.py::test_a", phase, 2.0)) + plugin.pytest_sessionfinish(MagicMock(spec=pytest.Session), exitstatus) + if exitstatus == 0: + assert json.loads(destination.read_text()) == {"api/tests/test_file.py": 6.0} + else: + assert not destination.exists()