ci(api): balance unit shards with shared test timings (#42593)

.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:

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
 

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.

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")

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")

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)

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()