Update

[ghstack-poisoned]

Author: poteto

.github/workflows/runtime_test.yml

@@ -7,16 +7,31 @@ on:
     paths-ignore:
       - 'compiler/**'
 
+env:
+  # Number of workers (one per shard) to spawn
+  SHARD_COUNT: 5
+
 jobs:
   # Define the various test parameters and parallelism for this workflow
   build_test_params:
     name: Build test params
     runs-on: ubuntu-latest
     outputs:
       params: ${{ steps.define-params.outputs.result }}
-      # How many chunks to group tests into
-      parallelism: 5
+      shard_id: ${{ steps.define-chunks.outputs.result }}
     steps:
+      - uses: actions/github-script@v7
+        id: define-chunks
+        with:
+          script: |
+            function range(from, to) {
+              const arr = [];
+              for (let n = from; n < to; n++) {
+                arr.push(n);
+              }
+              return arr;
+            }
+            return range(0, process.env.SHARD_COUNT);
       - uses: actions/github-script@v7
         id: define-params
         with:
@@ -43,51 +58,15 @@ jobs:
               "-r=experimental --env=development --persistent"
             ];
 
-  # Chunk tests into groups for parallelism
-  chunk_tests:
-    name: Chunk tests
-    runs-on: ubuntu-latest
-    needs: build_test_params
-    strategy:
-      matrix:
-        params: ${{ fromJSON(needs.build_test_params.outputs.params) }}
-    continue-on-error: true
-    outputs:
-      chunks: ${{ steps.chunks.outputs.chunks }}
-      chunk_ids: ${{ steps.chunk_ids.outputs.chunk_ids }}
-    steps:
-      - uses: actions/checkout@v4
-      - uses: actions/setup-node@v4
-        with:
-          node-version: 18.x
-          cache: "yarn"
-          cache-dependency-path: yarn.lock
-      - name: Restore cached node_modules
-        uses: actions/cache@v4
-        id: node_modules
-        with:
-          path: "**/node_modules"
-          key: ${{ runner.arch }}-${{ runner.os }}-modules-${{ hashFiles('yarn.lock') }}
-      - run: yarn install --frozen-lockfile
-      - id: chunks
-        name: Set chunks
-        run: echo "chunks=$(yarn --silent test ${{ matrix.params }} --listTests --json --ci=github | jq -cM '[_nwise(length / ${{ needs.build_test_params.outputs.parallelism }} | ceil)]')" >> $GITHUB_OUTPUT
-      - id: chunk_ids
-        name: Set chunk IDs
-        run: echo "chunk_ids=$(echo $CHUNKS | jq -cM 'to_entries | map(.key)')" >> $GITHUB_OUTPUT
-        env:
-          CHUNKS: ${{ steps.chunks.outputs.chunks }}
-
-  # Spawn a job for each set of test param and number of chunks
+  # Spawn a job for each shard for a given set of test params
   test:
-    name: yarn test ${{ matrix.params }} (Chunk ${{ matrix.chunk_ids }})
+    name: yarn test ${{ matrix.params }} (Chunk ${{ matrix.shard_id }})
     runs-on: ubuntu-latest
-    needs: [build_test_params, chunk_tests]
+    needs: build_test_params
     strategy:
       matrix:
         params: ${{ fromJSON(needs.build_test_params.outputs.params) }}
-        chunks: ${{ fromJSON(needs.chunk_tests.outputs.chunks) }}
-        chunk_ids: ${{ fromJSON(needs.chunk_tests.outputs.chunk_ids) }}
+        shard_id: ${{ fromJSON(needs.build_test_params.outputs.shard_id) }}
     continue-on-error: true
     steps:
       - uses: actions/checkout@v4
@@ -103,6 +82,4 @@ jobs:
           path: "**/node_modules"
           key: ${{ runner.arch }}-${{ runner.os }}-modules-${{ hashFiles('yarn.lock') }}
       - run: yarn install --frozen-lockfile
-      - run: echo $CHUNKS | jq '.[${{ matrix.chunk_ids }}] | .[] | @text' | xargs yarn test ${{ matrix.params }} --ci=github
-        env:
-          CHUNKS: ${{ matrix.chunks }}
+      - run: yarn test ${{ matrix.params }} --ci=github --shard=${{ matrix.shard_id }}/${{ env.SHARD_COUNT }}