Implement matrix sharding (astral-sh#492)

Since we've encountered the GitHub matrix size limit of 256 entries a
couple times now (recently in astral-sh#388), this adds sharding to our matrix
generation script.

There's now a `--max-shards` option which splits the matrix into (up to)
the requested number of shards. We don't split it evenly, we eagerly
fill each shard first. To consume the shards, we just copy the `build`
job into two identical child jobs which consume each matrix shard.
Ideally, we'd avoid code duplication here by using a re-usable action or
something for the build-steps — but that's overkill at this point.

We only need this for Linux for the foreseeable future.

Author: zanieb

.github/workflows/linux.yml

@@ -125,7 +125,8 @@ jobs:
   generate-matrix:
     runs-on: ubuntu-latest
     outputs:
-      matrix: ${{ steps.set-matrix.outputs.matrix }}
+      matrix-0: ${{ steps.set-matrix.outputs.matrix-0 }}
+      matrix-1: ${{ steps.set-matrix.outputs.matrix-1 }}
       any_builds: ${{ steps.set-matrix.outputs.any_builds }}
       pythonbuild_changed: ${{ steps.changed.outputs.pythonbuild_any_changed }}
     steps:
@@ -144,11 +145,19 @@ jobs:
       - name: Generate build matrix
         id: set-matrix
         run: |
-          uv run ci-matrix.py --platform linux --labels '${{ steps.get-labels.outputs.labels }}' > matrix.json && echo "matrix=$(cat matrix.json)" >> $GITHUB_OUTPUT
+          uv run ci-matrix.py \
+            --platform linux \
+            --labels '${{ steps.get-labels.outputs.labels }}' \
+            --max-shards 2 \
+            > matrix.json
+
+          echo "matrix-0=$(jq -c '.["0"]' matrix.json)" >> $GITHUB_OUTPUT
+          echo "matrix-1=$(jq -c '.["1"]' matrix.json)" >> $GITHUB_OUTPUT
+
           # Display the matrix for debugging too
           cat matrix.json | jq
 
-          if jq -e '.include | length > 0' matrix.json > /dev/null; then
+          if jq -e '.["0"].include | length > 0' matrix.json > /dev/null; then
             # Build matrix has entries
             echo "any_builds=true" >> $GITHUB_OUTPUT
           else
@@ -163,7 +172,7 @@ jobs:
             pythonbuild:
               - "src/*.rs"
 
-  build:
+  build-0:
     needs:
       - generate-matrix
       - pythonbuild
@@ -174,7 +183,81 @@ jobs:
       attestations: write
     runs-on: ${{ matrix.runner }}
     strategy:
-      matrix: ${{ fromJson(needs.generate-matrix.outputs.matrix) }}
+      matrix: ${{ fromJson(needs.generate-matrix.outputs.matrix-0) }}
+      fail-fast: false
+    name: ${{ matrix.target_triple }} / ${{ matrix.python }} / ${{ matrix.build_options }}
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - name: Install Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+
+      - name: Download pythonbuild
+        uses: actions/download-artifact@v4
+        with:
+          name: pythonbuild
+          path: build
+
+      - name: Download images
+        uses: actions/download-artifact@v4
+        with:
+          pattern: image-*
+          path: build
+          merge-multiple: true
+
+      - name: Load Docker Images
+        run: |
+          for f in build/image-*.tar.zst; do
+            echo "decompressing $f"
+            zstd -d --rm ${f}
+          done
+
+          for f in build/image-*.tar; do
+            echo "loading $f"
+            docker load --input $f
+          done
+
+      - name: Build
+        if: ${{ ! matrix.dry-run }}
+        run: |
+          # Do empty target so all generated files are touched.
+          ./build-linux.py --make-target empty
+
+          # Touch mtimes of all images so they are newer than autogenerated files above.
+          touch build/image-*
+
+          ./build-linux.py --target-triple ${{ matrix.target_triple }} --python cpython-${{ matrix.python }} --options ${{ matrix.build_options }}
+
+      - name: Validate Distribution
+        if: ${{ ! matrix.dry-run }}
+        run: |
+          chmod +x build/pythonbuild
+
+          if [ "${{ matrix.run }}" == "true" ]; then
+            EXTRA_ARGS="--run"
+          fi
+
+          build/pythonbuild validate-distribution ${EXTRA_ARGS} dist/*.tar.zst
+
+      - name: Upload Distribution
+        if: ${{ ! matrix.dry-run }}
+        uses: actions/upload-artifact@v4
+        with:
+          name: cpython-${{ matrix.python }}-${{ matrix.target_triple }}-${{ matrix.build_options }}
+          path: dist/*
+
+  build-1:
+    needs:
+      - generate-matrix
+      - pythonbuild
+      - image
+    runs-on: ${{ matrix.runner }}
+    strategy:
+      matrix: ${{ fromJson(needs.generate-matrix.outputs.matrix-1) }}
       fail-fast: false
     name: ${{ matrix.target_triple }} / ${{ matrix.python }} / ${{ matrix.build_options }}
     steps:

ci-matrix.py

@@ -8,6 +8,7 @@
 
 import argparse
 import json
+import sys
 from typing import Any, Optional
 
 import yaml
@@ -16,6 +17,7 @@
 CI_TARGETS_YAML = "ci-targets.yaml"
 CI_RUNNERS_YAML = "ci-runners.yaml"
 CI_EXTRA_SKIP_LABELS = ["documentation"]
+CI_MATRIX_SIZE_LIMIT = 256  # The maximum size of a matrix in GitHub Actions
 
 
 def meets_conditional_version(version: str, min_version: str) -> bool:
@@ -216,6 +218,12 @@ def parse_args() -> argparse.Namespace:
         choices=["darwin", "linux", "windows"],
         help="Filter matrix entries by platform",
     )
+    parser.add_argument(
+        "--max-shards",
+        type=int,
+        default=0,
+        help="The maximum number of shards allowed; set to zero to disable ",
+    )
     parser.add_argument(
         "--labels",
         help="Comma-separated list of labels to filter by (e.g., 'platform:darwin,python:3.13,build:debug'), all must match.",
@@ -246,14 +254,34 @@ def main() -> None:
             if runner_config.get("free")
         }
 
-    matrix = {
-        "include": generate_matrix_entries(
-            config,
-            runners,
-            args.platform,
-            labels,
-        )
-    }
+    entries = generate_matrix_entries(
+        config,
+        runners,
+        args.platform,
+        labels,
+    )
+
+    if args.max_shards:
+        matrix = {}
+        shards = (len(entries) // CI_MATRIX_SIZE_LIMIT) + 1
+        if shards > args.max_shards:
+            print(
+                f"error: matrix of size {len(entries)} requires {shards} shards, but the maximum is {args.max_shards}; consider increasing `--max-shards`",
+                file=sys.stderr,
+            )
+            sys.exit(1)
+        for shard in range(args.max_shards):
+            shard_entries = entries[
+                shard * CI_MATRIX_SIZE_LIMIT : (shard + 1) * CI_MATRIX_SIZE_LIMIT
+            ]
+            matrix[str(shard)] = {"include": shard_entries}
+    else:
+        if len(entries) > CI_MATRIX_SIZE_LIMIT:
+            print(
+                f"warning: matrix of size {len(entries)} exceeds limit of {CI_MATRIX_SIZE_LIMIT} but sharding is not enabled; consider setting `--max-shards`",
+                file=sys.stderr,
+            )
+        matrix = {"include": entries}
 
     print(json.dumps(matrix))