perf(ci): balance internal tests across two runners

Balance internal test files across packages and split the workload
across two runners to reduce CI waits. Preserve package test
configuration, refresh duration estimates, and combine all shard reports
for timing updates.

Mono-RevId: f43f740e7a8db6a812a157a8a02b2a15a7ac66fa

Author: ericallam

.github/workflows/pr_checks.yml

@@ -2,7 +2,7 @@ name: 🤖 PR Checks
 
 on:
   pull_request:
-    types: [opened, synchronize, reopened]
+    types: [opened, synchronize, reopened, ready_for_review]
   merge_group:
 
 concurrency:
@@ -91,6 +91,10 @@ jobs:
               - 'pnpm-lock.yaml'
               - 'pnpm-workspace.yaml'
               - 'turbo.json'
+              - 'vitest.internal.config.ts'
+              - 'vitest.internal-reports.config.mjs'
+              - 'internal-test-projects.mts'
+              - 'test-timings.json'
             # The whole webapp app tree, not just its routes, and that is the whole reason this
             # filter exists. Two tests in @internal/observability-map read it: integration.test.ts
             # scans the live route tree, and webappSymbols.test.ts walks all of apps/webapp/app and
@@ -110,8 +114,8 @@ jobs:
             # Deliberately NOT here: this package's own paths, and packages/plugins/src and
             # internal-packages/rbac/src, the other two trees webappSymbols.test.ts reads.
             # `internal` above already matches `internal-packages/**` and the shared packages, and
-            # `unit-tests-internal.yml` runs `turbo run test --filter "@internal/*"`, which picks up
-            # @internal/observability-map and runs the same vitest suite. Listing them here as well
+            # `unit-tests-internal.yml` loads every @internal/* test project through
+            # vitest.internal.config.ts, including the same observability-map suite. Listing them here as well
             # ran the suite twice on every PR touching them, which was this filter's own doing.
             #
             # Also deliberately NOT here: pr_checks.yml, package.json, pnpm-lock.yaml,

.github/workflows/unit-tests-internal.yml

@@ -14,14 +14,18 @@ on:
 jobs:
   unitTests:
     name: "🧪 Unit Tests: Internal"
-    # Single big machine instead of a 12-job matrix: the internal suites are serial
-    # (fileParallelism: false) and container-wait-bound, so 12 in-machine shard processes
-    # fit comfortably in 32 vCPUs while paying the setup cost (install, prisma generate,
-    # image pulls) once instead of 12 times.
+    # Two Docker daemons and disks, with twelve serial test processes per runner.
+    # Global shard indices ensure every file runs exactly once across both machines.
+    strategy:
+      fail-fast: false
+      matrix:
+        machine: [1, 2]
     runs-on: warp-ubuntu-latest-x64-32x
     env:
       DOCKERHUB_USERNAME: ${{ secrets.DOCKERHUB_USERNAME }}
-      SHARD_TOTAL: 12
+      MACHINE_INDEX: ${{ matrix.machine }}
+      SHARDS_PER_MACHINE: 12
+      SHARD_TOTAL: 24
     steps:
       - name: 🔧 Disable IPv6
         run: |
@@ -104,24 +108,32 @@ jobs:
         run: pnpm run generate
 
       - name: 🏗️ Build test dependencies
-        # Build once up-front so the parallel shard runs below (turbo --only) never race
-        # to build or cache-restore the same outputs concurrently.
+        # Build once up-front so shard processes never race to build or
+        # cache-restore the same outputs concurrently.
         run: pnpm exec turbo run build --filter "@internal/*..."
 
-      - name: 🧪 Run Internal Unit Tests (${{ env.SHARD_TOTAL }} in-machine shards)
+      - name: 🧪 Run Internal Unit Tests (machine ${{ matrix.machine }}/2)
         run: |
-          # Same shard partitioning as the old 12-job matrix (DurationShardingSequencer
-          # keys off --shard=i/N), but as parallel local processes. --only skips the
-          # ^build dependency handled by the step above.
+          # Each process discovers all internal projects, then takes its globally
+          # balanced share of files. Keep one test file active per process.
+          first=$(( (MACHINE_INDEX - 1) * SHARDS_PER_MACHINE + 1 ))
+          last=$(( MACHINE_INDEX * SHARDS_PER_MACHINE ))
           status=0
           declare -a pids
