shard split from timing file

packages/playwright/src/common/config.ts

@@ -57,6 +57,7 @@ export class FullConfigInternal {
   cliPassWithNoTests?: boolean;
   testIdMatcher?: Matcher;
   defineConfigWasUsed = false;
+  cliTimingFile?: string;
 
   constructor(location: ConfigLocation, userConfig: Config, configCLIOverrides: ConfigCLIOverrides) {
     if (configCLIOverrides.projects && userConfig.projects)

packages/playwright/src/program.ts

@@ -189,6 +189,7 @@ async function runTests(args: string[], opts: { [key: string]: any }) {
   config.cliListOnly = !!opts.list;
   config.cliProjectFilter = opts.project || undefined;
   config.cliPassWithNoTests = !!opts.passWithNoTests;
+  config.cliTimingFile = opts.timingFile || undefined;
 
   const runner = new Runner(config);
   let status: FullResult['status'];
@@ -336,6 +337,7 @@ const testOptions: [string, string][] = [
   ['--global-timeout <timeout>', `Maximum time this test suite can run in milliseconds (default: unlimited)`],
   ['-g, --grep <grep>', `Only run tests matching this regular expression (default: ".*")`],
   ['-gv, --grep-invert <grep>', `Only run tests that do not match this regular expression`],
+  ['--timing-file <file>', `Load JSON report to use as timing file for shard balancing`],
   ['--headed', `Run tests in headed browsers (default: headless)`],
   ['--ignore-snapshots', `Ignore screenshot and snapshot expectations`],
   ['--list', `Collect all the tests and report them, but do not run`],

packages/playwright/src/runner/loadUtils.ts

@@ -15,7 +15,7 @@
  */
 
 import path from 'path';
-import type { FullConfig, Reporter, TestError } from '../../types/testReporter';
+import type { FullConfig, JSONReport, Reporter, TestError } from '../../types/testReporter';
 import { InProcessLoaderHost, OutOfProcessLoaderHost } from './loaderHost';
 import { Suite } from '../common/test';
 import type { TestCase } from '../common/test';
@@ -27,10 +27,11 @@ import { buildProjectsClosure, collectFilesForProject, filterProjects } from './
 import type { TestRun } from './tasks';
 import { requireOrImport } from '../transform/transform';
 import { applyRepeatEachIndex, bindFileSuiteToProject, filterByFocusedLine, filterByTestIds, filterOnly, filterTestsRemoveEmptySuites } from '../common/suiteUtils';
-import { createTestGroups, filterForShard, type TestGroup } from './testGroups';
+import { createTestGroups, filterForShard, filterForShardFromTimingFile, type TestGroup } from './testGroups';
 import { dependenciesForTestFile } from '../transform/compilationCache';
 import { sourceMapSupport } from '../utilsBundle';
 import type { RawSourceMap } from 'source-map';
+import { readFileSync } from 'fs';
 
 export async function collectProjectsAndTestFiles(testRun: TestRun, doNotRunTestsOutsideProjectFilter: boolean, additionalFileMatcher?: Matcher) {
   const config = testRun.config;
@@ -182,7 +183,14 @@ export async function createRootSuite(testRun: TestRun, errors: TestError[], sho
       testGroups.push(...createTestGroups(projectSuite, config.config.workers));
 
     // Shard test groups.
-    const testGroupsInThisShard = filterForShard(config.config.shard, testGroups);
+    let testGroupsInThisShard: Set<TestGroup>;
+    if (config.cliTimingFile) {
+      const timingFile: JSONReport = JSON.parse(readFileSync(config.cliTimingFile, 'utf8'));
+      testGroupsInThisShard = filterForShardFromTimingFile(timingFile, config.config.shard, testGroups);
+    } else {
+      testGroupsInThisShard = filterForShard(config.config.shard, testGroups);
+    }
+
     const testsInThisShard = new Set<TestCase>();
     for (const group of testGroupsInThisShard) {
       for (const test of group.tests)

packages/playwright/src/runner/testGroups.ts

@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+import { JSONReport, JSONReportSuite, JSONReportTestResult } from 'packages/playwright-test/reporter';
 import type { Suite, TestCase } from '../common/test';
 
 export type TestGroup = {
@@ -24,6 +25,12 @@ export type TestGroup = {
   tests: TestCase[];
 };
 
+type testDetails = {
+  id: string,
+  name: string,
+  duration: number
+}
+
 export function createTestGroups(projectSuite: Suite, workers: number): TestGroup[] {
   // This function groups tests that can be run together.
   // Tests cannot be run together when:
@@ -162,3 +169,98 @@ export function filterForShard(shard: { total: number, current: number }, testGr
   }
   return result;
 }
+
+export function filterForShardFromTimingFile(timingFile: JSONReport, shard: { total: number, current: number }, testGroups: TestGroup[]): Set<TestGroup> {
+  const testDetails = getAllTestDetails(timingFile);
+  const shardsMapping = mapTestDetailsToShards(testDetails, shard);
+  const testIds = shardsMapping[shard.current - 1].map(shardDetails => shardDetails.id);
+  const recordedTests = testGroups.filter(group => testIds.includes(group.tests[0].id));
+  // console.log(shardsMapping);
+  // console.log(testDetails.length);
+  const result = new Set<TestGroup>();
+
+  recordedTests.forEach(group => result.add(group));
+
+  // If not all test groups have recorded timing information, filter the remaining groups based on shard distribution
+  if (testDetails.length != testGroups.length) {
+    const allTestIds = testDetails.map(shardDetails => shardDetails.id);
+    const UnrecordedTests = testGroups.filter(group => !allTestIds.includes(group.tests[0].id));
+
+    let shardableTotal = 0;
+    for (const group of UnrecordedTests)
+      shardableTotal += group.tests.length;
+  
+    // Each shard gets some tests.
+    const shardSize = Math.floor(shardableTotal / shard.total);
+    // First few shards get one more test each.
+    const extraOne = shardableTotal - shardSize * shard.total;
+  
+    const currentShard = shard.current - 1; // Make it zero-based for calculations.
+    const from = shardSize * currentShard + Math.min(extraOne, currentShard);
+    const to = from + shardSize + (currentShard < extraOne ? 1 : 0);
+  
+    let current = 0;
+    for (const group of UnrecordedTests) {
+      // Any test group goes to the shard that contains the first test of this group.
+      // So, this shard gets any group that starts at [from; to)
+      if (current >= from && current < to)
+        result.add(group);
+      current += group.tests.length;
+    }
+  }
+  return result;
+}
+
+// Extracts test details from a JSON report and returns an array of test details objects.
+function getAllTestDetails(report: JSONReport) {
+  const allTestDetails:testDetails[] = [];
+  report.suites.forEach((suite: JSONReportSuite) => {
+    getTestDetails(suite, allTestDetails);
+  });
+
+  return allTestDetails;
+}
+
+// Recursively traverses through test suites in a JSON report and extracts test details.
+function getTestDetails(suite: JSONReportSuite, allTestDetails: testDetails[]) {
+  if (suite.specs) {
+    suite.specs.forEach((spec) => {
+      spec.tests.forEach((test) => {
+        test.results.forEach((result) => {
+          allTestDetails.push({
+            id: spec.id,
+            name: spec.title,
+            duration: result.duration,
+          });
+        });
+      });
+    });
+  }
+
+  if (suite.suites) {
+    suite.suites.forEach((subSuite) => {
+      getTestDetails(subSuite, allTestDetails);
+    });
+  }
+}
+
+function mapTestDetailsToShards(testDetails: testDetails[], shard: {total: number, current: number}) {
+  testDetails.sort((a, b) => b.duration - a.duration);
+
+  const result:testDetails[][] = Array.from({ length: shard.total }, () => []);
+  const sums = Array(shard.total).fill(0);
+
+  // Distribute tests to shards based on their durations
+  for (let testDetailsObj of testDetails) {
+    // Find the shard with the smallest total duration and assign the test to that shard
+    let minIndex = sums.indexOf(Math.min(...sums));
+    result[minIndex].push(testDetailsObj);
+    sums[minIndex] += testDetailsObj.duration;
+  }
+
+  for (let i = 0; i < shard.total; i++) {
+    console.log(`Total duration for array ${i + 1}: ${sums[i]}`);
+  }
+
+  return result;
+}