Add ability to stop experiments running in the queue into stop experi…

…ments command (#3157)

* add ability to stop experiments running in the queue into stop button

* refactor set context from repositories

* add to existing integration test

extension/package.json

@@ -412,8 +412,8 @@
         "category": "DVC"
       },
       {
-        "title": "%command.stopRunningExperiment%",
-        "command": "dvc.stopRunningExperiment",
+        "title": "%command.stopAllRunningExperiments%",
+        "command": "dvc.stopAllRunningExperiments",
         "category": "DVC",
         "icon": "$(debug-stop)"
       },
@@ -828,7 +828,7 @@
           "when": "dvc.commands.available && dvc.project.available"
         },
         {
-          "command": "dvc.stopRunningExperiment",
+          "command": "dvc.stopAllRunningExperiments",
           "when": "dvc.commands.available && dvc.project.available && dvc.experiment.stoppable"
         },
         {
@@ -1000,7 +1000,7 @@
           "when": "dvc.experiments.webview.active && !dvc.experiment.running && dvc.commands.available && !dvc.experiment.checkpoints"
         },
         {
-          "command": "dvc.stopRunningExperiment",
+          "command": "dvc.stopAllRunningExperiments",
           "group": "navigation@0",
           "when": "dvc.experiment.stoppable && dvc.commands.available"
         },
@@ -1219,12 +1219,12 @@
           "group": "navigation@2"
         },
         {
-          "command": "dvc.stopRunningExperiment",
+          "command": "dvc.stopAllRunningExperiments",
           "when": "view == dvc.views.experimentsTree && !dvc.experiments.webview.active && dvc.experiment.stoppable",
           "group": "1_run@1"
         },
         {
-          "command": "dvc.stopRunningExperiment",
+          "command": "dvc.stopAllRunningExperiments",
           "when": "view == dvc.views.experimentsTree && dvc.experiments.webview.active && dvc.experiment.stoppable",
           "group": "navigation@1"
         },

extension/package.nls.json

@@ -56,9 +56,9 @@
   "command.showOutput": "Show DVC Output",
   "command.showPlots": "Show Plots",
   "command.startExperimentsQueue": "Start Processing Queued Experiments",
+  "command.stopAllRunningExperiments": "Stop All Running Experiments",
   "command.stopExperimentsQueue": "Stop Processing Queued Experiments Gracefully (Complete Running)",
   "command.stopQueuedExperiments": "Stop Running Queued Experiment(s)",
-  "command.stopRunningExperiment": "Stop Running Experiment",
   "command.views.experimentsColumnsTree.selectColumns": "Select Columns to Display in the Experiments Table",
   "command.views.experimentsFilterByTree.removeAllFilters": "Remove All Filters From Experiments Table",
   "command.views.experimentsFilterByTree.removeFilter": "Remove Filter From Experiments Table",

extension/src/cli/dvc/constants.ts

@@ -44,6 +44,7 @@ export enum Flag {
   GRANULAR = '--granular',
   JOBS = '-j',
   JSON = '--json',
+  KILL = '--kill',
   NUM_COMMIT = '-n',
   OUTPUT_PATH = '-o',
   SUBDIRECTORY = '--subdir',

extension/src/cli/dvc/executor.ts

@@ -149,8 +149,13 @@ export class DvcExecutor extends DvcCli {
     )
   }
 
-  public queueStop(cwd: string) {
-    return this.executeDvcProcess(cwd, Command.QUEUE, QueueSubCommand.STOP)
+  public queueStop(cwd: string, ...args: Args) {
+    return this.executeDvcProcess(
+      cwd,
+      Command.QUEUE,
+      QueueSubCommand.STOP,
+      ...args
+    )
   }
 
   public remove(cwd: string, ...args: Args) {

extension/src/commands/external.ts

@@ -66,7 +66,7 @@ export enum RegisteredCommands {
   EXPERIMENT_SORTS_REMOVE = 'dvc.removeExperimentsTableSorts',
   EXPERIMENT_SORTS_REMOVE_ALL = 'dvc.views.experimentsSortByTree.removeAllSorts',
   EXPERIMENT_TOGGLE = 'dvc.views.experiments.toggleStatus',
-  STOP_EXPERIMENT = 'dvc.stopRunningExperiment',
+  STOP_EXPERIMENTS = 'dvc.stopAllRunningExperiments',
 
   PLOTS_PATH_TOGGLE = 'dvc.views.plotsPathsTree.toggleStatus',
   PLOTS_SHOW = 'dvc.showPlots',

extension/src/context.ts

@@ -67,9 +67,14 @@ export class Context extends Disposable {
       (experiments: Experiments) => experiments.hasRunningExperiment()
     )
 
+    const hasRunningQueuedExperiment = repositories.some(experiments =>
+      experiments.hasRunningQueuedExperiment()
+    )
+
     void setContextValue(
       ContextKey.EXPERIMENT_STOPPABLE,
-      await this.experiments.hasDvcLiveOnlyExperimentRunning()
+      hasRunningQueuedExperiment ||
+        (await this.experiments.hasDvcLiveOnlyExperimentRunning())
     )
   }
 
@@ -78,12 +83,9 @@ export class Context extends Disposable {
     repositories: Experiments[],
     hasRequirement: (experiments: Experiments) => boolean
   ) {
-    for (const experiments of repositories) {
-      if (hasRequirement(experiments)) {
-        void setContextValue(contextKey, true)
-        return
-      }
-    }
-    void setContextValue(contextKey, false)
+    void setContextValue(
+      contextKey,
+      repositories.some(experiments => hasRequirement(experiments))
+    )
   }
 }

extension/src/experiments/index.ts

@@ -528,6 +528,10 @@ export class Experiments extends BaseRepository<TableData> {
     return this.experiments.hasRunningExperiment()
   }
 
+  public hasRunningQueuedExperiment() {
+    return this.experiments.getRunningQueueTasks().length > 0
+  }
+
   public getFirstThreeColumnOrder() {
     return this.columns.getFirstThreeColumnOrder()
   }

extension/src/experiments/workspace.ts

@@ -431,6 +431,12 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews<
     return pids
   }
 
