Allowing runing all cells in parallel (#1972)

* Allow a mix of yes, no, parallel, sequential for run all

* Test sequentially and parallel runs

src/extension/kernel.ts

@@ -119,7 +119,8 @@ import { SignedIn } from './signedIn'
 enum ConfirmationItems {
   Yes = 'Yes',
   No = 'No',
-  Skip = 'Skip confirmation and run all',
+  Skip = 'Run all sequentially (skip confirmations)',
+  Parallel = 'Run all in parallel (skip confirmations)',
   Cancel = 'Cancel',
 }
 
@@ -699,9 +700,12 @@ export class Kernel implements Disposable {
       0
     const totalCellsToExecute = cells.length
     let showConfirmPrompt = totalNotebookCells === totalCellsToExecute && totalNotebookCells > 1
+    let cellsToRunParallel: NotebookCell[] = []
+    let cellsToRunSerial: NotebookCell[] = []
     let cellsExecuted = 0
 
-    for (const cell of cells) {
+    outer: for (let i = 0; i < cells.length; i++) {
+      const cell = cells[i]
       const annotations = getAnnotations(cell)
 
       // Skip cells that are not assigned to requested category if requested
@@ -724,30 +728,51 @@ export class Kernel implements Disposable {
           annotations.name || cellText.length > 20 ? `${cellText.slice(0, 20)}...` : cellText
 
         const answer = (await window.showQuickPick(Object.values(ConfirmationItems), {
-          title: `Are you sure you like to run "${cellLabel}"?`,
+          title: `Are you sure you want to run "${cellLabel}"?`,
           ignoreFocusOut: true,
         })) as ConfirmationItems | undefined
 
-        if (answer === ConfirmationItems.No) {
-          continue
+        switch (answer) {
+          case ConfirmationItems.Yes:
+            cellsToRunSerial.push(cell)
+            break
+          case ConfirmationItems.No:
+            continue
+          case ConfirmationItems.Skip:
+            cellsToRunSerial.push(...cells.slice(i))
+            showConfirmPrompt = false
+            break
+          case ConfirmationItems.Parallel:
+            cellsToRunParallel.push(...cells.slice(i))
+            showConfirmPrompt = false
+            break
+          default:
+            break outer
         }
+      } else {
+        cellsToRunSerial = [cell]
+      }
 
-        if (answer === ConfirmationItems.Skip) {
-          showConfirmPrompt = false
-        }
+      for (const cell of cellsToRunSerial) {
+        cellsExecuted++
+        await this._doExecuteCell(cell)
+      }
+      cellsToRunSerial = []
 
-        if (answer === ConfirmationItems.Cancel) {
-          TelemetryReporter.sendTelemetryEvent('cells.executeAll', {
-            'cells.total': totalNotebookCells?.toString(),
-            'cells.executed': cellsExecuted?.toString(),
-          })
-          return
-        }
+      const cellExecs: Promise<void>[] = []
+      for (const cell of cellsToRunParallel) {
+        cellsExecuted++
+        cellExecs.push(this._doExecuteCell(cell))
       }
+      await Promise.all(cellExecs)
+      cellsToRunParallel = []
 
-      await this._doExecuteCell(cell)
-      cellsExecuted++
+      // if we've executed all cells, break out of the loop
+      if (cellsExecuted >= totalCellsToExecute) {
+        break outer
+      }
     }
+
     this.category = undefined
     const uri = cells[0] && cells[0].notebook.uri
     const categories = await getNotebookCategories(this.context, uri)

tests/extension/kernel.test.ts

@@ -298,17 +298,50 @@ suite('_executeAll', async () => {
     })
   })
 
-  test('skips prompt', async () => {
-    const k = new Kernel({} as any)
-    // @ts-ignore readonly
-    window.showQuickPick = vi.fn().mockResolvedValue('Skip confirmation and run all')
-    k['_doExecuteCell'] = vi.fn()
-    await k['_executeAll'](genCells(10))
-    expect(window.showQuickPick).toBeCalledTimes(1)
-    expect(k['_doExecuteCell']).toBeCalledTimes(10)
-    expect(TelemetryReporter.sendTelemetryEvent).lastCalledWith('cells.executeAll', {
-      'cells.executed': '10',
-      'cells.total': '10',
+  suite('run all', () => {
+    // fake gap between cell executions
+    const fakeGap = 10
+    const runAllTest = async (assertGapStdev: (stdev: number) => boolean) => {
+      const timestamps: number[] = []
+      const k = new Kernel({} as any)
+      k['_doExecuteCell'] = vi.fn().mockImplementation(() => {
+        return new Promise<void>((resolve) => {
+          setTimeout(() => {
+            timestamps.push(Date.now())
+            resolve()
+          }, fakeGap)
+        })
+      })
+      await k['_executeAll'](genCells(10))
+      const mean = timestamps.reduce((a, b) => a + b, 0) / timestamps.length
+      const stdev = Math.sqrt(
+        timestamps.map((t) => Math.pow(t - mean, 2)).reduce((a, b) => a + b, 0) / timestamps.length,
+      )
+      assertGapStdev(stdev)
+      expect(window.showQuickPick).toBeCalledTimes(1)
+      expect(k['_doExecuteCell']).toBeCalledTimes(10)
+      expect(TelemetryReporter.sendTelemetryEvent).lastCalledWith('cells.executeAll', {
+        'cells.executed': '10',
+        'cells.total': '10',
+      })
+    }
+
+    test('skips prompt and run sequentially', async () => {
+      // @ts-ignore readonly
+      window.showQuickPick = vi.fn().mockResolvedValue('Run all sequentially (skip confirmations)')
+      await runAllTest((stdev) => {
+        // sequential should at minimum run one stdev slower than fakeGap
+        return stdev > fakeGap
+      })
+    })
+
+    test('skips prompt and run parallel', async () => {
+      // @ts-ignore readonly
+      window.showQuickPick = vi.fn().mockResolvedValue('Run all in parallel (skip confirmations)')
+      await runAllTest((stdev) => {
+        // parallel should be faster than fakeGap since they roughly all run at the same time
+        return stdev < fakeGap
+      })
     })
   })