Merge pull request #633 from Mike-Dax/sharding

feat: Sharding
reg-viz · Feb 2, 2023 · a2e346b · a2e346b
2 parents 62bcea3 + ccf01c2
commit a2e346b
Show file tree

Hide file tree

Showing 6 changed files with 368 additions and 3 deletions.
diff --git a/README.md b/README.md
@@ -31,6 +31,7 @@ It is primarily responsible for image generation necessary for Visual Testing su
 - [Multiple PNGs from 1 story](#multiple-pngs-from-1-story)
   - [Basic usage](#basic-usage)
   - [Variants composition](#variants-composition)
+  - [Parallelisation across multiple computers](#parallelisation-across-multiple-computers)
 - [Tips](#tips)
   - [Run with Docker](#run-with-docker)
   - [Full control the screenshot timing](#full-control-the-screenshot-timing)
@@ -361,6 +362,9 @@ Options:
       --verbose                                                                               [boolean] [default: false]
       --serverCmd                  Command line to launch Storybook server.                       [string] [default: ""]
       --serverTimeout              Timeout [msec] for starting Storybook server.               [number] [default: 60000]
+      --shard                      The sharding options for this run. In the format <shardNumber>/<totalShards>.
+                                   <shardNumber> is a number between 1 and <totalShards>. <totalShards> is the total
+                                   number of computers working.                                [string] [default: "1/1"]
       --captureTimeout             Timeout [msec] for capture a story.                          [number] [default: 5000]
       --captureMaxRetryCount       Number of count to retry to capture.                            [number] [default: 3]
       --metricsWatchRetryCount     Number of count to retry until browser metrics stable.       [number] [default: 1000]
@@ -455,6 +459,16 @@ The above example generates the following:
 
 **Note:** You can extend some viewports with keys of `viewports` option because the `viewports` field is expanded to variants internally.
 
+### Parallelisation across multiple computers
+
+To process more stories in parallel across multiple computers, the `shard` argument can be used.
+
+The `shard` argument is a string of the format: `<shardNumber>/<totalShards>`. `<shardNumber>` is a number between 1 and `<totalShards>`, inclusive. `<totalShards>` is the total number of computers running the execution.
+
+For example, a run with `--shard 1/1` would be considered the default behaviour on a single computer. Two computers each running `--shard 1/2` and `--shard 2/2` respectively would split the stories across two computers.
+
+Stories are distributed across shards in a round robin fashion when ordered by their ID. If a series of stories 'close together' are slower to screenshot than others, they should be distributed evenly.
+
 ## Tips
 
 ### Run with Docker

diff --git a/packages/storycap/src/node/cli.ts b/packages/storycap/src/node/cli.ts
@@ -2,9 +2,10 @@
 
 import { time, ChromeChannel, getDeviceDescriptors } from 'storycrawler';
 import { main } from './main';
-import { MainOptions } from './types';
+import { MainOptions, ShardOptions } from './types';
 import yargs from 'yargs';
 import { Logger } from './logger';
+import { parseShardOptions } from './shard-utilities';
 
 function showDevices(logger: Logger) {
   getDeviceDescriptors().map(device => logger.log(device.name, JSON.stringify(device.viewport)));
@@ -41,6 +42,12 @@ function createOptions(): MainOptions {
       default: 60_000,
       description: 'Timeout [msec] for starting Storybook server.',
     })
+    .option('shard', {
+      string: true,
+      default: '1/1',
+      description:
+        'The sharding options for this run. In the format <shardNumber>/<totalShards>. <shardNumber> is a number between 1 and <totalShards>. <totalShards> is the total number of computers working.',
+    })
     .option('captureTimeout', { number: true, default: 5_000, description: 'Timeout [msec] for capture a story.' })
     .option('captureMaxRetryCount', { number: true, default: 3, description: 'Number of count to retry to capture.' })
     .option('metricsWatchRetryCount', {
@@ -109,6 +116,7 @@ function createOptions(): MainOptions {
     verbose,
     serverTimeout,
     serverCmd,
+    shard,
     captureTimeout,
     captureMaxRetryCount,
     metricsWatchRetryCount,
@@ -141,6 +149,14 @@ function createOptions(): MainOptions {
     throw error;
   }
 
+  let shardOptions: ShardOptions;
+  try {
+    shardOptions = parseShardOptions(shard);
+  } catch (error) {
+    logger.error(error);
+    throw error;
+  }
+
   const opt = {
     serverOptions: {
       storybookUrl,
@@ -154,6 +170,7 @@ function createOptions(): MainOptions {
     delay,
     viewports: viewport,
     parallel,
+    shard: shardOptions,
     captureTimeout,
     captureMaxRetryCount,
     metricsWatchRetryCount,

diff --git a/packages/storycap/src/node/main.ts b/packages/storycap/src/node/main.ts
@@ -4,6 +4,7 @@ import { CapturingBrowser } from './capturing-browser';
 import { MainOptions, RunMode } from './types';
 import { FileSystem } from './file';
 import { createScreenshotService } from './screenshot-service';
+import { shardStories, sortStories } from './shard-utilities';
 
 async function detectRunMode(storiesBrowser: StoriesBrowser, opt: MainOptions) {
   // Reuse `storiesBrowser` instance to avoid cost of re-launching another Puppeteer process.
@@ -60,21 +61,41 @@ export async function main(mainOptions: MainOptions) {
   storiesBrowser.close();
 
   const stories = filterStories(allStories, mainOptions.include, mainOptions.exclude);
+
   if (stories.length === 0) {
     logger.warn('There is no matched story. Check your include/exclude options.');
     return 0;
   }
 
-  logger.log(`Found ${logger.color.green(stories.length + '')} stories.`);
+  const sortedStories = sortStories(stories);
+  const shardedStories = shardStories(sortedStories, mainOptions.shard.shardNumber, mainOptions.shard.totalShards);
+
+  if (shardedStories.length === 0) {
+    logger.log('This shard has no stories to screenshot.');
+    return 0;
+  }
+
+  if (mainOptions.shard.totalShards === 1) {
+    logger.log(`Found ${logger.color.green(String(stories.length))} stories.`);
+  } else {
+    logger.log(
+      `Found ${logger.color.green(String(stories.length))} stories. ${logger.color.green(
+        String(shardedStories.length),
+      )} are being processed by this shard (number ${mainOptions.shard.shardNumber} of ${
+        mainOptions.shard.totalShards
+      }).`,
+    );
+  }
 
   // Launce Puppeteer processes to capture each story.
   const { workers, closeWorkers } = await bootCapturingBrowserAsWorkers(connection, mainOptions, mode);
   logger.debug('Created workers.');
 
   try {
     // Execution caputuring procedure.
-    return await createScreenshotService({ workers, stories, fileSystem, logger }).execute();
+    const captured = await createScreenshotService({ workers, stories: shardedStories, fileSystem, logger }).execute();
     logger.debug('Ended ScreenshotService execution.');
+    return captured;
   } catch (error) {
     if (error instanceof ChromiumNotFoundError) {
       throw new Error(

diff --git a/packages/storycap/src/node/shard-utilities.test.ts b/packages/storycap/src/node/shard-utilities.test.ts
@@ -0,0 +1,244 @@
+import { Story } from 'storycrawler';
+import { parseShardOptions, sortStories, shardStories } from './shard-utilities';
+
+describe(parseShardOptions, () => {
+  it('should accept correct arguments', () => {
+    expect(parseShardOptions('1/1')).toMatchObject({ shardNumber: 1, totalShards: 1 });
+    expect(parseShardOptions('1/2')).toMatchObject({ shardNumber: 1, totalShards: 2 });
+    expect(parseShardOptions('2/2')).toMatchObject({ shardNumber: 2, totalShards: 2 });
+    expect(parseShardOptions('1/3')).toMatchObject({ shardNumber: 1, totalShards: 3 });
+    expect(parseShardOptions('2/3')).toMatchObject({ shardNumber: 2, totalShards: 3 });
+    expect(parseShardOptions('3/3')).toMatchObject({ shardNumber: 3, totalShards: 3 });
+  });
+  it('should be resiliant to whitespace', () => {
+    expect(parseShardOptions(' 1/2')).toMatchObject({ shardNumber: 1, totalShards: 2 });
+    expect(parseShardOptions('1/2 ')).toMatchObject({ shardNumber: 1, totalShards: 2 });
+    expect(parseShardOptions(' 1/2 ')).toMatchObject({ shardNumber: 1, totalShards: 2 });
+    expect(parseShardOptions('1 /2')).toMatchObject({ shardNumber: 1, totalShards: 2 });
+    expect(parseShardOptions('1/ 2')).toMatchObject({ shardNumber: 1, totalShards: 2 });
+    expect(parseShardOptions('1 / 2')).toMatchObject({ shardNumber: 1, totalShards: 2 });
+    expect(parseShardOptions(' 1 /2')).toMatchObject({ shardNumber: 1, totalShards: 2 });
+    expect(parseShardOptions('1/ 2 ')).toMatchObject({ shardNumber: 1, totalShards: 2 });
+    expect(parseShardOptions(' 1 / 2 ')).toMatchObject({ shardNumber: 1, totalShards: 2 });
+  });
+  it('errors for incorrect arguments', () => {
+    expect(() => parseShardOptions('0')).toThrowError();
+    expect(() => parseShardOptions('1')).toThrowError();
+    expect(() => parseShardOptions('text')).toThrowError();
+    expect(() => parseShardOptions('0/1')).toThrowError();
+    expect(() => parseShardOptions('-1/1')).toThrowError();
+    expect(() => parseShardOptions('2/1')).toThrowError();
+    expect(() => parseShardOptions('0/3')).toThrowError();
+    expect(() => parseShardOptions('/3')).toThrowError();
+    expect(() => parseShardOptions('4/')).toThrowError();
+    expect(() => parseShardOptions('4/3')).toThrowError();
+    expect(() => parseShardOptions('ab/c')).toThrowError();
+  });
+});
+
+describe(sortStories, () => {
+  it('should sort stories alphabetically based on their ID', () => {
+    const stories: Story[] = [
+      {
+        id: 'simple-tooltip--with-component-content',
+        kind: 'simple/Tooltip',
+        story: 'with component content',
+        version: 'v5',
+      },
+      {
+        id: 'complex-scene--for-table-of-contents',
+        kind: 'complex/Scene',
+        story: 'for table-of-contents',
+        version: 'v5',
+      },
+      {
+        id: 'complex-scene--basic-usage',
+        kind: 'complex/Scene',
+        story: 'basic-usage',
+        version: 'v5',
+      },
+      {
+        id: 'complex-scene--verticalannotation',
+        kind: 'complex/Scene',
+        story: 'verticalannotation',
+        version: 'v5',
+      },
+    ];
+
+    const sortedStories = sortStories(stories);
+
+    let prev: Story | null = null;
+
+    for (const next of sortedStories) {
+      if (!prev) {
+        prev = next;
+        continue;
+      }
+      expect(next.id > prev.id).toBeTruthy();
+
+      prev = next;
+    }
+  });
+});
+
+describe(shardStories, () => {
+  it('a single shard gets all the stories', () => {
+    const stories: Story[] = [
+      {
+        id: 'simple-tooltip--with-component-content',
+        kind: 'simple/Tooltip',
+        story: 'with component content',
+        version: 'v5',
+      },
+      {
+        id: 'complex-scene--for-table-of-contents',
+        kind: 'complex/Scene',
+        story: 'for table-of-contents',
+        version: 'v5',
+      },
+      {
+        id: 'complex-scene--basic-usage',
+        kind: 'complex/Scene',
+        story: 'basic-usage',
+        version: 'v5',
+      },
+      {
+        id: 'complex-scene--verticalannotation',
+        kind: 'complex/Scene',
+        story: 'verticalannotation',
+        version: 'v5',
+      },
+    ];
+
+    const sortedStories = sortStories(stories);
+    const shardedStories = shardStories(sortedStories, 1, 1);
+
+    expect(shardedStories).toMatchObject(sortedStories);
+  });
+  it('two shards get equal amounts of stories when the number of them is even', () => {
+    const stories: Story[] = [
+      {
+        id: 'simple-tooltip--with-component-content',
+        kind: 'simple/Tooltip',
+        story: 'with component content',
+        version: 'v5',
+      },
+      {
+        id: 'complex-scene--for-table-of-contents',
+        kind: 'complex/Scene',
+        story: 'for table-of-contents',
+        version: 'v5',
+      },
+      {
+        id: 'complex-scene--basic-usage',
+        kind: 'complex/Scene',
+        story: 'basic-usage',
+        version: 'v5',
+      },
+      {
+        id: 'complex-scene--verticalannotation',
+        kind: 'complex/Scene',
+        story: 'verticalannotation',
+        version: 'v5',
+      },
+    ];
+
+    const sortedStories = sortStories(stories);
+    const shardedStoriesA = shardStories(sortedStories, 1, 2);
+    const shardedStoriesB = shardStories(sortedStories, 2, 2);
+
+    expect(shardedStoriesA.length).toBe(shardedStoriesB.length);
+  });
+
+  it('two shards get roughly equal amounts of stories when the number of them is odd', () => {
+    const stories: Story[] = [
+      {
+        id: 'simple-tooltip--with-component-content',
+        kind: 'simple/Tooltip',
+        story: 'with component content',
+        version: 'v5',
+      },
+      {
+        id: 'complex-scene--for-table-of-contents',
+        kind: 'complex/Scene',
+        story: 'for table-of-contents',
+        version: 'v5',
+      },
+      {
+        id: 'complex-scene--verticalannotation',
+        kind: 'complex/Scene',
+        story: 'verticalannotation',
+        version: 'v5',
+      },
+    ];
+
+    const sortedStories = sortStories(stories);
+    const shardedStoriesA = shardStories(sortedStories, 1, 2);
+    const shardedStoriesB = shardStories(sortedStories, 2, 2);
+
+    expect(Math.abs(shardedStoriesA.length - shardedStoriesB.length)).toBeLessThanOrEqual(1);
+  });
+
+  it("stories aren't duplicated when there are more shards than stories", () => {
+    const stories: Story[] = [
+      {
+        id: 'simple-tooltip--with-component-content',
+        kind: 'simple/Tooltip',
+        story: 'with component content',
+        version: 'v5',
+      },
+      {
+        id: 'complex-scene--for-table-of-contents',
+        kind: 'complex/Scene',
+        story: 'for table-of-contents',
+        version: 'v5',
+      },
+    ];
+
+    const sortedStories = sortStories(stories);
+    const shardedStoriesA = shardStories(sortedStories, 1, 4);
+    const shardedStoriesB = shardStories(sortedStories, 2, 4);
+    const shardedStoriesC = shardStories(sortedStories, 3, 4);
+    const shardedStoriesD = shardStories(sortedStories, 4, 4);
+
+    expect(shardedStoriesA.length + shardedStoriesB.length + shardedStoriesC.length + shardedStoriesD.length).toBe(
+      sortedStories.length,
+    );
+  });
+
+  it('complex and simple stories are distributed evenly across shards', () => {
+    function makeDummyStory(index: number, complex: boolean): Story {
+      return {
+        id: `${complex ? 'complex' : 'simple'}-component--${index}`,
+        kind: `${complex ? 'complex' : 'simple'}/Component`,
+        story: `${index}`,
+        version: 'v5',
+      } as const;
+    }
+
+    const stories: Story[] = [
+      makeDummyStory(0, true),
+      makeDummyStory(1, true),
+      makeDummyStory(2, true),
+      makeDummyStory(3, true),
+      makeDummyStory(4, false),
+      makeDummyStory(5, false),
+      makeDummyStory(6, false),
+      makeDummyStory(7, false),
+      makeDummyStory(8, false),
+      makeDummyStory(9, false),
+      makeDummyStory(10, false),
+      makeDummyStory(11, false),
+    ];
+
+    const sortedStories = sortStories(stories);
+    const shardedStoriesA = shardStories(sortedStories, 1, 2);
+    const shardedStoriesB = shardStories(sortedStories, 2, 2);
+
+    const numComplexOnA = shardedStoriesA.filter(story => story.id.startsWith('complex')).length;
+    const numComplexOnB = shardedStoriesB.filter(story => story.id.startsWith('complex')).length;
+
+    expect(shardedStoriesA.length).toBe(shardedStoriesB.length);
+    expect(numComplexOnA).toBe(numComplexOnB);
+  });
+});