diff --git a/extension/package.json b/extension/package.json index 6d6db48594..9c469334ad 100644 --- a/extension/package.json +++ b/extension/package.json @@ -341,6 +341,11 @@ "category": "DVC", "icon": "$(run-all)" }, + { + "title": "%command.stopExperimentsQueue%", + "command": "dvc.stopExperimentsQueue", + "category": "DVC" + }, { "title": "%command.resetAndRunCheckpointExperiment%", "command": "dvc.resetAndRunCheckpointExperiment", @@ -775,7 +780,11 @@ }, { "command": "dvc.startExperimentsQueue", - "when": "dvc.commands.available && dvc.project.available && !dvc.experiment.running" + "when": "dvc.commands.available && dvc.project.available" + }, + { + "command": "dvc.stopExperimentsQueue", + "when": "dvc.commands.available && dvc.project.available" }, { "command": "dvc.selectForCompare", @@ -1257,9 +1266,14 @@ }, { "command": "dvc.startExperimentsQueue", - "when": "view == dvc.views.experimentsTree && !dvc.experiment.running", + "when": "view == dvc.views.experimentsTree", "group": "3_queue@2" }, + { + "command": "dvc.stopExperimentsQueue", + "when": "view == dvc.views.experimentsTree", + "group": "3_queue@3" + }, { "command": "dvc.addExperimentsTableSort", "when": "view == dvc.views.experimentsSortByTree", diff --git a/extension/package.nls.json b/extension/package.nls.json index 7a750a2a80..5398696baf 100644 --- a/extension/package.nls.json +++ b/extension/package.nls.json @@ -45,6 +45,7 @@ "command.runExperiment": "Run Experiment", "command.resumeCheckpointExperiment": "Resume Experiment", "command.startExperimentsQueue": "Start the Experiments Queue", + "command.stopExperimentsQueue": "Stop the Experiments Queue", "command.resetAndRunCheckpointExperiment": "Run Experiment", "command.selectForCompare": "Select for Compare", "command.selectFocusedProjects": "Select Project(s) to Focus (set dvc.focusedProjects)", diff --git a/extension/src/cli/dvc/constants.ts b/extension/src/cli/dvc/constants.ts index 355bec9fa4..193076c775 100644 --- a/extension/src/cli/dvc/constants.ts +++ b/extension/src/cli/dvc/constants.ts @@ -62,7 +62,8 @@ export enum ExperimentSubCommand { } export enum QueueSubCommand { - START = 'start' + START = 'start', + STOP = 'stop' } export enum ExperimentFlag { diff --git a/extension/src/cli/dvc/executor.test.ts b/extension/src/cli/dvc/executor.test.ts index 303f3f417d..0994bfc4c1 100644 --- a/extension/src/cli/dvc/executor.test.ts +++ b/extension/src/cli/dvc/executor.test.ts @@ -585,6 +585,27 @@ describe('CliExecutor', () => { }) }) + describe('queueStop', () => { + it("should call createProcess with the correct parameters to stop the experiment's queue", async () => { + const cwd = __dirname + + const stdout = 'Queue workers will stop after running tasks finish.' + + mockedCreateProcess.mockReturnValueOnce(getMockedProcess(stdout)) + + const output = await dvcExecutor.queueStop(cwd) + + expect(output).toStrictEqual(stdout) + + expect(mockedCreateProcess).toHaveBeenCalledWith({ + args: ['queue', 'stop'], + cwd, + env: mockedEnv, + executable: 'dvc' + }) + }) + }) + describe('remove', () => { it('should call createProcess with the correct parameters to remove a .dvc file', async () => { const cwd = __dirname diff --git a/extension/src/cli/dvc/executor.ts b/extension/src/cli/dvc/executor.ts index 4ba1c4b1df..cfec926f7d 100644 --- a/extension/src/cli/dvc/executor.ts +++ b/extension/src/cli/dvc/executor.ts @@ -27,6 +27,7 @@ export const autoRegisteredCommands = { PULL: 'pull', PUSH: 'push', QUEUE_START: 'queueStart', + QUEUE_STOP: 'queueStop', REMOVE: 'remove' } as const @@ -138,6 +139,10 @@ export class DvcExecutor extends DvcCli { ) } + public queueStop(cwd: string) { + return this.executeDvcProcess(cwd, Command.QUEUE, QueueSubCommand.STOP) + } + public remove(cwd: string, ...args: Args) { return this.blockAndExecuteProcess(cwd, Command.REMOVE, ...args) } diff --git a/extension/src/commands/external.ts b/extension/src/commands/external.ts index 4069bde666..f5142a6593 100644 --- a/extension/src/commands/external.ts +++ b/extension/src/commands/external.ts @@ -12,6 +12,7 @@ export enum RegisteredCliCommands { EXPERIMENT_SHARE_AS_COMMIT = 'dvc.shareExperimentAsCommit', QUEUE_EXPERIMENT = 'dvc.queueExperiment', QUEUE_START = 'dvc.startExperimentsQueue', + QUEUE_STOP = 'dvc.stopExperimentsQueue', EXPERIMENT_VIEW_APPLY = 'dvc.views.experiments.applyExperiment', EXPERIMENT_VIEW_BRANCH = 'dvc.views.experiments.branchExperiment', diff --git a/extension/src/experiments/commands/register.ts b/extension/src/experiments/commands/register.ts index ea53589583..70114e57db 100644 --- a/extension/src/experiments/commands/register.ts +++ b/extension/src/experiments/commands/register.ts @@ -27,6 +27,11 @@ const registerExperimentCwdCommands = ( ) ) + internalCommands.registerExternalCliCommand( + RegisteredCliCommands.QUEUE_STOP, + () => experiments.getCwdThenReport(AvailableCommands.QUEUE_STOP) + ) + internalCommands.registerExternalCliCommand( RegisteredCliCommands.MODIFY_EXPERIMENT_PARAMS_AND_QUEUE, () => diff --git a/extension/src/telemetry/constants.ts b/extension/src/telemetry/constants.ts index 94ac61b3fb..7da09f3b25 100644 --- a/extension/src/telemetry/constants.ts +++ b/extension/src/telemetry/constants.ts @@ -155,6 +155,7 @@ export interface IEventNamePropertyMapping { [EventName.EXPERIMENT_VIEW_SHARE_AS_COMMIT]: undefined [EventName.QUEUE_EXPERIMENT]: undefined [EventName.QUEUE_START]: undefined + [EventName.QUEUE_STOP]: undefined [EventName.EXPERIMENT_VIEW_QUEUE]: undefined [EventName.EXPERIMENT_VIEW_RESUME]: undefined diff --git a/extension/src/test/suite/experiments/workspace.test.ts b/extension/src/test/suite/experiments/workspace.test.ts index e6f6118a23..25ed3554ed 100644 --- a/extension/src/test/suite/experiments/workspace.test.ts +++ b/extension/src/test/suite/experiments/workspace.test.ts @@ -438,7 +438,7 @@ suite('Workspace Experiments Test Suite', () => { }) describe('dvc.startExperimentsQueue', () => { - it('should be able to execute all experiments in the run queue', async () => { + it('should be able to start the experiments queue with the selected number of workers', async () => { const mockQueueStart = stub(DvcExecutor.prototype, 'queueStart').resolves( undefined ) @@ -454,11 +454,29 @@ suite('Workspace Experiments Test Suite', () => { await commands.executeCommand(RegisteredCliCommands.QUEUE_START) expect(mockQueueStart).to.be.calledOnce - expect(mockQueueStart).to.be.calledWith(dvcDemoPath, dDosNumberOfJobs) + expect(mockQueueStart).to.be.calledWithExactly( + dvcDemoPath, + dDosNumberOfJobs + ) expect(mockInputBox) }) }) + describe('dvc.stopExperimentsQueue', () => { + it('should be able to stop the experiments queue', async () => { + const mockQueueStop = stub(DvcExecutor.prototype, 'queueStop').resolves( + undefined + ) + + stubWorkspaceExperimentsGetters(dvcDemoPath) + + await commands.executeCommand(RegisteredCliCommands.QUEUE_STOP) + + expect(mockQueueStop).to.be.calledOnce + expect(mockQueueStop).to.be.calledWithExactly(dvcDemoPath) + }) + }) + describe('dvc.applyExperiment', () => { it('should ask the user to pick an experiment and then apply that experiment to the workspace', async () => { const selectedExperiment = 'test-branch'