-          for i in $(seq 1 "$SHARD_TOTAL"); do
-            pnpm exec turbo run test --only --concurrency=1 --filter "@internal/*" -- \
-              --run --reporter=default --reporter=blob --shard="$i/$SHARD_TOTAL" --passWithNoTests \
-              > "/tmp/internal-shard-$i.log" 2>&1 &
+          for i in $(seq "$first" "$last"); do
+            (
+              started=$SECONDS
+              result=0
+              pnpm exec vitest run --config vitest.internal.config.ts \
+                --sequence.concurrent=false --reporter=default --reporter=blob \
+                --shard="$i/$SHARD_TOTAL" --passWithNoTests \
+                > "/tmp/internal-shard-$i.log" 2>&1 || result=$?
+              echo "::notice::internal shard $i/$SHARD_TOTAL finished in $((SECONDS - started))s (exit $result)"
+              exit "$result"
+            ) &
             pids[i]=$!
           done
-          for i in $(seq 1 "$SHARD_TOTAL"); do
+          for i in $(seq "$first" "$last"); do
             if ! wait "${pids[i]}"; then
               status=1
               echo "::error::internal unit test shard $i/$SHARD_TOTAL failed"
@@ -132,13 +144,67 @@ jobs:
           done
           exit "$status"
 
-      - name: Gather all reports
+      - name: Upload blob reports
         if: ${{ !cancelled() }}
+        uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
+        with:
+          name: internal-blob-report-${{ matrix.machine }}
+          path: .vitest-reports/*
+          include-hidden-files: true
+          if-no-files-found: error
+          retention-days: 1
+
+  merge-reports:
+    name: "📊 Merge Internal Reports"
+    if: ${{ !cancelled() }}
+    needs: [unitTests]
+    runs-on: warp-ubuntu-latest-x64-2x
+    steps:
+      - name: ⬇️ Checkout repo
+        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
+        with:
+          fetch-depth: 1
+          persist-credentials: false
+
+      - name: ⎔ Setup pnpm
+        uses: pnpm/action-setup@ea17c68df8912ef543352723c149a84f56e3d413 # v6.1.0
+        with:
+          version: 10.33.2
+
+      - name: ⎔ Setup node
+        uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7.0.0
+        with:
+          node-version: 24.18.0
+
+      - name: Download blob reports
+        uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1
+        with:
+          path: .vitest-reports
+          pattern: internal-blob-report-*
+          merge-multiple: true
+
+      - name: Check all shard reports are present
         run: |
-          mkdir -p .vitest-reports
-          find . -type f -path '*/.vitest-reports/blob-*.json' \
-            -exec bash -c 'src="$1"; basename=$(basename "$src"); pkg=$(dirname "$src" | sed "s|^\./||;s|/\.vitest-reports$||;s|/|_|g"); cp "$src" ".vitest-reports/${pkg}-${basename}"' _ {} \;
+          for i in $(seq 1 24); do
+            test -s ".vitest-reports/blob-$i-24.json"
+          done
 
       - name: 📊 Merge reports
+        run: >-
+          pnpm dlx vitest@4.1.7 run --config vitest.internal-reports.config.mjs --merge-reports
+          --reporter=default --reporter=json --outputFile.json=internal-test-results.json
+
+      - name: Validate merged test timings
+        run: |
+          cp test-timings.json "$RUNNER_TEMP/checked-test-timings.json"
+          node internal-packages/testcontainers/scripts/update-test-timings.mjs \
+            internal-test-results.json "$RUNNER_TEMP/checked-test-timings.json"
+
+      - name: Upload internal test timings
         if: ${{ !cancelled() }}
-        run: pnpm dlx vitest@4.1.7 run --merge-reports --pass-with-no-tests
+        uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
+        with:
+          name: internal-test-results
+          path: internal-test-results.json
+          if-no-files-found: error
+          retention-days: 14

internal-packages/observability-map/src/integration.test.ts

