Enable the user to specify a number of concurrent jobs when starting the experiments queue (#3048)

* switch runner for executor

* wire up starting the queue with input box

* reorder command and event lists

* refactor input box validator

Author: mattseddon

extension/src/cli/dvc/constants.ts

@@ -25,6 +25,7 @@ export enum Command {
   PLOTS = 'plots',
   PULL = 'pull',
   PUSH = 'push',
+  QUEUE = 'queue',
   REMOVE = 'remove',
   ROOT = 'root',
   PARAMS = 'params',
@@ -41,6 +42,7 @@ export enum Flag {
   ALL_COMMITS = '-A',
   FORCE = '-f',
   GRANULAR = '--granular',
+  JOBS = '-j',
   JSON = '--json',
   NUM_COMMIT = '-n',
   OUTPUT_PATH = '-o',
@@ -59,11 +61,14 @@ export enum ExperimentSubCommand {
   RUN = 'run'
 }
 
+export enum QueueSubCommand {
+  START = 'start'
+}
+
 export enum ExperimentFlag {
   NO_FETCH = '--no-fetch',
   QUEUE = '--queue',
-  RESET = '--reset',
-  RUN_ALL = '--run-all'
+  RESET = '--reset'
 }
 
 export enum GcPreserveFlag {

extension/src/cli/dvc/executor.test.ts

@@ -563,6 +563,28 @@ describe('CliExecutor', () => {
     })
   })
 
+  describe('queueStart', () => {
+    it("should call createProcess with the correct parameters to start the experiment's queue", async () => {
+      const cwd = __dirname
+      const jobs = '91231324'
+
+      const stdout = `Started '${jobs}' new experiments task queue workers.`
+
+      mockedCreateProcess.mockReturnValueOnce(getMockedProcess(stdout))
+
+      const output = await dvcExecutor.queueStart(cwd, jobs)
+
+      expect(output).toStrictEqual(stdout)
+
+      expect(mockedCreateProcess).toHaveBeenCalledWith({
+        args: ['queue', 'start', '-j', jobs],
+        cwd,
+        env: mockedEnv,
+        executable: 'dvc'
+      })
+    })
+  })
+
   describe('remove', () => {
     it('should call createProcess with the correct parameters to remove a .dvc file', async () => {
       const cwd = __dirname

extension/src/cli/dvc/executor.ts

@@ -5,7 +5,8 @@ import {
   ExperimentFlag,
   ExperimentSubCommand,
   Flag,
-  GcPreserveFlag
+  GcPreserveFlag,
+  QueueSubCommand
 } from './constants'
 import { typeCheckCommands } from '..'
 import { ContextKey, setContextValue } from '../../vscode/context'
@@ -25,6 +26,7 @@ export const autoRegisteredCommands = {
   MOVE: 'move',
   PULL: 'pull',
   PUSH: 'push',
+  QUEUE_START: 'queueStart',
   REMOVE: 'remove'
 } as const
 
@@ -126,6 +128,16 @@ export class DvcExecutor extends DvcCli {
     return this.blockAndExecuteProcess(cwd, Command.PUSH, ...args)
   }
 
+  public queueStart(cwd: string, jobs: string) {
+    return this.executeDvcProcess(
+      cwd,
+      Command.QUEUE,
+      QueueSubCommand.START,
+      Flag.JOBS,
+      jobs
+    )
+  }
+
   public remove(cwd: string, ...args: Args) {
     return this.blockAndExecuteProcess(cwd, Command.REMOVE, ...args)
   }

extension/src/cli/dvc/runner.ts

@@ -19,8 +19,7 @@ import { Disposable } from '../../class/dispose'
 
 export const autoRegisteredCommands = {
   EXPERIMENT_RESET_AND_RUN: 'runExperimentReset',
-  EXPERIMENT_RUN: 'runExperiment',
-  EXPERIMENT_RUN_QUEUED: 'runExperimentQueue'
+  EXPERIMENT_RUN: 'runExperiment'
 } as const
 
 export class DvcRunner extends Disposable implements ICli {
@@ -112,10 +111,6 @@ export class DvcRunner extends Disposable implements ICli {
     return this.runExperiment(dvcRoot, ExperimentFlag.RESET, ...args)
   }
 
-  public runExperimentQueue(dvcRoot: string) {
-    return this.runExperiment(dvcRoot, ExperimentFlag.RUN_ALL)
-  }
-
   public async run(cwd: string, ...args: Args) {
     await this.pseudoTerminal.openCurrentInstance()
     if (!this.pseudoTerminal.isBlocked()) {

extension/src/commands/external.ts

@@ -5,13 +5,13 @@ export enum RegisteredCliCommands {
   EXPERIMENT_REMOVE = 'dvc.removeExperiment',
   EXPERIMENT_REMOVE_QUEUE = 'dvc.removeExperimentQueue',
   EXPERIMENT_REMOVE_QUEUED = 'dvc.removeQueuedExperiment',
+  EXPERIMENT_RESET_AND_RUN = 'dvc.resetAndRunCheckpointExperiment',
   EXPERIMENT_RESUME = 'dvc.resumeCheckpointExperiment',
   EXPERIMENT_RUN = 'dvc.runExperiment',
-  EXPERIMENT_RUN_QUEUED = 'dvc.startExperimentsQueue',
-  EXPERIMENT_RESET_AND_RUN = 'dvc.resetAndRunCheckpointExperiment',
   EXPERIMENT_SHARE_AS_BRANCH = 'dvc.shareExperimentAsBranch',
   EXPERIMENT_SHARE_AS_COMMIT = 'dvc.shareExperimentAsCommit',
   QUEUE_EXPERIMENT = 'dvc.queueExperiment',
+  QUEUE_START = 'dvc.startExperimentsQueue',
 
   EXPERIMENT_VIEW_APPLY = 'dvc.views.experiments.applyExperiment',
   EXPERIMENT_VIEW_BRANCH = 'dvc.views.experiments.branchExperiment',

extension/src/experiments/commands/register.ts

@@ -294,8 +294,17 @@ const registerExperimentRunCommands = (
   )
 
   internalCommands.registerExternalCliCommand(
-    RegisteredCliCommands.EXPERIMENT_RUN_QUEUED,
-    () => experiments.getCwdThenRun(AvailableCommands.EXPERIMENT_RUN_QUEUED)
+    RegisteredCliCommands.QUEUE_START,
+    () =>
+      experiments.getCwdIntegerInputAndRun(
+        AvailableCommands.QUEUE_START,
+        Title.ENTER_EXPERIMENT_WORKER_COUNT,
+        {
+          prompt:
+            'Input the maximum number of concurrent queue workers to start.',
+          value: '1'
+        }
+      )
   )
 
   internalCommands.registerExternalCommand(

extension/src/experiments/workspace.ts

@@ -6,7 +6,7 @@ import { Args, DVCLIVE_ONLY_RUNNING_SIGNAL_FILE } from '../cli/dvc/constants'
 import { CommandId, InternalCommands } from '../commands/internal'
 import { ResourceLocator } from '../resourceLocator'
 import { Toast } from '../vscode/toast'
-import { getInput } from '../vscode/inputBox'
+import { getInput, getPositiveIntegerInput } from '../vscode/inputBox'
 import { BaseWorkspaceWebviews } from '../webview/workspace'
 import { Title } from '../vscode/title'
 import { ContextKey, setContextValue } from '../vscode/context'
@@ -264,6 +264,25 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews<
     return this.runCommand(commandId, cwd, name)
   }
 
+  public async getCwdIntegerInputAndRun(
+    commandId: CommandId,
+    title: Title,
+    options: { prompt: string; value: string }
+  ) {
+    const cwd = await this.getFocusedOrOnlyOrPickProject()
+    if (!cwd) {
+      return
+    }
+
+    const integer = await getPositiveIntegerInput(title, options)
+
+    if (!integer) {
+      return
+    }
+
+    return this.runCommand(commandId, cwd, integer)
+  }
+
   public runCommand(commandId: CommandId, cwd: string, ...args: Args) {
     return Toast.showOutput(
       this.internalCommands.executeCommand(commandId, cwd, ...args)

extension/src/telemetry/constants.ts

@@ -137,30 +137,30 @@ export interface IEventNamePropertyMapping {
   [EventName.EXPERIMENT_REMOVE_QUEUED]: undefined
   [EventName.EXPERIMENT_RESUME]: undefined
   [EventName.EXPERIMENT_RUN]: undefined
-  [EventName.EXPERIMENT_RUN_QUEUED]: undefined
   [EventName.EXPERIMENT_RESET_AND_RUN]: undefined
   [EventName.EXPERIMENT_SELECT]: undefined
   [EventName.EXPERIMENT_SHARE_AS_BRANCH]: undefined
   [EventName.EXPERIMENT_SHARE_AS_COMMIT]: undefined
   [EventName.EXPERIMENT_SHOW]: undefined
-  [EventName.EXPERIMENT_SORT_ADD]: undefined
   [EventName.EXPERIMENT_SORT_ADD_STARRED]: undefined
+  [EventName.EXPERIMENT_SORT_ADD]: undefined
   [EventName.EXPERIMENT_SORT_REMOVE]: undefined
-  [EventName.EXPERIMENT_SORTS_REMOVE]: undefined
   [EventName.EXPERIMENT_SORTS_REMOVE_ALL]: undefined
+  [EventName.EXPERIMENT_SORTS_REMOVE]: undefined
+  [EventName.EXPERIMENT_TOGGLE]: undefined
   [EventName.EXPERIMENT_VIEW_APPLY]: undefined
   [EventName.EXPERIMENT_VIEW_BRANCH]: undefined
   [EventName.EXPERIMENT_VIEW_REMOVE]: undefined
   [EventName.EXPERIMENT_VIEW_SHARE_AS_BRANCH]: undefined
   [EventName.EXPERIMENT_VIEW_SHARE_AS_COMMIT]: undefined
-  [EventName.EXPERIMENT_TOGGLE]: undefined
+  [EventName.QUEUE_EXPERIMENT]: undefined
+  [EventName.QUEUE_START]: undefined
 
   [EventName.EXPERIMENT_VIEW_QUEUE]: undefined
   [EventName.EXPERIMENT_VIEW_RESUME]: undefined
   [EventName.EXPERIMENT_VIEW_RUN]: undefined
   [EventName.EXPERIMENT_VIEW_RESET_AND_RUN]: undefined
 
-  [EventName.QUEUE_EXPERIMENT]: undefined
   [EventName.MODIFY_EXPERIMENT_PARAMS_AND_QUEUE]: undefined
   [EventName.MODIFY_EXPERIMENT_PARAMS_AND_RESUME]: undefined
   [EventName.MODIFY_EXPERIMENT_PARAMS_AND_RUN]: undefined

extension/src/test/suite/experiments/workspace.test.ts

@@ -439,17 +439,23 @@ suite('Workspace Experiments Test Suite', () => {
 
   describe('dvc.startExperimentsQueue', () => {
     it('should be able to execute all experiments in the run queue', async () => {
-      const mockRunExperimentQueue = stub(
-        DvcRunner.prototype,
-        'runExperimentQueue'
-      ).resolves(undefined)
+      const mockQueueStart = stub(DvcExecutor.prototype, 'queueStart').resolves(
+        undefined
+      )
+
+      const dDosNumberOfJobs = '10000'
+
+      const mockInputBox = stub(window, 'showInputBox').resolves(
+        dDosNumberOfJobs
+      )
 
       stubWorkspaceExperimentsGetters(dvcDemoPath)
 
-      await commands.executeCommand(RegisteredCliCommands.EXPERIMENT_RUN_QUEUED)
+      await commands.executeCommand(RegisteredCliCommands.QUEUE_START)
 
-      expect(mockRunExperimentQueue).to.be.calledOnce
-      expect(mockRunExperimentQueue).to.be.calledWith(dvcDemoPath)
+      expect(mockQueueStart).to.be.calledOnce
+      expect(mockQueueStart).to.be.calledWith(dvcDemoPath, dDosNumberOfJobs)
+      expect(mockInputBox)
     })
   })
 

extension/src/util/number.ts

@@ -10,18 +10,17 @@ export const formatNumber = (value: number): string => {
   return automatic
 }
 
+export const isValidStringInteger = (
+  input: string | undefined
+): input is string =>
+  !!input && Number.parseInt(input) === Number.parseFloat(input)
+
 export const createValidInteger = (
   input: string | number | undefined
 ): number | undefined => {
-  if (!input) {
-    return
-  }
-
   if (typeof input === 'number') {
     return validateNumericInteger(input)
   }
 
-  return Number.parseInt(input) === Number.parseFloat(input)
-    ? Number.parseInt(input)
-    : undefined
+  return isValidStringInteger(input) ? Number.parseInt(input) : undefined
 }