diff --git a/extension/package.json b/extension/package.json index 79abbf7b1a..6b17f8c93d 100644 --- a/extension/package.json +++ b/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" }, diff --git a/extension/package.nls.json b/extension/package.nls.json index 8519723654..8921c9171d 100644 --- a/extension/package.nls.json +++ b/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", diff --git a/extension/src/cli/dvc/constants.ts b/extension/src/cli/dvc/constants.ts index 59095d4303..16872f4c28 100644 --- a/extension/src/cli/dvc/constants.ts +++ b/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', diff --git a/extension/src/cli/dvc/executor.ts b/extension/src/cli/dvc/executor.ts index 9fb9da9830..7c62153119 100644 --- a/extension/src/cli/dvc/executor.ts +++ b/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) { diff --git a/extension/src/commands/external.ts b/extension/src/commands/external.ts index 51e5ebe0ea..5e5c96fbee 100644 --- a/extension/src/commands/external.ts +++ b/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', diff --git a/extension/src/context.ts b/extension/src/context.ts index daed14f2bc..889e4bcdf3 100644 --- a/extension/src/context.ts +++ b/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)) + ) } } diff --git a/extension/src/experiments/index.ts b/extension/src/experiments/index.ts index 942e7e8ce6..604cccfad5 100644 --- a/extension/src/experiments/index.ts +++ b/extension/src/experiments/index.ts @@ -528,6 +528,10 @@ export class Experiments extends BaseRepository { return this.experiments.hasRunningExperiment() } + public hasRunningQueuedExperiment() { + return this.experiments.getRunningQueueTasks().length > 0 + } + public getFirstThreeColumnOrder() { return this.columns.getFirstThreeColumnOrder() } diff --git a/extension/src/experiments/workspace.ts b/extension/src/experiments/workspace.ts index a99e97d5c8..29c3d7cee0 100644 --- a/extension/src/experiments/workspace.ts +++ b/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: ( diff --git a/extension/src/extension.ts b/extension/src/extension.ts index 6d85b24891..761510f537 100644 --- a/extension/src/extension.ts +++ b/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( diff --git a/extension/src/telemetry/constants.ts b/extension/src/telemetry/constants.ts index 2c480407fd..ed6b947679 100644 --- a/extension/src/telemetry/constants.ts +++ b/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 diff --git a/extension/src/test/suite/context.test.ts b/extension/src/test/suite/context.test.ts index 27622f7cad..b035e9fc9f 100644 --- a/extension/src/test/suite/context.test.ts +++ b/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, diff --git a/extension/src/test/suite/extension.test.ts b/extension/src/test/suite/extension.test.ts index 59a5152b1c..48c0b1a784 100644 --- a/extension/src/test/suite/extension.test.ts +++ b/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) }) })