Implement matrix sharding

astral-sh · Jan 10, 2025 · a69bb3d · a69bb3d
1 parent e985fa1
commit a69bb3d
Show file tree

Hide file tree

Showing 2 changed files with 118 additions and 13 deletions.
diff --git a/.github/workflows/linux.yml b/.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,14 +172,88 @@ jobs:
             pythonbuild:
               - "src/*.rs"
 
-  build:
+  build-0:
+    needs:
+      - generate-matrix
+      - pythonbuild
+      - image
+    runs-on: ${{ matrix.runner }}
+    strategy:
+      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) }}
+      matrix: ${{ fromJson(needs.generate-matrix.outputs.matrix-0) }}
       fail-fast: false
     name: ${{ matrix.target_triple }} / ${{ matrix.python }} / ${{ matrix.build_options }}
     steps:

diff --git a/ci-matrix.py b/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,28 @@ 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))