Add new stop queue action (#3054)

* add kill and stop queue

* remove kill as it has a different function

Author: mattseddon

extension/package.json

@@ -341,6 +341,11 @@
         "category": "DVC",
         "icon": "$(run-all)"
       },
+      {
+        "title": "%command.stopExperimentsQueue%",
+        "command": "dvc.stopExperimentsQueue",
+        "category": "DVC"
+      },
       {
         "title": "%command.resetAndRunCheckpointExperiment%",
         "command": "dvc.resetAndRunCheckpointExperiment",
@@ -775,7 +780,11 @@
         },
         {
           "command": "dvc.startExperimentsQueue",
-          "when": "dvc.commands.available && dvc.project.available && !dvc.experiment.running"
+          "when": "dvc.commands.available && dvc.project.available"
+        },
+        {
+          "command": "dvc.stopExperimentsQueue",
+          "when": "dvc.commands.available && dvc.project.available"
         },
         {
           "command": "dvc.selectForCompare",
@@ -1257,9 +1266,14 @@
         },
         {
           "command": "dvc.startExperimentsQueue",
-          "when": "view == dvc.views.experimentsTree && !dvc.experiment.running",
+          "when": "view == dvc.views.experimentsTree",
           "group": "3_queue@2"
         },
+        {
+          "command": "dvc.stopExperimentsQueue",
+          "when": "view == dvc.views.experimentsTree",
+          "group": "3_queue@3"
+        },
         {
           "command": "dvc.addExperimentsTableSort",
           "when": "view == dvc.views.experimentsSortByTree",

extension/package.nls.json

@@ -45,6 +45,7 @@
   "command.runExperiment": "Run Experiment",
   "command.resumeCheckpointExperiment": "Resume Experiment",
   "command.startExperimentsQueue": "Start the Experiments Queue",
+  "command.stopExperimentsQueue": "Stop the Experiments Queue",
   "command.resetAndRunCheckpointExperiment": "Run Experiment",
   "command.selectForCompare": "Select for Compare",
   "command.selectFocusedProjects": "Select Project(s) to Focus (set dvc.focusedProjects)",

extension/src/cli/dvc/constants.ts

@@ -62,7 +62,8 @@ export enum ExperimentSubCommand {
 }
 
 export enum QueueSubCommand {
-  START = 'start'
+  START = 'start',
+  STOP = 'stop'
 }
 
 export enum ExperimentFlag {

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

@@ -585,6 +585,27 @@ describe('CliExecutor', () => {
     })
   })
 
+  describe('queueStop', () => {
+    it("should call createProcess with the correct parameters to stop the experiment's queue", async () => {
+      const cwd = __dirname
+
+      const stdout = 'Queue workers will stop after running tasks finish.'
+
+      mockedCreateProcess.mockReturnValueOnce(getMockedProcess(stdout))
+
+      const output = await dvcExecutor.queueStop(cwd)
+
+      expect(output).toStrictEqual(stdout)
+
+      expect(mockedCreateProcess).toHaveBeenCalledWith({
+        args: ['queue', 'stop'],
+        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

@@ -27,6 +27,7 @@ export const autoRegisteredCommands = {
   PULL: 'pull',
   PUSH: 'push',
   QUEUE_START: 'queueStart',
+  QUEUE_STOP: 'queueStop',
   REMOVE: 'remove'
 } as const
 
@@ -138,6 +139,10 @@ export class DvcExecutor extends DvcCli {
     )
   }
 
+  public queueStop(cwd: string) {
+    return this.executeDvcProcess(cwd, Command.QUEUE, QueueSubCommand.STOP)
+  }
+
   public remove(cwd: string, ...args: Args) {
     return this.blockAndExecuteProcess(cwd, Command.REMOVE, ...args)
   }

extension/src/commands/external.ts

@@ -12,6 +12,7 @@ export enum RegisteredCliCommands {
   EXPERIMENT_SHARE_AS_COMMIT = 'dvc.shareExperimentAsCommit',
   QUEUE_EXPERIMENT = 'dvc.queueExperiment',
   QUEUE_START = 'dvc.startExperimentsQueue',
+  QUEUE_STOP = 'dvc.stopExperimentsQueue',
 
   EXPERIMENT_VIEW_APPLY = 'dvc.views.experiments.applyExperiment',
   EXPERIMENT_VIEW_BRANCH = 'dvc.views.experiments.branchExperiment',

extension/src/experiments/commands/register.ts

@@ -27,6 +27,11 @@ const registerExperimentCwdCommands = (
       )
   )
 
+  internalCommands.registerExternalCliCommand(
+    RegisteredCliCommands.QUEUE_STOP,
+    () => experiments.getCwdThenReport(AvailableCommands.QUEUE_STOP)
+  )
+
   internalCommands.registerExternalCliCommand(
     RegisteredCliCommands.MODIFY_EXPERIMENT_PARAMS_AND_QUEUE,
     () =>

extension/src/telemetry/constants.ts

@@ -155,6 +155,7 @@ export interface IEventNamePropertyMapping {
   [EventName.EXPERIMENT_VIEW_SHARE_AS_COMMIT]: undefined
   [EventName.QUEUE_EXPERIMENT]: undefined
   [EventName.QUEUE_START]: undefined
+  [EventName.QUEUE_STOP]: undefined
 
   [EventName.EXPERIMENT_VIEW_QUEUE]: undefined
   [EventName.EXPERIMENT_VIEW_RESUME]: undefined

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

@@ -438,7 +438,7 @@ suite('Workspace Experiments Test Suite', () => {
   })
 
   describe('dvc.startExperimentsQueue', () => {
-    it('should be able to execute all experiments in the run queue', async () => {
+    it('should be able to start the experiments queue with the selected number of workers', async () => {
       const mockQueueStart = stub(DvcExecutor.prototype, 'queueStart').resolves(
         undefined
       )
@@ -454,11 +454,29 @@ suite('Workspace Experiments Test Suite', () => {
       await commands.executeCommand(RegisteredCliCommands.QUEUE_START)
 
       expect(mockQueueStart).to.be.calledOnce
-      expect(mockQueueStart).to.be.calledWith(dvcDemoPath, dDosNumberOfJobs)
+      expect(mockQueueStart).to.be.calledWithExactly(
+        dvcDemoPath,
+        dDosNumberOfJobs
+      )
       expect(mockInputBox)
     })
   })
 
+  describe('dvc.stopExperimentsQueue', () => {
+    it('should be able to stop the experiments queue', async () => {
+      const mockQueueStop = stub(DvcExecutor.prototype, 'queueStop').resolves(
+        undefined
+      )
+
+      stubWorkspaceExperimentsGetters(dvcDemoPath)
+
+      await commands.executeCommand(RegisteredCliCommands.QUEUE_STOP)
+
+      expect(mockQueueStop).to.be.calledOnce
+      expect(mockQueueStop).to.be.calledWithExactly(dvcDemoPath)
+    })
+  })
+
   describe('dvc.applyExperiment', () => {
     it('should ask the user to pick an experiment and then apply that experiment to the workspace', async () => {
       const selectedExperiment = 'test-branch'