From cbd8a51a8d58f7110adf37814229afc375ce5319 Mon Sep 17 00:00:00 2001 From: Michael Orenstein Date: Sat, 15 Oct 2022 18:43:12 +1030 Subject: [PATCH 1/2] feat: Sharding --- README.md | 14 + packages/storycap/src/node/cli.ts | 19 +- packages/storycap/src/node/main.ts | 25 +- .../storycap/src/node/shard-utilities.test.ts | 252 ++++++++++++++++++ packages/storycap/src/node/shard-utilities.ts | 58 ++++ packages/storycap/src/node/types.ts | 11 + 6 files changed, 376 insertions(+), 3 deletions(-) create mode 100644 packages/storycap/src/node/shard-utilities.test.ts create mode 100644 packages/storycap/src/node/shard-utilities.ts diff --git a/README.md b/README.md index 998817637..e655b137d 100644 --- a/README.md +++ b/README.md @@ -32,6 +32,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) @@ -362,6 +363,9 @@ Options: --verbose [boolean] [default: false] --serverCmd Command line to launch Storybook server. [string] [default: ""] --serverTimeout Timeout [msec] for starting Storybook server. [number] [default: 20000] + --shard The sharding options for this run. In the format /. + is a number between 1 and . 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] @@ -456,6 +460,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: `/`. `` is a number between 1 and ``, inclusive. `` 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 index cfb33fcb6..4d2130a17 100644 --- 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: 20_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 /. is a number between 1 and . 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 index 49f15ee86..484106989 100644 --- 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,12 +61,31 @@ 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); @@ -73,8 +93,9 @@ export async function main(mainOptions: MainOptions) { 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 new file mode 100644 index 000000000..d71e9d1aa --- /dev/null +++ b/packages/storycap/src/node/shard-utilities.test.ts @@ -0,0 +1,252 @@ +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); + + console.log( + shardedStoriesA.length, + shardedStoriesB.length, + shardedStoriesC.length, + shardedStoriesD.length, + sortedStories.length, + ); + + 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); + }); +}); diff --git a/packages/storycap/src/node/shard-utilities.ts b/packages/storycap/src/node/shard-utilities.ts new file mode 100644 index 000000000..9e791223c --- /dev/null +++ b/packages/storycap/src/node/shard-utilities.ts @@ -0,0 +1,58 @@ +import { Story } from 'storycrawler'; +import { ShardOptions } from './types'; + +export const parseShardOptions = (arg: string): ShardOptions => { + const split = arg.split('/'); + + const shardNumber = parseInt(split[0].trim(), 10); + const totalShards = parseInt(split[1].trim(), 10); + + if (split.length !== 2 || Number.isNaN(shardNumber) || Number.isNaN(totalShards)) { + throw new Error(`The shard argument must be in the format /.`); + } + + if (shardNumber === 0 || totalShards === 0) { + throw new Error(`The shard arguments cannot be 0.`); + } + + if (shardNumber < 0 || totalShards < 0) { + throw new Error(`The shard arguments cannot be negative.`); + } + + if (shardNumber > totalShards) { + throw new Error(`The shard number cannot be greater than the total number of shards.`); + } + + return { + shardNumber, + totalShards, + }; +}; + +/** + * + * Sort the stories by their ID. + * + **/ +export const sortStories = (stories: Story[]): Story[] => { + return stories.sort((a, b) => { + if (a.id < b.id) { + return -1; + } + if (a.id > b.id) { + return 1; + } + return 0; + }); +}; + +/** + * + * Shard stories in a round robin fashion based on their index in the sorted list. + * + **/ +export const shardStories = (stories: Story[], shardNumber: number, totalShards: number): Story[] => { + const shardIndex = shardNumber - 1; + + return stories.filter((_, index) => index % totalShards === shardIndex); +}; diff --git a/packages/storycap/src/node/types.ts b/packages/storycap/src/node/types.ts index e619ef7cd..647b04319 100644 --- a/packages/storycap/src/node/types.ts +++ b/packages/storycap/src/node/types.ts @@ -9,6 +9,16 @@ import { StorybookConnectionOptions, BaseBrowserOptions, ChromeChannel } from 's **/ export type RunMode = 'simple' | 'managed'; +/** + * + * Parameters for sharding. + * + **/ +export type ShardOptions = { + shardNumber: number; + totalShards: number; +}; + /** * * Parameters for main procedure. @@ -31,6 +41,7 @@ export interface MainOptions extends BaseBrowserOptions { disableCssAnimation: boolean; disableWaitAssets: boolean; parallel: number; + shard: ShardOptions; metricsWatchRetryCount: number; chromiumChannel: ChromeChannel; chromiumPath: string; From 387c99b922fc85dda08e9554dd19ac1fe90b1995 Mon Sep 17 00:00:00 2001 From: Michael Orenstein Date: Sat, 15 Oct 2022 23:08:57 +1030 Subject: [PATCH 2/2] chore: Removed console statement --- packages/storycap/src/node/shard-utilities.test.ts | 8 -------- 1 file changed, 8 deletions(-) diff --git a/packages/storycap/src/node/shard-utilities.test.ts b/packages/storycap/src/node/shard-utilities.test.ts index d71e9d1aa..0e9fc286e 100644 --- a/packages/storycap/src/node/shard-utilities.test.ts +++ b/packages/storycap/src/node/shard-utilities.test.ts @@ -201,14 +201,6 @@ describe(shardStories, () => { const shardedStoriesC = shardStories(sortedStories, 3, 4); const shardedStoriesD = shardStories(sortedStories, 4, 4); - console.log( - shardedStoriesA.length, - shardedStoriesB.length, - shardedStoriesC.length, - shardedStoriesD.length, - sortedStories.length, - ); - expect(shardedStoriesA.length + shardedStoriesB.length + shardedStoriesC.length + shardedStoriesD.length).toBe( sortedStories.length, );