diff --git a/extension/src/commands/external.ts b/extension/src/commands/external.ts index d3de70e56b..4bd973f718 100644 --- a/extension/src/commands/external.ts +++ b/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', diff --git a/extension/src/experiments/commands/register.ts b/extension/src/experiments/commands/register.ts index 25efce68ff..da8f52cc0d 100644 --- a/extension/src/experiments/commands/register.ts +++ b/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( diff --git a/extension/src/experiments/index.ts b/extension/src/experiments/index.ts index 6d5c943383..f428a9c28d 100644 --- a/extension/src/experiments/index.ts +++ b/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 { ) } - 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 { 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() + 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 { () => 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, diff --git a/extension/src/experiments/model/collect.ts b/extension/src/experiments/model/collect.ts index 16155ed90b..2b3ff03a36 100644 --- a/extension/src/experiments/model/collect.ts +++ b/extension/src/experiments/model/collect.ts @@ -208,12 +208,13 @@ 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 @@ -221,6 +222,9 @@ const collectExecutorInfo = ( if (state && state !== ExperimentStatus.SUCCESS) { experiment.status = state } + if (local?.pid) { + experiment.executorPid = local.pid + } if (experiment.status === ExperimentStatus.RUNNING) { acc.runningExperiments.push({ diff --git a/extension/src/experiments/model/index.ts b/extension/src/experiments/model/index.ts index 26a40e35ba..a203e0843f 100644 --- a/extension/src/experiments/model/index.ts +++ b/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) ) } diff --git a/extension/src/experiments/processExecution/index.ts b/extension/src/experiments/processExecution/index.ts index 702f974fca..fa14c73f0d 100644 --- a/extension/src/experiments/processExecution/index.ts +++ b/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 } diff --git a/extension/src/experiments/webview/contract.ts b/extension/src/experiments/webview/contract.ts index 957a80c6b4..df2c104720 100644 --- a/extension/src/experiments/webview/contract.ts +++ b/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 diff --git a/extension/src/experiments/webview/messages.ts b/extension/src/experiments/webview/messages.ts index e3403f23d4..4e241a75b4 100644 --- a/extension/src/experiments/webview/messages.ts +++ b/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 | undefined private readonly notifyChanged: () => void private readonly selectColumns: () => Promise - private readonly stopQueuedExperiments: ( - dvcRoot: string, - ...ids: string[] - ) => Promise + private readonly stopExperiments: (ids: string[]) => void private readonly hasStages: () => Promise @@ -62,10 +56,7 @@ export class WebviewMessages { getWebview: () => BaseWebview | undefined, notifyChanged: () => void, selectColumns: () => Promise, - stopQueuedExperiments: ( - dvcRoot: string, - ...ids: string[] - ) => Promise, + stopExperiments: (ids: string[]) => void, hasStages: () => Promise, addStage: () => Promise, 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() - for (const { executor, id } of experiments) { - if (executor === EXPERIMENT_WORKSPACE_ID) { - runningInWorkspace = true - } - if (executor === Executor.DVC_TASK) { - runningInQueueIds.add(id) - } - } - return { runningInQueueIds, runningInWorkspace } - } } diff --git a/extension/src/experiments/workspace.ts b/extension/src/experiments/workspace.ts index 6d4758e225..2367d3e69f 100644 --- a/extension/src/experiments/workspace.ts +++ b/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 | undefined diff --git a/extension/src/telemetry/constants.ts b/extension/src/telemetry/constants.ts index a22a2302ac..4d6381757c 100644 --- a/extension/src/telemetry/constants.ts +++ b/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 diff --git a/extension/src/test/fixtures/expShow/base/output.ts b/extension/src/test/fixtures/expShow/base/output.ts index a4761b41e4..eaca7564eb 100644 --- a/extension/src/test/fixtures/expShow/base/output.ts +++ b/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 } }, diff --git a/extension/src/test/fixtures/expShow/base/rows.ts b/extension/src/test/fixtures/expShow/base/rows.ts index 5f7ca0b623..6e150cced2 100644 --- a/extension/src/test/fixtures/expShow/base/rows.ts +++ b/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: { diff --git a/extension/src/test/suite/experiments/index.test.ts b/extension/src/test/suite/experiments/index.test.ts index f0fa9503a8..6b326ba5ae 100644 --- a/extension/src/test/suite/experiments/index.test.ts +++ b/extension/src/test/suite/experiments/index.test.ts @@ -63,7 +63,7 @@ import * as Telemetry from '../../../telemetry' import { EventName } from '../../../telemetry/constants' import * as VscodeContext from '../../../vscode/context' import { Title } from '../../../vscode/title' -import { EXP_RWLOCK_FILE, ExperimentFlag } from '../../../cli/dvc/constants' +import { ExperimentFlag } from '../../../cli/dvc/constants' import { DvcExecutor } from '../../../cli/dvc/executor' import { WorkspacePlots } from '../../../plots/workspace' import { @@ -1245,32 +1245,23 @@ suite('Experiments Test Suite', () => { const webview = await experiments.showWebview() const mockMessageReceived = getMessageReceivedEmitter(webview) - const mockExperimentIds = ['exp-e7a67', 'test-branch'] + const mockExperimentIds = ['exp-e7a67', 'exp-83425'] stubWorkspaceExperimentsGetters(dvcDemoPath, experiments) const mockPid = 1234 - const mockGetPidFromFile = stub(FileSystem, 'getPidFromFile').resolves( - mockPid - ) const mockProcessExists = stub( ProcessExecution, 'processExists' ).resolves(true) mockMessageReceived.fire({ - payload: [ - ...mockExperimentIds.map(id => ({ executor: Executor.DVC_TASK, id })), - { executor: Executor.WORKSPACE, id: EXPERIMENT_WORKSPACE_ID } - ], - type: MessageFromWebviewType.STOP_EXPERIMENT + payload: mockExperimentIds, + type: MessageFromWebviewType.STOP_EXPERIMENTS }) await Promise.all([experimentsKilled, workspaceStopped]) - expect(mockQueueKill).to.be.calledWith(dvcDemoPath, ...mockExperimentIds) - expect(mockGetPidFromFile).to.be.calledWithExactly( - join(dvcDemoPath, EXP_RWLOCK_FILE) - ) + expect(mockQueueKill).to.be.calledWith(dvcDemoPath, 'exp-e7a67') expect(mockProcessExists).to.be.calledWithExactly(mockPid) expect(mockStopProcesses).to.be.calledWithExactly([mockPid]) }).timeout(WEBVIEW_TEST_TIMEOUT) diff --git a/extension/src/test/suite/experiments/workspace.test.ts b/extension/src/test/suite/experiments/workspace.test.ts index df51617bd5..387cc15204 100644 --- a/extension/src/test/suite/experiments/workspace.test.ts +++ b/extension/src/test/suite/experiments/workspace.test.ts @@ -447,8 +447,8 @@ suite('Workspace Experiments Test Suite', () => { }) }) - describe('dvc.stopQueuedExperiments', () => { - it('should be able to kill running queue tasks', async () => { + describe('dvc.stopExperiments', () => { + it('should be able to stop any running experiment', async () => { const mockQueueKill = stub(DvcExecutor.prototype, 'queueKill').resolves( undefined ) @@ -472,7 +472,7 @@ suite('Workspace Experiments Test Suite', () => { stubWorkspaceExperimentsGetters(dvcDemoPath, experiments) - await commands.executeCommand(RegisteredCliCommands.QUEUE_KILL) + await commands.executeCommand(RegisteredCommands.EXPERIMENT_STOP) expect(mockShowQuickPick).to.be.calledWithExactly( [ @@ -483,13 +483,21 @@ suite('Workspace Experiments Test Suite', () => { )}, loss:2.0205045, accuracy:0.37241668`, label: '4fb124a', value: queueTaskId + }, + { + description: '[exp-83425]', + detail: `Created:${formatDate( + '2020-12-29T15:27:02' + )}, loss:1.7750162, accuracy:0.59265000`, + label: 'workspace', + value: 'exp-83425' } ], { canPickMany: true, matchOnDescription: true, matchOnDetail: true, - title: Title.SELECT_QUEUE_KILL + title: Title.SELECT_EXPERIMENTS_STOP } ) expect(mockQueueKill).to.be.calledOnce diff --git a/extension/src/vscode/title.ts b/extension/src/vscode/title.ts index 51398ffb21..f5eb35800a 100644 --- a/extension/src/vscode/title.ts +++ b/extension/src/vscode/title.ts @@ -31,7 +31,7 @@ export enum Title { SELECT_CUSTOM_PLOTS_TO_REMOVE = 'Select Custom Plot(s) to Remove', SELECT_PARAM_TO_MODIFY = 'Select Param(s) to Modify', SELECT_PLOTS = 'Select Plots to Display', - SELECT_QUEUE_KILL = 'Select Queue Task(s) to Kill', + SELECT_EXPERIMENTS_STOP = 'Select Experiments to Stop', SELECT_SORT_DIRECTION = 'Select Sort Direction', SELECT_SORTS_TO_REMOVE = 'Select Sort(s) to Remove', SELECT_TRAINING_SCRIPT = 'Select your training script', diff --git a/extension/src/webview/contract.ts b/extension/src/webview/contract.ts index c02972e181..e39233d047 100644 --- a/extension/src/webview/contract.ts +++ b/extension/src/webview/contract.ts @@ -38,7 +38,7 @@ export enum MessageFromWebviewType { RESIZE_PLOTS = 'resize-plots', SAVE_STUDIO_TOKEN = 'save-studio-token', SHOW_EXPERIMENT_LOGS = 'show-experiment-logs', - STOP_EXPERIMENT = 'stop-experiment', + STOP_EXPERIMENTS = 'stop-experiments', SORT_COLUMN = 'sort-column', TOGGLE_EXPERIMENT = 'toggle-experiment', TOGGLE_EXPERIMENT_STAR = 'toggle-experiment-star', @@ -147,8 +147,8 @@ export type MessageFromWebview = payload: SortDefinition } | { - type: MessageFromWebviewType.STOP_EXPERIMENT - payload: { id: string; executor?: string | null }[] + type: MessageFromWebviewType.STOP_EXPERIMENTS + payload: string[] } | { type: MessageFromWebviewType.SHOW_EXPERIMENT_LOGS; payload: string } | { diff --git a/webview/src/experiments/components/App.test.tsx b/webview/src/experiments/components/App.test.tsx index 8578763c9c..099d7314a1 100644 --- a/webview/src/experiments/components/App.test.tsx +++ b/webview/src/experiments/components/App.test.tsx @@ -18,7 +18,7 @@ import { } from 'dvc/src/experiments/webview/contract' import { buildMetricOrParamPath } from 'dvc/src/experiments/columns/paths' import dataTypesTableFixture from 'dvc/src/test/fixtures/expShow/dataTypes/tableData' -import { EXPERIMENT_WORKSPACE_ID, Executor } from 'dvc/src/cli/dvc/contract' +import { EXPERIMENT_WORKSPACE_ID } from 'dvc/src/cli/dvc/contract' import { useIsFullyContained } from './overflowHoverTooltip/useIsFullyContained' import styles from './table/styles.module.scss' import { vsCodeApi } from '../../shared/api' @@ -1144,8 +1144,8 @@ describe('App', () => { stopOption && fireEvent.click(stopOption) expect(sendMessage).toHaveBeenCalledWith({ - payload: [{ executor: Executor.DVC_TASK, id: 'exp-e7a67' }], - type: MessageFromWebviewType.STOP_EXPERIMENT + payload: ['exp-e7a67'], + type: MessageFromWebviewType.STOP_EXPERIMENTS }) }) @@ -1175,7 +1175,7 @@ describe('App', () => { expect(sendMessage).toHaveBeenCalledWith({ payload: ['exp-e7a67', 'exp-83425'], - type: MessageFromWebviewType.STOP_EXPERIMENT + type: MessageFromWebviewType.STOP_EXPERIMENTS }) }) @@ -1218,10 +1218,8 @@ describe('App', () => { stopOption && fireEvent.click(stopOption) expect(sendMessage).toHaveBeenCalledWith({ - payload: [ - { executor: Executor.WORKSPACE, id: EXPERIMENT_WORKSPACE_ID } - ], - type: MessageFromWebviewType.STOP_EXPERIMENT + payload: [EXPERIMENT_WORKSPACE_ID], + type: MessageFromWebviewType.STOP_EXPERIMENTS }) }) diff --git a/webview/src/experiments/components/table/body/RowContextMenu.tsx b/webview/src/experiments/components/table/body/RowContextMenu.tsx index ae7e4cd751..c050987abd 100644 --- a/webview/src/experiments/components/table/body/RowContextMenu.tsx +++ b/webview/src/experiments/components/table/body/RowContextMenu.tsx @@ -135,7 +135,7 @@ const getMultiSelectMenuOptions = ( experimentMenuOption( selectedIds, 'Stop', - MessageFromWebviewType.STOP_EXPERIMENT, + MessageFromWebviewType.STOP_EXPERIMENTS, disableStopOption, true ), @@ -270,9 +270,9 @@ const getSingleSelectMenuOptions = ( !hasRunningWorkspaceExperiment ), experimentMenuOption( - [{ executor, id }], + [id], 'Stop', - MessageFromWebviewType.STOP_EXPERIMENT, + MessageFromWebviewType.STOP_EXPERIMENTS, !isRunning(status), id !== EXPERIMENT_WORKSPACE_ID ),