ci-test-shard.py

#!/usr/bin/env python3
"""Run one shard of ci.yml's `test`, `test-release` or `coverage` job:
`cargo test --workspace` split over `COUNT` runners. Each runner compiles
and executes only its assigned integration targets, shortening both the
debug and release-mode jobs without changing their compilation profiles.

Shard 1 runs the unit tests (`--lib --bins`, ~4 min of the run on Windows)
and the doctests; the integration-test targets (`cargo metadata`, kind
`test`) are dealt out over the shards by name, with shard 1 taking a smaller
share to balance its unit tests. Every target lands in exactly one shard, so
the union of the shards is the old single `cargo test --workspace` run.

Windows debug CI sets CI_TEST_TIMINGS to the checked-in runtime sample.
Those shards balance measured test time plus linking work, with shard 1's
unit tests reserved up front. Other runners retain the count-based split.

CI_TEST_COVERAGE=1 runs the shard under cargo-llvm-cov for ci.yml's
`coverage` job: `cargo llvm-cov --no-report` in place of `cargo test`, and
no doctests, which cargo-llvm-cov leaves out without `--doctests`. A later
`cargo llvm-cov report` exports the shard's profiles.

Extra arguments after `SHARD COUNT` go to every `cargo` invocation.
Each invocation runs with `--no-fail-fast`; the exit status is non-zero if
any of them failed.

    python3 scripts/ci-test-shard.py 1 2
    python3 scripts/ci-test-shard.py 1 3 --locked --profile ci-release
"""

import json
import os
from pathlib import Path
import subprocess
import sys

# Shard 1's share of the integration targets, relative to every other
# shard's 1.0 (it also carries the unit tests and doctests).
FIRST_SHARD_WEIGHT = 0.5


def integration_targets(metadata):
    """Sorted, de-duplicated names of the workspace's integration tests."""
    members = set(metadata["workspace_members"])
    return sorted({t["name"] for p in metadata["packages"] if p["id"] in members
                   for t in p["targets"] if "test" in t["kind"]})


def partition(names, count, timings=None):
    """`count` lists covering `names` exactly once, in order, with the first
    list weighted by FIRST_SHARD_WEIGHT."""
    if count < 1:
        raise ValueError("count must be >= 1")
    if timings is not None:
        # Longest processing time first, reserving shard 1's unit-test cost.
        # Link overhead keeps a pile of fast tests from becoming a slow build.
        def cost(name):
            return timings["compile_seconds"] + timings["targets"].get(name, timings["default_seconds"])
        shards = [[] for _ in range(count)]
        load = [timings["unit_seconds"]] + [0.0] * (count - 1)
        for name in sorted(names, key=lambda n: (-cost(n), n)):
            i = min(range(count), key=lambda k: (load[k], k))
            shards[i].append(name)
            load[i] += cost(name)
        return [sorted(shard) for shard in shards]
    weights = [FIRST_SHARD_WEIGHT if count > 1 else 1.0] + [1.0] * (count - 1)
    shards = [[] for _ in range(count)]
    load = [0.0] * count
    for name in names:
        # The least-loaded shard relative to its weight; ties go to the lower index.
        i = min(range(count), key=lambda k: ((load[k] + 1) / weights[k], k))
        shards[i].append(name)
        load[i] += 1
    return shards


def invocations(shard, count, names, extra=(), timings=None, coverage=False):
    """The `cargo test` argument lists shard `shard` (1-based) runs."""
    if not 1 <= shard <= count:
        raise ValueError(f"shard {shard} not in 1..{count}")
    runner = ["llvm-cov", "--no-report"] if coverage else ["test"]
    base = ["cargo", *runner, "--workspace", "--no-fail-fast", *extra]
    mine = partition(names, count, timings)[shard - 1]
    runs = []
    if shard == 1:
        runs.append(base + ["--lib", "--bins"] + [a for n in mine for a in ("--test", n)])
        if not coverage:
            runs.append(base + ["--doc"])
    elif mine:
        runs.append(base + [a for n in mine for a in ("--test", n)])
    return runs


def main(argv):
    if len(argv) < 2:
        print(__doc__, file=sys.stderr)
        return 2
    shard, count = int(argv[0]), int(argv[1])
    metadata = json.loads(subprocess.run(
        ["cargo", "metadata", "--no-deps", "--format-version", "1"],
        check=True, capture_output=True, text=True).stdout)
    names = integration_targets(metadata)
    timing_path = os.environ.get("CI_TEST_TIMINGS")
    timings = json.loads(Path(timing_path).read_text(encoding="utf-8")) if timing_path else None
    coverage = os.environ.get("CI_TEST_COVERAGE") == "1"
    print(f"ci-test-shard: shard {shard}/{count}: "
          f"{len(partition(names, count, timings)[shard - 1])} of {len(names)} integration targets", flush=True)
    status = 0
    for args in invocations(shard, count, names, argv[2:], timings, coverage):
        print("+ " + " ".join(args), flush=True)
        if subprocess.run(args).returncode != 0:
            status = 1
    return status


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))