feat: custom sharding weights

docs/src/test-sharding-js.md

@@ -45,6 +45,12 @@ Without the fullyParallel setting, Playwright Test defaults to file-level granul
 - **Without** `fullyParallel`: Tests are split at the file level, so to balance the shards, it's important to keep your test files small and evenly sized.
 - To ensure the most effective use of sharding, especially in CI environments, it is recommended to use `fullyParallel: true` when aiming for balanced distribution across shards. Otherwise, you may need to manually organize your test files to avoid imbalances.
 
+## Rebalancing Shards
+
+```bash
+npx playwright test --shard=1/4 --shard-weights=26:24:25:25
+```
+
 ## Merging reports from multiple shards
 
 In the previous example, each test shard has its own test report. If you want to have a combined report showing all the test results from all the shards, you can merge them.

packages/playwright/src/common/ipc.ts

@@ -37,6 +37,7 @@ export type ConfigCLIOverrides = {
   reporter?: ReporterDescription[];
   additionalReporters?: ReporterDescription[];
   shard?: { current: number, total: number };
+  shardWeights?: number[];
   timeout?: number;
   tsconfig?: string;
   ignoreSnapshots?: boolean;

packages/playwright/src/program.ts

@@ -300,6 +300,7 @@ function overridesFromOptions(options: { [key: string]: any }): ConfigCLIOverrid
     retries: options.retries ? parseInt(options.retries, 10) : undefined,
     reporter: resolveReporterOption(options.reporter),
     shard: resolveShardOption(options.shard),
+    shardWeights: resolveShardWeightsOption(options.shardWeights),
     timeout: options.timeout ? parseInt(options.timeout, 10) : undefined,
     tsconfig: options.tsconfig ? path.resolve(process.cwd(), options.tsconfig) : undefined,
     ignoreSnapshots: options.ignoreSnapshots ? !!options.ignoreSnapshots : undefined,
@@ -372,6 +373,18 @@ function resolveShardOption(shard?: string): ConfigCLIOverrides['shard'] {
   return { current, total };
 }
 
+function resolveShardWeightsOption(shardWeights?: string): ConfigCLIOverrides['shardWeights'] {
+  if (!shardWeights)
+    return undefined;
+
+  return shardWeights.split(':').map(w => {
+    const weight = parseInt(w, 10);
+    if (isNaN(weight) || weight < 0)
+      throw new Error(`--shard-weights "${shardWeights}" weights must be non-negative numbers`);
+    return weight;
+  });
+}
+
 function resolveReporter(id: string) {
   if (builtInReporters.includes(id as any))
     return id;
@@ -411,6 +424,7 @@ const testOptions: [string, { description: string, choices?: string[], preset?:
   ['--retries <retries>', { description: `Maximum retry count for flaky tests, zero for no retries (default: no retries)` }],
   ['--run-agents', { description: `Run agents to generate the code for page.perform` }],
   ['--shard <shard>', { description: `Shard tests and execute only the selected shard, specify in the form "current/all", 1-based, for example "3/5"` }],
+  ['--shard-weights <weights>', { description: `Weights for each shard, colon-separated, for example "2:1:1" for 3 shards where the first shard should be allocated half of the work` }],
   ['--test-list <file>', { description: `Path to a file containing a list of tests to run. See https://playwright.dev/docs/test-cli for more details.` }],
   ['--test-list-invert <file>', { description: `Path to a file containing a list of tests to skip. See https://playwright.dev/docs/test-cli for more details.` }],
   ['--timeout <timeout>', { description: `Specify test timeout threshold in milliseconds, zero for unlimited (default: ${defaultTimeout})` }],

packages/playwright/src/runner/loadUtils.ts

@@ -189,7 +189,7 @@ export async function createRootSuite(testRun: TestRun, errors: TestError[], sho
     }
 
     // Shard test groups.
-    const testGroupsInThisShard = filterForShard(config.config.shard, testGroups);
+    const testGroupsInThisShard = filterForShard(config.config.shard, config.configCLIOverrides.shardWeights, testGroups);
     const testsInThisShard = new Set<TestCase>();
     for (const group of testGroupsInThisShard) {
       for (const test of group.tests)

packages/playwright/src/runner/testGroups.ts

@@ -130,7 +130,12 @@ export function createTestGroups(projectSuite: Suite, expectedParallelism: numbe
   return result;
 }
 
-export function filterForShard(shard: { total: number, current: number }, testGroups: TestGroup[]): Set<TestGroup> {
+export function filterForShard(shard: { total: number, current: number }, weights: number[] | undefined, testGroups: TestGroup[]): Set<TestGroup> {
+  weights ??= Array.from({ length: shard.total }, () => 1);
+  if (weights.length !== shard.total)
+    throw new Error(`--shard-weights number of weights must match the shard total of ${shard.total}`);
+
+  const totalWeight = weights.reduce((a, b) => a + b, 0);
   // Note that sharding works based on test groups.
   // This means parallel files will be sharded by single tests,
   // while non-parallel files will be sharded by the whole file.
@@ -143,13 +148,17 @@ export function filterForShard(shard: { total: number, current: number }, testGr
     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 shardSizes = weights.map(w => Math.floor(w * shardableTotal / totalWeight));
+  const remainder = shardableTotal - shardSizes.reduce((a, b) => a + b, 0);
+  for (let i = 0; i < remainder; i++) {
+    // First few shards get one more test each.
+    shardSizes[i % shardSizes.length]++;
+  }
 
-  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 from = 0;
+  for (let i = 0; i < shard.current - 1; i++)
+    from += shardSizes[i];
+  const to = from + shardSizes[shard.current - 1];
 
   let current = 0;
   const result = new Set<TestGroup>();

tests/playwright-test/shard.spec.ts

@@ -16,6 +16,8 @@
 
 import { test, expect } from './playwright-test-fixtures';
 
+test.describe.configure({ mode: 'parallel' });
+
 const tests = {
   'a1.spec.ts': `
     import { test } from '@playwright/test';
@@ -330,3 +332,28 @@ test('should shard tests with beforeAll based on shards total instead of workers
     expect(result.outputLines).toEqual(['beforeAll', 'test7']);
   }
 });
+
+test('should respect custom shard weights', async ({ runInlineTest }) => {
+  await test.step('shard 1', async () => {
+    const result = await runInlineTest(tests, { 'shard': '1/2', 'shard-weights': '40:60', 'workers': 1 });
+    expect.soft(result.exitCode).toBe(0);
+    expect.soft(result.outputLines).toEqual([
+      'a1-test1-done',
+      'a1-test2-done',
+      'a1-test3-done',
+      'a1-test4-done',
+    ]);
+  });
+  await test.step('shard 2', async () => {
+    const result = await runInlineTest(tests, { 'shard': '2/2', 'shard-weights': '40:60', 'workers': 1 });
+    expect.soft(result.exitCode).toBe(0);
+    expect.soft(result.outputLines).toEqual([
+      'a2-test1-done',
+      'a2-test2-done',
+      'a3-test1-done',
+      'a3-test2-done',
+      'a4-test1-done',
+      'a4-test2-done',
+    ]);
+  });
+});