+  public hasQueuedExperimentsRunning() {
+    return Object.values(this.repositories).some(experiments =>
+      experiments.hasRunningQueuedExperiment()
+    )
+  }
+
   private async pickExpThenRun(
     commandId: CommandId,
     pickFunc: (

extension/src/extension.ts

@@ -49,6 +49,7 @@ import { GitReader } from './cli/git/reader'
 import { Setup } from './setup'
 import { definedAndNonEmpty } from './util/array'
 import { stopProcesses } from './processExecution'
+import { Flag } from './cli/dvc/constants'
 
 export class Extension extends Disposable {
   protected readonly internalCommands: InternalCommands
@@ -194,35 +195,42 @@ export class Extension extends Disposable {
     )
 
     this.dispose.track(
-      commands.registerCommand(RegisteredCommands.STOP_EXPERIMENT, async () => {
-        const stopWatch = new StopWatch()
-        const dvcLiveOnlyPids = await this.experiments.getDvcLiveOnlyPids()
-        const wasRunning =
-          this.dvcRunner.isExperimentRunning() ||
-          definedAndNonEmpty(dvcLiveOnlyPids)
-        try {
-          const [dvcLiveOnlyStopped, runnerStopped] = await Promise.all([
-            stopProcesses(dvcLiveOnlyPids),
-            this.dvcRunner.stop()
-          ])
-
-          const stopped = dvcLiveOnlyStopped && runnerStopped
-          sendTelemetryEvent(
-            RegisteredCommands.STOP_EXPERIMENT,
-            { stopped, wasRunning },
-            {
-              duration: stopWatch.getElapsedTime()
-            }
-          )
-          return stopped
-        } catch (error: unknown) {
-          return sendTelemetryEventAndThrow(
-            RegisteredCommands.STOP_EXPERIMENT,
-            error as Error,
-            stopWatch.getElapsedTime()
-          )
+      commands.registerCommand(
+        RegisteredCommands.STOP_EXPERIMENTS,
+        async () => {
+          const stopWatch = new StopWatch()
+          const dvcLiveOnlyPids = await this.experiments.getDvcLiveOnlyPids()
+          const wasRunning =
+            this.dvcRunner.isExperimentRunning() ||
+            definedAndNonEmpty(dvcLiveOnlyPids) ||
+            this.experiments.hasQueuedExperimentsRunning()
+          try {
+            const allStopped = await Promise.all([
+              stopProcesses(dvcLiveOnlyPids),
+              this.dvcRunner.stop(),
+              ...this.getRoots().map(dvcRoot =>
+                this.dvcExecutor.queueStop(dvcRoot, Flag.KILL)
+              )
+            ])
+
+            const stopped = allStopped.every(Boolean)
+            sendTelemetryEvent(
+              RegisteredCommands.STOP_EXPERIMENTS,
+              { stopped, wasRunning },
+              {
+                duration: stopWatch.getElapsedTime()
+              }
+            )
+            return stopped
+          } catch (error: unknown) {
+            return sendTelemetryEventAndThrow(
+              RegisteredCommands.STOP_EXPERIMENTS,
+              error as Error,
+              stopWatch.getElapsedTime()
+            )
+          }
         }
-      })
+      )
     )
 
     this.internalCommands.registerExternalCommand(

extension/src/telemetry/constants.ts

@@ -167,7 +167,7 @@ export interface IEventNamePropertyMapping {
   [EventName.MODIFY_EXPERIMENT_PARAMS_AND_RESUME]: undefined
   [EventName.MODIFY_EXPERIMENT_PARAMS_AND_RUN]: undefined
   [EventName.MODIFY_EXPERIMENT_PARAMS_RESET_AND_RUN]: undefined
-  [EventName.STOP_EXPERIMENT]: { stopped: boolean; wasRunning: boolean }
+  [EventName.STOP_EXPERIMENTS]: { stopped: boolean; wasRunning: boolean }
 
   [EventName.PLOTS_PATH_TOGGLE]: undefined
   [EventName.PLOTS_SHOW]: undefined

extension/src/test/suite/context.test.ts

@@ -67,12 +67,14 @@ suite('Context Test Suite', () => {
   const buildMockExperiments = (
     filters: FilterDefinition[] = [],
     sorts: SortDefinition[] = [],
-    experimentRunning = false
+    experimentRunning = false,
+    queuedExperimentRunning = false
   ) => {
     return {
       getFilters: () => filters,
       getSorts: () => sorts,
-      hasRunningExperiment: () => experimentRunning
+      hasRunningExperiment: () => experimentRunning,
+      hasRunningQueuedExperiment: () => queuedExperimentRunning
     }
   }
 
@@ -173,7 +175,7 @@ suite('Context Test Suite', () => {
       )
     })
 
-    it('should set the dvc.experiment.stoppable context to whether or not there is a DVCLive only experiment running if an experiment is not running in the runner', async () => {
+    it('should set the dvc.experiment.stoppable context to whether or not there is a DVCLive only experiment running if an experiment is not running in the runner or queue', async () => {
       const {
         executeCommandSpy,
         experimentsChanged,
@@ -220,6 +222,53 @@ suite('Context Test Suite', () => {
       )
     })
 
+    it('should set the dvc.experiment.stoppable context to true if an experiment is running in the queue', async () => {
+      const {
+        executeCommandSpy,
+        experimentsChanged,
+        mockGetDvcRoots,
+        mockGetRepository,
+        mockHasDvcLiveOnlyExperimentRunning,
+        onDidChangeExperiments
+      } = buildContext(false)
+
+      const queuedExperimentRunningEvent = new Promise(resolve =>
+        disposable.track(onDidChangeExperiments(() => resolve(undefined)))
+      )
+
+      const mockDvcRoot = resolve('first', 'root')
+      mockGetDvcRoots.returns([mockDvcRoot])
+      let hasRunningQueuedExperiment = true
+      mockGetRepository.callsFake(() =>
+        buildMockExperiments([], [], false, hasRunningQueuedExperiment)
+      )
+
+      mockHasDvcLiveOnlyExperimentRunning.resolves(false)
+
+      experimentsChanged.fire()
+      await queuedExperimentRunningEvent
+
+      expect(executeCommandSpy).to.be.calledWith(
+        'setContext',
+        'dvc.experiment.stoppable',
+        true
+      )
+      executeCommandSpy.resetHistory()
+
+      const queuedExperimentStoppedEvent = new Promise(resolve =>
+        disposable.track(onDidChangeExperiments(() => resolve(undefined)))
+      )
+      hasRunningQueuedExperiment = false
+      experimentsChanged.fire()
+      await queuedExperimentStoppedEvent
+
+      expect(executeCommandSpy).to.be.calledWith(
+        'setContext',
+        'dvc.experiment.stoppable',
+        false
+      )
+    })
+
     it('should set the dvc.experiments.filtered context correctly depending on whether there are filters applied', async () => {
       const {
         executeCommandSpy,

extension/src/test/suite/extension.test.ts

@@ -25,6 +25,10 @@ import { WorkspaceExperiments } from '../../experiments/workspace'
 import { GitReader } from '../../cli/git/reader'
 import { MIN_CLI_VERSION } from '../../cli/dvc/contract'
 import { ConfigKey, setConfigValue } from '../../vscode/config'
+import { DvcExecutor } from '../../cli/dvc/executor'
+import { dvcDemoPath } from '../util'
+import { Setup } from '../../setup'
+import { Flag } from '../../cli/dvc/constants'
 
 suite('Extension Test Suite', () => {
   const disposable = Disposable.fn()
@@ -240,17 +244,23 @@ suite('Extension Test Suite', () => {
     }).timeout(25000)
   })
 
-  describe('dvc.stopRunningExperiment', () => {
+  describe('dvc.stopAllRunningExperiments', () => {
     it('should send a telemetry event containing properties relating to the event', async () => {
       const duration = 1234
+      const otherRoot = resolve('other', 'root')
       mockDuration(duration)
 
       const mockSendTelemetryEvent = stub(Telemetry, 'sendTelemetryEvent')
+      const mockQueueStop = stub(DvcExecutor.prototype, 'queueStop').resolves(
+        undefined
+      )
+
+      stub(Setup.prototype, 'getRoots').returns([dvcDemoPath, otherRoot])
 
-      await commands.executeCommand(RegisteredCommands.STOP_EXPERIMENT)
+      await commands.executeCommand(RegisteredCommands.STOP_EXPERIMENTS)
 
       expect(mockSendTelemetryEvent).to.be.calledWith(
-        RegisteredCommands.STOP_EXPERIMENT,
+        RegisteredCommands.STOP_EXPERIMENTS,
         {
           stopped: false,
           wasRunning: false
@@ -259,6 +269,9 @@ suite('Extension Test Suite', () => {
           duration
         }
       )
+
+      expect(mockQueueStop).to.be.calledWith(dvcDemoPath, Flag.KILL)
+      expect(mockQueueStop).to.be.calledWith(otherRoot, Flag.KILL)
     })
   })