refactor stop experiment commands

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',
 
@@ -62,6 +61,7 @@ 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',
   STOP_EXPERIMENTS = 'dvc.stopAllRunningExperiments',
 

extension/src/experiments/commands/register.ts

@@ -200,9 +200,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,
@@ -41,7 +42,11 @@ import { DecorationProvider } from './model/decorationProvider'
 import { starredFilter } from './model/filterBy/constants'
 import { ResourceLocator } from '../resourceLocator'
 import { AvailableCommands, InternalCommands } from '../commands/internal'
-import { ExpShowOutput } from '../cli/dvc/contract'
+import {
+  EXPERIMENT_WORKSPACE_ID,
+  Executor,
+  ExpShowOutput
+} from '../cli/dvc/contract'
 import { ViewKey } from '../webview/constants'
 import { BaseRepository } from '../webview/repository'
 import { Title } from '../vscode/title'
@@ -378,11 +383,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 +491,36 @@ export class Experiments extends BaseRepository<TableData> {
     return this.notifyChanged()
   }
 
+  // eslint-disable-next-line sonarjs/cognitive-complexity
+  public stopExperiments(ids: string[]) {
+    const experiments = this.experiments.getExperiments()
+    const idSet = new Set(ids)
+    let pidRunningInWorkspace
+    const runningInQueueIds = new Set<string>()
+    for (const { executor, id, executorPid } of experiments) {
+      if (!idSet.has(id)) {
+        continue
+      }
+      if (executor === EXPERIMENT_WORKSPACE_ID && executorPid) {
+        pidRunningInWorkspace = executorPid
+      }
+      if (executor === Executor.DVC_TASK) {
+        runningInQueueIds.add(id)
+      }
+    }
+
+    if (runningInQueueIds.size > 0) {
+      void this.internalCommands.executeCommand(
+        AvailableCommands.QUEUE_KILL,
+        this.dvcRoot,
+        ...runningInQueueIds
+      )
+      if (pidRunningInWorkspace) {
+        void stopWorkspaceExperiment(pidRunningInWorkspace)
+      }
+    }
+  }
+
   public hasRunningExperiment() {
     return this.experiments.hasRunningExperiment()
   }
