Experiments Commands #2: Add "Run All Queued Experiments" Command

extension/package.json

@@ -60,6 +60,16 @@
         "command": "dvc.queueExperiment",
         "category": "DVC"
       },
+      {
+        "title": "%command.runQueuedExperiments%",
+        "command": "dvc.runQueuedExperiments",
+        "category": "DVC"
+      },
+      {
+        "title": "%command.experimentGarbageCollect%",
+        "command": "dvc.experimentGarbageCollect",
+        "category": "DVC"
+      },
       {
         "title": "%command.selectDvcPath%",
         "command": "dvc.selectDvcPath",

extension/package.nls.json

@@ -13,6 +13,8 @@
   "command.pushTarget": "Push Target",
   "command.runExperiment": "Run Experiment",
   "command.queueExperiment": "Queue Experiment",
+  "command.runQueuedExperiments": "Run Queued Experiments",
+  "command.experimentGarbageCollect": "Garbage Collect Experiments",
   "command.selectDvcPath": "Select DVC CLI Path",
   "command.showExperiments": "Show Experiments",
   "config.dvcPath.description": "Call DVC from this path. Follows Python Extension when blank.",

extension/src/IntegratedTerminal.test.ts

@@ -1,4 +1,8 @@
-import { IntegratedTerminal, runExperiment } from './IntegratedTerminal'
+import {
+  IntegratedTerminal,
+  runExperiment,
+  runQueuedExperiments
+} from './IntegratedTerminal'
 
 describe('runExperiment', () => {
   it('should run the correct command in the IntegratedTerminal', async () => {
@@ -9,6 +13,19 @@ describe('runExperiment', () => {
     const undef = await runExperiment()
     expect(undef).toBeUndefined()
 
-    expect(terminalSpy).toBeCalledWith('dvc exp run ')
+    expect(terminalSpy).toBeCalledWith('dvc exp run')
+  })
+})
+
+describe('runQueuedExperiments', () => {
+  it('should run the correct command in the IntegratedTerminal', async () => {
+    const terminalSpy = jest
+      .spyOn(IntegratedTerminal, 'run')
+      .mockResolvedValueOnce(undefined)
+
+    const returnValue = await runQueuedExperiments()
+    expect(returnValue).toBeUndefined()
+
+    expect(terminalSpy).toBeCalledWith('dvc exp run --run-all')
   })
 })

extension/src/IntegratedTerminal.ts

@@ -1,5 +1,5 @@
 import { Terminal, window, workspace } from 'vscode'
-import { getRunExperimentCommand } from './cli/commands'
+import { Commands } from './cli/commands'
 import { getReadyPythonExtension } from './extensions/python'
 import { delay } from './util'
 
@@ -21,6 +21,10 @@ export class IntegratedTerminal {
     return currentTerminal?.sendText(command, true)
   }
 
+  static runCommand = async (command: string): Promise<void> => {
+    return IntegratedTerminal.run(`dvc ${command}`)
+  }
+
   static dispose = (): void => {
     const currentTerminal = IntegratedTerminal.instance
     if (currentTerminal) {
@@ -65,6 +69,9 @@ export class IntegratedTerminal {
 }
 
 export const runExperiment = (): Promise<void> => {
-  const runExperimentCommand = getRunExperimentCommand()
-  return IntegratedTerminal.run(runExperimentCommand)
+  return IntegratedTerminal.runCommand(Commands.EXPERIMENT_RUN)
+}
+
+export const runQueuedExperiments = (): Promise<void> => {
+  return IntegratedTerminal.runCommand(Commands.EXPERIMENT_RUN_ALL)
 }

extension/src/__mocks__/vscode.ts

@@ -8,7 +8,8 @@ export const ThemeColor = jest.fn()
 export const Terminal = jest.fn()
 export const window = {
   showInformationMessage: jest.fn(),
-  showErrorMessage: jest.fn()
+  showErrorMessage: jest.fn(),
+  showQuickPick: jest.fn()
 }
 export const workspace = {
   workspaceFolders: [

extension/src/cli/commands.ts

@@ -2,21 +2,22 @@ export enum Commands {
   ADD = 'add',
   CHECKOUT = 'checkout',
   CHECKOUT_RECURSIVE = 'checkout --recursive',
-  EXPERIMENT_RUN = 'exp run',
-  EXPERIMENT_SHOW = 'exp show --show-json',
   INITIALIZE_SUBDIRECTORY = 'init --subdir',
   PULL = 'pull',
   PUSH = 'push',
   STATUS = 'status --show-json',
-  QUEUE_EXPERIMENT = 'exp run --queue'
-}
-
-const getCliCommand = (command: string, ...options: string[]): string => {
-  return `dvc ${command} ${options.join(' ')}`
+  EXPERIMENT_RUN = 'exp run',
+  EXPERIMENT_SHOW = 'exp show --show-json',
+  EXPERIMENT_QUEUE = 'exp run --queue',
+  EXPERIMENT_RUN_ALL = 'exp run --run-all',
+  EXPERIMENT_GC = 'exp gc -f -w'
 }
 
-export const getRunExperimentCommand = (): string => {
-  return getCliCommand(Commands.EXPERIMENT_RUN)
+export enum GcPreserveFlag {
+  ALL_BRANCHES = '--all-branches',
+  ALL_TAGS = '--all-tags',
+  ALL_COMMITS = '--all-commits',
+  QUEUED = '--queued'
 }
 
 export const getCommandWithTarget = (

extension/src/cli/index.test.ts

@@ -1,10 +1,15 @@
 import { Config } from '../Config'
-import { queueExperimentCommand } from './index'
+import {
+  GcQuickPickItem,
+  experimentGcCommand,
+  queueExperimentCommand
+} from './index'
 import { mocked } from 'ts-jest/utils'
 import { execPromise } from '../util'
 import { basename, resolve } from 'path'
 import { addTarget } from '.'
-import { window } from 'vscode'
+import { QuickPickOptions, window } from 'vscode'
+import { GcPreserveFlag } from './commands'
 
 jest.mock('fs')
 jest.mock('../util')
@@ -13,6 +18,12 @@ jest.mock('vscode')
 const mockedExecPromise = mocked(execPromise)
 const mockedShowErrorMessage = mocked(window.showErrorMessage)
 const mockedShowInformationMessage = mocked(window.showInformationMessage)
+const mockedShowQuickPick = mocked<
+  (
+    items: GcQuickPickItem[],
+    options: QuickPickOptions
+  ) => Thenable<GcQuickPickItem[] | undefined>
+>(window.showQuickPick)
 
 beforeEach(() => {
   jest.resetAllMocks()
@@ -54,17 +65,126 @@ describe('queueExperimentCommand', () => {
     cwd: resolve()
   } as unknown) as Config
 
-  test('it displays an info message with the contents of stdout when the command succeeds', async () => {
+  it('displays an info message with the contents of stdout when the command succeeds', async () => {
     const stdout = 'Example stdout that will be resolved literally\n'
     mockedExecPromise.mockResolvedValue({ stdout, stderr: '' })
     await queueExperimentCommand(exampleConfig)
     expect(mockedShowInformationMessage).toBeCalledWith(stdout)
   })
 
-  test('it displays an error message with the contents of stderr when the command fails', async () => {
+  it('displays an error message with the contents of stderr when the command fails', async () => {
     const stderr = 'Example stderr that will be resolved literally\n'
     mockedExecPromise.mockRejectedValue({ stderr, stdout: '' })
     await queueExperimentCommand(exampleConfig)
     expect(mockedShowErrorMessage).toBeCalledWith(stderr)
   })
 })
+
+describe('experimentGcCommand', () => {
+  const exampleConfig = ({
+    dvcPath: 'dvc',
+    cwd: resolve()
+  } as unknown) as Config
+
+  it('invokes a QuickPick with snapshotted options', async () => {
+    await experimentGcCommand(exampleConfig)
+    expect(mockedShowQuickPick.mock.calls).toMatchInlineSnapshot(`
+      Array [
+        Array [
+          Array [
+            Object {
+              "detail": "Preserve Experiments derived from all Git branches",
+              "flag": "--all-branches",
+              "label": "All Branches",
+            },
+            Object {
+              "detail": "Preserve Experiments derived from all Git tags",
+              "flag": "--all-tags",
+              "label": "All Tags",
+            },
+            Object {
+              "detail": "Preserve Experiments derived from all Git commits",
+              "flag": "--all-commits",
+              "label": "All Commits",
+            },
+            Object {
+              "detail": "Preserve all queued Experiments",
+              "flag": "--queued",
+              "label": "Queued Experiments",
+            },
+          ],
+          Object {
+            "canPickMany": true,
+            "placeHolder": "Select which Experiments to preserve",
+          },
+        ],
+      ]
+    `)
+  })
+
+  it('executes the proper command given a mocked selection', async () => {
+    mockedShowQuickPick.mockResolvedValue([
+      {
+        detail: 'Preserve Experiments derived from all Git tags',
+        flag: GcPreserveFlag.ALL_TAGS,
+        label: 'All Tags'
+      },
+      {
+        detail: 'Preserve Experiments derived from all Git commits',
+        flag: GcPreserveFlag.ALL_COMMITS,
+        label: 'All Commits'
+      }
+    ])
+
+    await experimentGcCommand(exampleConfig)
+
+    expect(mockedExecPromise).toBeCalledWith(
+      'dvc exp gc -f -w --all-tags --all-commits',
+      {
+        cwd: exampleConfig.workspaceRoot
+      }
+    )
+  })
+
+  it('reports stdout from the executed command via showInformationMessage', async () => {
+    const stdout = 'example stdout that will be passed on'
+    mockedShowQuickPick.mockResolvedValue([])
+    mockedExecPromise.mockResolvedValue({ stdout, stderr: '' })
+    await experimentGcCommand(exampleConfig)
+    expect(mockedShowInformationMessage).toBeCalledWith(stdout)
+  })
+
+  it('reports stderr from the executed command via showInformationMessage', async () => {
+    const stderr = 'example stderr that will be passed on'
+    mockedShowQuickPick.mockResolvedValue([])
+    mockedExecPromise.mockRejectedValue({ stderr, stdout: '' })
+    await experimentGcCommand(exampleConfig)
+    expect(mockedShowErrorMessage).toBeCalledWith(stderr)
+  })
+
+  it('reports the message from a non-shell Exception', async () => {
+    const message = 'example message that will be passed on'
+    mockedShowQuickPick.mockResolvedValue([])
+    mockedExecPromise.mockImplementation(() => {
+      throw new Error(message)
+    })
+    await experimentGcCommand(exampleConfig)
+    expect(mockedShowErrorMessage).toBeCalledWith(message)
+  })
+
+  it('executes the proper default command given no selections', async () => {
+    mockedShowQuickPick.mockResolvedValue([])
+
+    await experimentGcCommand(exampleConfig)
+
+    expect(mockedExecPromise).toBeCalledWith('dvc exp gc -f -w', {
+      cwd: exampleConfig.workspaceRoot
+    })
+  })
+
+  it('does not execute a command if the QuickPick is dismissed', async () => {
+    mockedShowQuickPick.mockResolvedValue(undefined)
+    await experimentGcCommand(exampleConfig)
+    expect(mockedExecPromise).not.toBeCalled()
+  })
+})

extension/src/cli/index.ts

@@ -1,14 +1,15 @@
 import { basename, dirname } from 'path'
-import { commands, window } from 'vscode'
+import { commands, QuickPickItem, window } from 'vscode'
 import { Disposer } from '@hediet/std/disposable'
 import { Config } from '../Config'
-import { Commands, getCommandWithTarget } from './commands'
+import { Commands, GcPreserveFlag, getCommandWithTarget } from './commands'
 import {
   execCommand,
   initializeDirectory,
   checkout,
   checkoutRecursive,
-  queueExperiment
+  queueExperiment,
+  experimentGarbageCollect
 } from './reader'
 
 const runTargetCommand = async (
@@ -25,8 +26,7 @@ const runTargetCommand = async (
   const target = basename(fsPath)
   const commandWithTarget = getCommandWithTarget(command, target)
 
-  const { stdout } = await execCommand({ cwd, cliPath }, commandWithTarget)
-  return stdout
+  return execCommand({ cwd, cliPath }, commandWithTarget)
 }
 
 export const queueExperimentCommand = async (config: Config) => {
@@ -41,6 +41,54 @@ export const queueExperimentCommand = async (config: Config) => {
     return window.showErrorMessage(e.stderr || e.message)
   }
 }
+
+export interface GcQuickPickItem extends QuickPickItem {
+  flag: GcPreserveFlag
+}
+
+export const experimentGcCommand = async (config: Config) => {
+  const quickPickResult = await window.showQuickPick<GcQuickPickItem>(
+    [
+      {
+        label: 'All Branches',
+        detail: 'Preserve Experiments derived from all Git branches',
+        flag: GcPreserveFlag.ALL_BRANCHES
+      },
+      {
+        label: 'All Tags',
+        detail: 'Preserve Experiments derived from all Git tags',
+        flag: GcPreserveFlag.ALL_TAGS
+      },
+      {
+        label: 'All Commits',
+        detail: 'Preserve Experiments derived from all Git commits',
+        flag: GcPreserveFlag.ALL_COMMITS
+      },
+      {
+        label: 'Queued Experiments',
+        detail: 'Preserve all queued Experiments',
+        flag: GcPreserveFlag.QUEUED
+      }
+    ],
+    { canPickMany: true, placeHolder: 'Select which Experiments to preserve' }
+  )
+
+  if (quickPickResult) {
+    try {
+      const stdout = await experimentGarbageCollect(
+        {
+          cwd: config.workspaceRoot,
+          cliPath: config.dvcPath
+        },
+        quickPickResult.map(({ flag }) => flag)
+      )
+      window.showInformationMessage(stdout)
+    } catch (e) {
+      window.showErrorMessage(e.stderr || e.message)
+    }
+  }
+}
+
 export const addTarget = async (options: {
   fsPath: string
   cliPath: string | undefined
@@ -97,8 +145,14 @@ export const registerCommands = (config: Config, disposer: Disposer) => {
   )
 
   disposer.track(
-    commands.registerCommand('dvc.queueExperiment', async () => {
+    commands.registerCommand('dvc.queueExperiment', () => {
       return queueExperimentCommand(config)
     })
   )
+
+  disposer.track(
+    commands.registerCommand('dvc.experimentGarbageCollect', () => {
+      return experimentGcCommand(config)
+    })
+  )
 }