Combine/enhance commands used to stop experiments (#3840)

* refactor stop experiment commands

* fix stop command in package.json

* use usual pattern for registering commands against the experiments tree

* rework implementation

* fix some bugs found in demos

* fix broken integration test

* add integration test for inline stop action

* add some more defensive code to stop workspace experiment

extension/package.json

@@ -362,8 +362,8 @@
         "category": "DVC"
       },
       {
-        "title": "Stop Running Queued Experiment(s)",
-        "command": "dvc.stopQueuedExperiments",
+        "title": "Stop Running Experiment(s)",
+        "command": "dvc.stopExperiments",
         "category": "DVC"
       },
       {
@@ -429,6 +429,18 @@
         "command": "dvc.showSetup",
         "category": "DVC"
       },
+      {
+        "title": "Stop",
+        "command": "dvc.views.experimentsTree.stopExperiment",
+        "category": "DVC",
+        "icon": "$(debug-stop)"
+      },
+      {
+        "title": "Stop",
+        "command": "dvc.views.experiments.stopExperiment",
+        "category": "DVC",
+        "icon": "$(debug-stop)"
+      },
       {
         "title": "Stop All Running Experiments",
         "command": "dvc.stopAllRunningExperiments",
@@ -827,7 +839,7 @@
           "when": "dvc.commands.available && dvc.project.available"
         },
         {
-          "command": "dvc.stopQueuedExperiments",
+          "command": "dvc.stopExperiments",
           "when": "dvc.commands.available && dvc.project.available"
         },
         {
@@ -914,6 +926,14 @@
           "command": "dvc.views.experiments.showLogs",
           "when": "false"
         },
+        {
+          "command": "dvc.views.experimentsTree.stopExperiment",
+          "when": "false"
+        },
+        {
+          "command": "dvc.views.experiments.stopExperiment",
+          "when": "false"
+        },
         {
           "command": "dvc.views.experimentsFilterByTree.removeAllFilters",
           "when": "false"
@@ -1160,6 +1180,11 @@
           "group": "inline",
           "when": "view == dvc.views.experimentsFilterByTree && dvc.commands.available && viewItem != dvcRoot"
         },
+        {
+          "command": "dvc.views.experimentsTree.stopExperiment",
+          "group": "inline@0",
+          "when": "view == dvc.views.experimentsTree && dvc.commands.available && viewItem == running"
+        },
         {
           "command": "dvc.views.experiments.applyExperiment",
           "group": "inline@1",
@@ -1217,6 +1242,11 @@
           "when": "view == dvc.views.experimentsColumnsTree",
           "group": "navigation@2"
         },
+        {
+          "command": "dvc.stopAllRunningExperiments",
+          "when": "view == dvc.views.experimentsTree && dvc.experiments.webview.active && dvc.experiment.running",
+          "group": "navigation@0"
+        },
         {
           "command": "dvc.runExperiment",
           "when": "view == dvc.views.experimentsTree && !dvc.experiment.running.workspace && !dvc.experiment.checkpoints",
@@ -1237,11 +1267,6 @@
           "when": "view == dvc.views.experimentsTree && !dvc.experiments.webview.active && dvc.experiment.running",
           "group": "1_run@1"
         },
-        {
-          "command": "dvc.stopAllRunningExperiments",
-          "when": "view == dvc.views.experimentsTree && dvc.experiments.webview.active && dvc.experiment.running",
-          "group": "navigation@1"
-        },
         {
           "command": "dvc.views.experimentsTree.selectExperiments",
           "when": "view == dvc.views.experimentsTree && dvc.plots.webview.active",
@@ -1287,11 +1312,6 @@
           "when": "view == dvc.views.experimentsTree",
           "group": "3_queue@3"
         },