@@ -562,12 +597,7 @@ export class Experiments extends BaseRepository<TableData> {
       () => this.getWebview(),
       () => this.notifyChanged(),
       () => this.selectColumns(),
-      (dvcRoot: string, ...ids: string[]) =>
-        this.internalCommands.executeCommand(
-          AvailableCommands.QUEUE_KILL,
-          dvcRoot,
-          ...ids
-        ),
+      (ids: string[]) => this.stopExperiments(ids),
       () =>
         this.internalCommands.executeCommand(
           AvailableCommands.STAGE_LIST,

extension/src/experiments/model/collect.ts

@@ -208,19 +208,23 @@ const collectExecutorInfo = (
   acc: ExperimentsAccumulator,
   experiment: Experiment,
   executor: ExecutorState
+  // eslint-disable-next-line sonarjs/cognitive-complexity
 ): void => {
   if (!executor) {
     return
   }
 
-  const { name, state } = executor
+  const { name, state, local } = executor
 
   if (name && state === ExperimentStatus.RUNNING) {
     experiment.executor = name
   }
   if (state && state !== ExperimentStatus.SUCCESS) {
     experiment.status = state
   }
+  if (local?.pid) {
+    experiment.executorPid = local.pid
+  }
 
   if (experiment.status === ExperimentStatus.RUNNING) {
     acc.runningExperiments.push({

extension/src/experiments/model/index.ts

@@ -334,9 +334,9 @@ export class ExperimentsModel extends ModelWithPersistence {
     )
   }
 
-  public getRunningQueueTasks() {
+  public getRunningExperiments() {
     return this.getExperimentsAndQueued().filter(experiment =>
-      isRunningInQueue(experiment)
+      isRunning(experiment.status)
     )
   }
 

extension/src/experiments/processExecution/index.ts

@@ -1,14 +1,7 @@
-import { join } from 'path'
-import { EXP_RWLOCK_FILE } from '../../cli/dvc/constants'
-import { getPidFromFile } from '../../fileSystem'
 import { Toast } from '../../vscode/toast'
 import { processExists, stopProcesses } from '../../process/execution'
 
-export const stopWorkspaceExperiment = async (dvcRoot: string) => {
-  const pid = await getPidFromFile(join(dvcRoot, EXP_RWLOCK_FILE))
-  if (!pid) {
-    return
-  }
+export const stopWorkspaceExperiment = async (pid: number) => {
   if (!(await processExists(pid))) {
     return
   }

extension/src/experiments/webview/contract.ts

@@ -33,6 +33,7 @@ export type Experiment = {
   description?: string
   error?: string
   executor?: Executor
+  executorPid?: number
   id: string
   label: string
   metrics?: MetricOrParamColumns

extension/src/experiments/webview/messages.ts

@@ -20,9 +20,6 @@ import { SortDefinition } from '../model/sortBy'
 import { getPositiveIntegerInput } from '../../vscode/inputBox'
 import { Title } from '../../vscode/title'
 import { ConfigKey, setConfigValue } from '../../vscode/config'
-import { Toast } from '../../vscode/toast'
-import { Executor, EXPERIMENT_WORKSPACE_ID } from '../../cli/dvc/contract'
-import { stopWorkspaceExperiment } from '../processExecution'
 import { hasDvcYamlFile } from '../../fileSystem'
 import { NUM_OF_COMMITS_TO_INCREASE } from '../../cli/dvc/constants'
 
@@ -35,10 +32,7 @@ export class WebviewMessages {
   private readonly getWebview: () => BaseWebview<TableData> | undefined
   private readonly notifyChanged: () => void
   private readonly selectColumns: () => Promise<void>
-  private readonly stopQueuedExperiments: (
-    dvcRoot: string,
-    ...ids: string[]
-  ) => Promise<string | undefined>
+  private readonly stopExperiments: (ids: string[]) => void
 
   private readonly hasStages: () => Promise<string | undefined>
 
@@ -62,10 +56,7 @@ export class WebviewMessages {
     getWebview: () => BaseWebview<TableData> | undefined,
     notifyChanged: () => void,
     selectColumns: () => Promise<void>,
-    stopQueuedExperiments: (
-      dvcRoot: string,
-      ...ids: string[]
-    ) => Promise<string | undefined>,
+    stopExperiments: (ids: string[]) => void,
     hasStages: () => Promise<string>,
     addStage: () => Promise<boolean>,
     selectBranches: (
@@ -80,7 +71,7 @@ export class WebviewMessages {
     this.getWebview = getWebview
     this.notifyChanged = notifyChanged
     this.selectColumns = selectColumns
-    this.stopQueuedExperiments = stopQueuedExperiments
+    this.stopExperiments = stopExperiments
     this.hasStages = hasStages
     this.addStage = addStage
     this.selectBranches = selectBranches
@@ -187,8 +178,8 @@ export class WebviewMessages {
         return this.setMaxTableHeadDepth()
       }
 
-      case MessageFromWebviewType.STOP_EXPERIMENT: {
-        return this.stopExperiments(message.payload)
+      case MessageFromWebviewType.STOP_EXPERIMENTS: {
+        return this.stopExperimentsAndSend(message.payload)
       }
 
       case MessageFromWebviewType.ADD_CONFIGURATION: {
@@ -443,37 +434,9 @@ export class WebviewMessages {
     )
   }
 
-  private stopExperiments(
-    runningExperiments: { id: string; executor?: string | null }[]
-  ) {
-    const { runningInQueueIds, runningInWorkspace } =
-      this.groupRunningExperiments(runningExperiments)
-
-    if (runningInQueueIds.size > 0) {
-      void Toast.showOutput(
-        this.stopQueuedExperiments(this.dvcRoot, ...runningInQueueIds)
-      )
-    }
-    if (runningInWorkspace) {
-      void stopWorkspaceExperiment(this.dvcRoot)
-    }
+  private stopExperimentsAndSend(ids: string[]) {
+    this.stopExperiments(ids)
 
     sendTelemetryEvent(EventName.EXPERIMENT_VIEW_STOP, undefined, undefined)
   }
-
-  private groupRunningExperiments(
-    experiments: { executor?: string | null; id: string }[]
-  ) {
-    let runningInWorkspace = false
-    const runningInQueueIds = new Set<string>()
-    for (const { executor, id } of experiments) {
-      if (executor === EXPERIMENT_WORKSPACE_ID) {
-        runningInWorkspace = true
-      }
-      if (executor === Executor.DVC_TASK) {
-        runningInQueueIds.add(id)
-      }
-    }
-    return { runningInQueueIds, runningInWorkspace }
-  }
 }

extension/src/experiments/workspace.ts

@@ -122,11 +122,20 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews<
     return this.getRepositoryThenUpdate('selectColumns', overrideRoot)
   }
 
-  public selectQueueTasksToKill() {
-    return this.pickIdsThenRun(
-      'pickQueueTasksToKill',
-      AvailableCommands.QUEUE_KILL
-    )
+  public async selectExperimentsToStop() {
+    const cwd = await this.getFocusedOrOnlyOrPickProject()
+    if (!cwd) {
+      return
+    }
+
+    const repository = this.getRepository(cwd)
+
+    const ids = await repository.pickRunningExperiments()
+
+    if (!ids || isEmpty(ids)) {
+      return
+    }
+    return repository.stopExperiments(ids)
   }
 
   public async selectExperimentsToPush(setup: Setup) {
@@ -145,11 +154,18 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews<
     return pushCommand({ dvcRoot, ids })
   }
 
-  public selectExperimentsToRemove() {
-    return this.pickIdsThenRun(
-      'pickExperimentsToRemove',
-      AvailableCommands.EXP_REMOVE
-    )
+  public async selectExperimentsToRemove() {
+    const cwd = await this.getFocusedOrOnlyOrPickProject()
+    if (!cwd) {
+      return
+    }
+
+    const ids = await this.getRepository(cwd).pickExperimentsToRemove()
+
+    if (!ids || isEmpty(ids)) {
+      return
+    }
+    return this.runCommand(AvailableCommands.EXP_REMOVE, cwd, ...ids)
   }
 
   public async modifyExperimentParamsAndRun(
@@ -511,25 +527,6 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews<
     return { command, enteredManually, trainingScript }
   }
 
-  private async pickIdsThenRun(
-    pickMethod: 'pickQueueTasksToKill' | 'pickExperimentsToRemove',
-    commandId:
-      | typeof AvailableCommands.QUEUE_KILL
-      | typeof AvailableCommands.EXP_REMOVE
-  ) {
-    const cwd = await this.getFocusedOrOnlyOrPickProject()
-    if (!cwd) {
-      return
-    }
-
-    const ids = await this.getRepository(cwd)[pickMethod]()
-
-    if (!ids || isEmpty(ids)) {
-      return
-    }
-    return this.runCommand(commandId, cwd, ...ids)
-  }
-
   private async pickExpThenRun(
     commandId: CommandId,
     pickFunc: (cwd: string) => Thenable<string | undefined> | undefined

extension/src/telemetry/constants.ts

@@ -148,7 +148,7 @@ export interface IEventNamePropertyMapping {
   [EventName.EXPERIMENT_VIEW_SHOW_LOGS]: undefined
   [EventName.EXPERIMENT_VIEW_STOP]: undefined
   [EventName.QUEUE_EXPERIMENT]: undefined
-  [EventName.QUEUE_KILL]: undefined
+  [EventName.EXPERIMENT_STOP]: undefined
   [EventName.QUEUE_START]: undefined
   [EventName.QUEUE_STOP]: undefined
 

extension/src/test/fixtures/expShow/base/output.ts

@@ -579,7 +579,7 @@ const data: ExpShowOutput = [
         ],
         executor: {
           name: Executor.WORKSPACE,
-          local: null,
+          local: { pid: 1234, root: null, log: null, returncode: null },
           state: ExperimentStatus.RUNNING
         }
       },

extension/src/test/fixtures/expShow/base/rows.ts

@@ -272,6 +272,7 @@ export const rowsFixtureWithBranches: Commit[] = [
         displayColor: undefined,
         description: '[exp-83425]',
         executor: Executor.WORKSPACE,
+        executorPid: 1234,
         id: 'exp-83425',
         label: EXPERIMENT_WORKSPACE_ID,
         metrics: {