@@ -9,6 +9,7 @@ import {
 } from "node:fs";
 import { tmpdir } from "node:os";
 import { join, resolve } from "node:path";
+import internalTestConfig from "../../../vitest.internal.config.js";
 import { isScannableFile, scanDirectory, scanFile } from "./scan.js";
 import { buildReport } from "./score.js";
 import { SCORED_CHECK_IDS } from "./checks/index.js";
@@ -306,7 +307,10 @@ describe("the package's tests are wired into the gate", () => {
     const internal = text.split("            internal:")[1]!.split("            obsmap:")[0]!;
     expect(internal).toContain(`'${PATH_PREFIX}internal-packages/**'`);
     expect(internal).not.toContain(`!${PATH_PREFIX}internal-packages/observability-map`);
-    expect(read(resolve(WORKFLOWS, "unit-tests-internal.yml"))).toContain('--filter "@internal/*"');
+    expect(read(resolve(WORKFLOWS, "unit-tests-internal.yml"))).toContain(
+      "--config vitest.internal.config.ts"
+    );
+    expect(internalTestConfig.test?.projects).toContain(resolve(__dirname, ".."));
   });
 
   // The test above only checks the package's own source path, a different overlap that was already

internal-packages/testcontainers/scripts/update-test-timings.mjs

@@ -0,0 +1,17 @@
+// Node 24 strips the helper's types, so the weekly job needs no dependency install.
+import { readFileSync, writeFileSync } from "node:fs";
+import { updateTestTimings } from "../src/test-timings.ts";
+
+const [reportPath, timingsPath] = process.argv.slice(2);
+if (!reportPath || !timingsPath) {
+  throw new Error(
+    "Usage: update-test-timings.mjs <internal-test-results.json> <test-timings.json>"
+  );
+}
+const previous = JSON.parse(readFileSync(timingsPath, "utf8"));
+const report = JSON.parse(readFileSync(reportPath, "utf8"));
+const updated = updateTestTimings(previous, report);
+writeFileSync(timingsPath, `${JSON.stringify(updated, null, 2)}\n`);
+console.log(
+  `Refreshed ${Object.keys(updated).filter((file) => updated[file] !== previous[file]).length} timings`
+);

internal-packages/testcontainers/src/internal-report-merge.test.ts

