Enable the user to specify a number of concurrent jobs when starting the experiments queue

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

@@ -7,7 +7,7 @@ export enum RegisteredCliCommands {
   EXPERIMENT_REMOVE_QUEUED = 'dvc.removeQueuedExperiment',
   EXPERIMENT_RESUME = 'dvc.resumeCheckpointExperiment',
   EXPERIMENT_RUN = 'dvc.runExperiment',
-  EXPERIMENT_RUN_QUEUED = 'dvc.startExperimentsQueue',
+  QUEUE_START = 'dvc.startExperimentsQueue',
   EXPERIMENT_RESET_AND_RUN = 'dvc.resetAndRunCheckpointExperiment',
   EXPERIMENT_SHARE_AS_BRANCH = 'dvc.shareExperimentAsBranch',
   EXPERIMENT_SHARE_AS_COMMIT = 'dvc.shareExperimentAsCommit',

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,7 +137,7 @@ export interface IEventNamePropertyMapping {
   [EventName.EXPERIMENT_REMOVE_QUEUED]: undefined
   [EventName.EXPERIMENT_RESUME]: undefined
   [EventName.EXPERIMENT_RUN]: undefined
-  [EventName.EXPERIMENT_RUN_QUEUED]: undefined
+  [EventName.QUEUE_START]: undefined
   [EventName.EXPERIMENT_RESET_AND_RUN]: undefined
   [EventName.EXPERIMENT_SELECT]: undefined
   [EventName.EXPERIMENT_SHARE_AS_BRANCH]: 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/vscode/inputBox.ts

@@ -18,3 +18,27 @@ export const getValidInput = (
     validateInput,
     value: options?.value
   })
+
+export const getPositiveIntegerInput = async (
+  title: Title,
+  options: { prompt: string; value: string }
+) => {
+  const input = await getValidInput(
+    title,
+    val => {
+      const number = Number(val)
+
+      if (!Number.isInteger(number) || number <= 0) {
+        return 'Input needs to be a positive integer'
+      }
+
+      return ''
+    },
+    options
+  )
+
+  if (!input) {
+    return
+  }
+  return Number.parseInt(input).toString()
+}

extension/src/vscode/title.ts

@@ -2,6 +2,7 @@ export enum Title {
   CHOOSE_RESOURCES = 'Choose Resources to Add to the Dataset',
   ENTER_BRANCH_NAME = 'Enter a Name for the New Branch',
   ENTER_COMMIT_MESSAGE = 'Enter a Commit Message',
+  ENTER_EXPERIMENT_WORKER_COUNT = 'Enter the Number of Queue Workers',
   ENTER_FILTER_VALUE = 'Enter a Filter Value',
   ENTER_RELATIVE_DESTINATION = 'Enter a Destination Relative to the Root',
   GARBAGE_COLLECT_EXPERIMENTS = 'Garbage Collect Experiments',