-        {
-          "command": "dvc.stopQueuedExperiments",
-          "when": "view == dvc.views.experimentsTree",
-          "group": "3_queue@5"
-        },
         {
           "command": "dvc.addExperimentsTableSort",
           "when": "view == dvc.views.experimentsSortByTree",

extension/src/commands/external.ts

@@ -9,7 +9,6 @@ export enum RegisteredCliCommands {
   EXPERIMENT_RESUME = 'dvc.resumeCheckpointExperiment',
   EXPERIMENT_RUN = 'dvc.runExperiment',
   QUEUE_EXPERIMENT = 'dvc.queueExperiment',
-  QUEUE_KILL = 'dvc.stopQueuedExperiments',
   QUEUE_START = 'dvc.startExperimentsQueue',
   QUEUE_STOP = 'dvc.stopExperimentsQueue',
 
@@ -18,7 +17,6 @@ export enum RegisteredCliCommands {
   EXPERIMENT_VIEW_PUSH = 'dvc.views.experiments.pushExperiment',
   EXPERIMENT_VIEW_REMOVE = 'dvc.views.experiments.removeExperiment',
   EXPERIMENT_VIEW_SHOW_LOGS = 'dvc.views.experiments.showLogs',
-  EXPERIMENT_VIEW_STOP = 'dvc.views.experiments.stopQueueExperiment',
 
   EXPERIMENT_VIEW_QUEUE = 'dvc.views.experiments.queueExperiment',
   EXPERIMENT_VIEW_RESUME = 'dvc.views.experiments.resumeCheckpointExperiment',
@@ -62,7 +60,9 @@ export enum RegisteredCommands {
   EXPERIMENT_SORT_REMOVE = 'dvc.views.experimentsSortByTree.removeSort',
   EXPERIMENT_SORTS_REMOVE = 'dvc.removeExperimentsTableSorts',
   EXPERIMENT_SORTS_REMOVE_ALL = 'dvc.views.experimentsSortByTree.removeAllSorts',
+  EXPERIMENT_STOP = 'dvc.stopExperiments',
   EXPERIMENT_TOGGLE = 'dvc.views.experiments.toggleStatus',
+  EXPERIMENT_VIEW_STOP = 'dvc.views.experiments.stopExperiment',
   STOP_EXPERIMENTS = 'dvc.stopAllRunningExperiments',
 
   PLOTS_PATH_TOGGLE = 'dvc.views.plotsPathsTree.toggleStatus',

extension/src/experiments/commands/register.ts

@@ -115,6 +115,12 @@ const registerExperimentNameCommands = (
     ({ dvcRoot, ids }: { dvcRoot: string; ids: string[] }) =>
       experiments.runCommand(AvailableCommands.EXP_REMOVE, dvcRoot, ...ids)
   )
+
+  internalCommands.registerExternalCommand(
+    RegisteredCommands.EXPERIMENT_VIEW_STOP,
+    ({ dvcRoot, ids }: { dvcRoot: string; ids: string[] }) =>
+      experiments.stopExperiments(dvcRoot, ...ids)
+  )
 }
 
 const registerExperimentInputCommands = (
@@ -200,9 +206,9 @@ const registerExperimentQuickPickCommands = (
       experiments.selectColumns(getDvcRootFromContext(context))
   )
 
-  internalCommands.registerExternalCliCommand(
-    RegisteredCliCommands.QUEUE_KILL,
-    () => experiments.selectQueueTasksToKill()
+  internalCommands.registerExternalCommand(
+    RegisteredCommands.EXPERIMENT_STOP,
+    () => experiments.selectExperimentsToStop()
   )
 
   internalCommands.registerExternalCliCommand(

extension/src/experiments/index.ts

@@ -30,6 +30,7 @@ import { starredSort } from './model/sortBy/constants'
 import { pickSortsToRemove, pickSortToAdd } from './model/sortBy/quickPick'
 import { ColumnsModel } from './columns/model'
 import { ExperimentsData } from './data'
+import { stopWorkspaceExperiment } from './processExecution'
 import {
   Experiment,
   ColumnType,
@@ -378,11 +379,11 @@ export class Experiments extends BaseRepository<TableData> {
     )
   }
 
-  public pickQueueTasksToKill() {
+  public pickRunningExperiments() {
     return pickExperiments(
-      this.experiments.getRunningQueueTasks(),
+      this.experiments.getRunningExperiments(),
       this.getFirstThreeColumnOrder(),
-      Title.SELECT_QUEUE_KILL
+      Title.SELECT_EXPERIMENTS_STOP
     )
   }
 
@@ -486,6 +487,30 @@ export class Experiments extends BaseRepository<TableData> {
     return this.notifyChanged()
   }
 
+  public stopExperiments(ids: string[]) {
+    const { runningInQueueIds, runningInWorkspaceId } =
+      this.experiments.getStopDetails(ids)
+
+    const promises: Promise<string | void>[] = []
+
+    if (runningInQueueIds) {
+      promises.push(
+        this.internalCommands.executeCommand(
+          AvailableCommands.QUEUE_KILL,
+          this.dvcRoot,
+          ...runningInQueueIds
+        )
+      )
+    }
+    if (runningInWorkspaceId) {
+      promises.push(stopWorkspaceExperiment(this.dvcRoot, runningInWorkspaceId))
+    }
+
+    return Toast.showOutput(
+      Promise.all(promises).then(output => output.filter(Boolean).join('\n'))
+    )
+  }
+
   public hasRunningExperiment() {
     return this.experiments.hasRunningExperiment()
   }
@@ -562,12 +587,6 @@ export class Experiments extends BaseRepository<TableData> {
       () => this.getWebview(),
       () => this.notifyChanged(),
       () => this.selectColumns(),
-      (dvcRoot: string, ...ids: string[]) =>
-        this.internalCommands.executeCommand(
-          AvailableCommands.QUEUE_KILL,
-          dvcRoot,
-          ...ids
-        ),
       () =>
         this.internalCommands.executeCommand(
           AvailableCommands.STAGE_LIST,

extension/src/experiments/model/collect.ts

@@ -10,7 +10,8 @@ import {
   Experiment,
   CommitData,
   RunningExperiment,
-  isQueued
+  isQueued,
+  isRunning
 } from '../webview/contract'
 import {
   EXPERIMENT_WORKSPACE_ID,
@@ -205,7 +206,6 @@ const getExecutor = (experiment: Experiment): Executor => {
 }
 
 const collectExecutorInfo = (
-  acc: ExperimentsAccumulator,
   experiment: Experiment,
   executor: ExecutorState
 ): void => {
@@ -221,13 +221,19 @@ const collectExecutorInfo = (
   if (state && state !== ExperimentStatus.SUCCESS) {
     experiment.status = state
   }
+}
 
-  if (experiment.status === ExperimentStatus.RUNNING) {
-    acc.runningExperiments.push({
-      executor: getExecutor(experiment),
-      id: experiment.id
-    })
+const collectRunningExperiment = (
+  acc: ExperimentsAccumulator,
+  experiment: Experiment
+): void => {
+  if (!isRunning(experiment.status)) {
+    return
   }
+  acc.runningExperiments.push({
+    executor: getExecutor(experiment),
+    id: experiment.id
+  })
 }
 
 const collectExpRange = (
@@ -261,7 +267,8 @@ const collectExpRange = (
     experiment.description = `[${name}]`
   }
 
-  collectExecutorInfo(acc, experiment, executor)
+  collectExecutorInfo(experiment, executor)
+  collectRunningExperiment(acc, experiment)
 
   addToMapArray(acc.experimentsByCommit, baseline.id, experiment)
 }
@@ -402,3 +409,33 @@ export const collectOrderedCommitsAndExperiments = (
   }
   return acc
 }
+
+export const collectRunningInQueue = (
+  ids: Set<string>,
+  runningExperiments: RunningExperiment[]
+): string[] | undefined => {
+  const runningInQueueIds = new Set<string>()
+  for (const { executor, id } of runningExperiments) {
+    if (!ids.has(id)) {
+      continue
+    }
+    if (executor === Executor.DVC_TASK) {
+      runningInQueueIds.add(id)
+    }
+  }
+  return runningInQueueIds.size > 0 ? [...runningInQueueIds] : undefined
+}
+
+export const collectRunningInWorkspace = (
+  ids: Set<string>,
+  runningExperiments: RunningExperiment[]
+): string | undefined => {
+  for (const { executor, id } of runningExperiments) {
+    if (!ids.has(id)) {
+      continue
+    }
+    if (executor === EXPERIMENT_WORKSPACE_ID) {
+      return id
+    }
+  }
+}

extension/src/experiments/model/index.ts

@@ -3,7 +3,9 @@ import { SortDefinition, sortExperiments } from './sortBy'
 import { FilterDefinition, filterExperiment, getFilterId } from './filterBy'
 import {
   collectExperiments,
-  collectOrderedCommitsAndExperiments
+  collectOrderedCommitsAndExperiments,
+  collectRunningInQueue,
+  collectRunningInWorkspace
 } from './collect'
 import {
   collectColoredStatus,
@@ -284,7 +286,11 @@ export class ExperimentsModel extends ModelWithPersistence {
       {
         ...this.addDetails(this.workspace),
         hasChildren: false,
-        type: ExperimentType.WORKSPACE
+        type: this.running.some(
+          ({ executor }) => executor === Executor.WORKSPACE
+        )
+          ? ExperimentType.RUNNING
+          : ExperimentType.WORKSPACE
       },
       ...this.commits.map(commit => {
         return {
@@ -334,12 +340,21 @@ export class ExperimentsModel extends ModelWithPersistence {
     )
   }
 
-  public getRunningQueueTasks() {
+  public getRunningExperiments() {
     return this.getExperimentsAndQueued().filter(experiment =>
-      isRunningInQueue(experiment)
+      isRunning(experiment.status)
     )
   }
 
+  public getStopDetails(idsToStop: string[]) {
+    const running = [...this.running]
+    const ids = new Set(idsToStop)
+    return {
+      runningInQueueIds: collectRunningInQueue(ids, running),
+      runningInWorkspaceId: collectRunningInWorkspace(ids, running)
+    }
+  }
+
   public getRowData() {
     return [
       this.addDetails(this.workspace),