@@ -0,0 +1,63 @@
+import { execFileSync } from "node:child_process";
+import { mkdtempSync, mkdirSync, readFileSync, rmSync, writeFileSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { resolve } from "node:path";
+import { pathToFileURL } from "node:url";
+import { expect, test } from "vitest";
+
+const monorepoRoot = resolve(__dirname, "../../..");
+const vitestCli = resolve(monorepoRoot, "node_modules/vitest/vitest.mjs");
+const projectHelper = pathToFileURL(resolve(monorepoRoot, "internal-test-projects.mts")).href;
+
+test("replays named project blobs into a nonempty JSON report without package dependencies", () => {
+  const root = mkdtempSync(resolve(tmpdir(), "internal-report-merge-"));
+  try {
+    const projects = ["first", "second"].map((name) => {
+      const directory = resolve(root, "internal-packages", name);
+      mkdirSync(directory, { recursive: true });
+      writeFileSync(
+        resolve(directory, "package.json"),
+        JSON.stringify({ name: `@internal/${name}`, scripts: { test: "vitest" } })
+      );
+      writeFileSync(
+        resolve(directory, "vitest.config.mjs"),
+        "export default { test: { globals: true } };"
+      );
+      writeFileSync(
+        resolve(directory, "example.test.js"),
+        "test('passes', () => expect(2 + 2).toBe(4));"
+      );
+      return directory;
+    });
+    writeFileSync(
+      resolve(root, "run.config.mjs"),
+      `export default { test: { projects: ${JSON.stringify(projects)} } };`
+    );
+    writeFileSync(
+      resolve(root, "merge.config.mjs"),
+      `import { getInternalTestReportConfig } from ${JSON.stringify(projectHelper)};
+export default getInternalTestReportConfig(${JSON.stringify(root)});`
+    );
+    const run = (args: string[]) =>
+      execFileSync(process.execPath, [vitestCli, "run", ...args], {
+        cwd: root,
+        encoding: "utf8",
+        timeout: 10_000,
+      });
+    run(["--config=run.config.mjs", "--maxWorkers=1", "--reporter=blob"]);
+    run([
+      "--config=merge.config.mjs",
+      "--merge-reports",
+      "--reporter=json",
+      "--outputFile.json=results.json",
+    ]);
+    const merged = JSON.parse(readFileSync(resolve(root, "results.json"), "utf8"));
+    expect(merged.success).toBe(true);
+    expect(merged.numPassedTests).toBe(2);
+    expect(merged.testResults.map((file: { name: string }) => file.name).sort()).toEqual(
+      projects.map((directory) => resolve(directory, "example.test.js")).sort()
+    );
+  } finally {
+    rmSync(root, { recursive: true, force: true });
+  }
+}, 30_000);

internal-packages/testcontainers/src/sequencer.cjs

@@ -56,8 +56,15 @@ function median(nums) {
 // It's a rotation of the bin->shard mapping, so coverage stays exact (each file runs once).
 function packageOffset(specs, count) {
   if (specs.length === 0) return 0;
-  const rel = path.relative(REPO_ROOT, specs[0].moduleId);
-  const key = rel.split(path.sep).slice(0, 2).join("/");
+  const keys = new Set(
+    specs.map((spec) =>
+      path.relative(REPO_ROOT, spec.moduleId).split(path.sep).slice(0, 2).join("/")
+    )
+  );
+  // A multi-project run is already balanced globally. Its mapping must not depend
+  // on which package happened to finish discovering files first.
+  if (keys.size !== 1) return 0;
+  const [key] = keys;
   // FNV-1a - spreads similar sibling package names (e.g. internal-packages/*) far better than a
   // simple polynomial hash mod count, which collided run-engine + schedule-engine onto one shard.
   let h = 2166136261;

internal-packages/testcontainers/src/sequencer.test.ts

@@ -0,0 +1,48 @@
+import { readFileSync } from "node:fs";
+import { resolve } from "node:path";
+import { describe, expect, test } from "vitest";
+import type { TestSpecification, Vitest } from "vitest/node";
+import { DurationShardingSequencer } from "./sequencer.cjs";
+
+const root = resolve(__dirname, "../../..");
+const timings: Record<string, number> = JSON.parse(
+  readFileSync(resolve(root, "test-timings.json"), "utf8")
+);
+const paths = Object.keys(timings).filter((file) => file.startsWith("internal-packages/"));
+// Sequencing only consumes moduleId; these are inputs to the scheduling algorithm.
+const specs = [...paths, "internal-packages/new-package/src/new.test.ts"].map(
+  (file) => ({ moduleId: resolve(root, file) }) as TestSpecification
+);
+
+function sequencer(index: number, count: number) {
+  return new DurationShardingSequencer({ config: { shard: { index, count } } } as Vitest);
+}
+
+describe("multi-project duration sharding", () => {
+  test.each([1, 2, 3, 12, 24])(
+    "assigns every file exactly once across %i shards",
+    async (count) => {
+      const shards = await Promise.all(
+        Array.from({ length: count }, (_, index) => sequencer(index + 1, count).shard(specs))
+      );
+      const assigned = shards.flat().map((spec) => spec.moduleId);
+      expect(assigned).toHaveLength(specs.length);
+      expect(new Set(assigned).size).toBe(specs.length);
+      expect(assigned.sort()).toEqual(specs.map((spec) => spec.moduleId).sort());
+    }
+  );
+
+  test("keeps the same assignment when project discovery order changes", async () => {
+    for (let index = 1; index <= 12; index++) {
+      const original = await sequencer(index, 12).shard(specs);
+      const reversed = await sequencer(index, 12).shard([...specs].reverse());
+      expect(reversed.map((spec) => spec.moduleId)).toEqual(original.map((spec) => spec.moduleId));
+    }
+  });
+
+  test("retains all files without sharding and accepts an empty project", async () => {
+    const unsharded = new DurationShardingSequencer({ config: {} } as Vitest);
+    expect(await unsharded.shard(specs)).toBe(specs);
+    expect(await sequencer(1, 12).shard([])).toEqual([]);
+  });
+});