Skip to content

Commit 9198d9c

Browse files
mattseddonmattseddon
authored
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
1 parent 9da00b1 commit 9198d9c

File tree

22 files changed

+321
-193
lines changed

22 files changed

+321
-193
lines changed

‎extension/package.json‎

Lines changed: 33 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -362,8 +362,8 @@
362362
"category": "DVC"
363363
},
364364
{
365-
"title": "Stop Running Queued Experiment(s)",
366-
"command": "dvc.stopQueuedExperiments",
365+
"title": "Stop Running Experiment(s)",
366+
"command": "dvc.stopExperiments",
367367
"category": "DVC"
368368
},
369369
{
@@ -429,6 +429,18 @@
429429
"command": "dvc.showSetup",
430430
"category": "DVC"
431431
},
432+
{
433+
"title": "Stop",
434+
"command": "dvc.views.experimentsTree.stopExperiment",
435+
"category": "DVC",
436+
"icon": "$(debug-stop)"
437+
},
438+
{
439+
"title": "Stop",
440+
"command": "dvc.views.experiments.stopExperiment",
441+
"category": "DVC",
442+
"icon": "$(debug-stop)"
443+
},
432444
{
433445
"title": "Stop All Running Experiments",
434446
"command": "dvc.stopAllRunningExperiments",
@@ -827,7 +839,7 @@
827839
"when": "dvc.commands.available && dvc.project.available"
828840
},
829841
{
830-
"command": "dvc.stopQueuedExperiments",
842+
"command": "dvc.stopExperiments",
831843
"when": "dvc.commands.available && dvc.project.available"
832844
},
833845
{
@@ -914,6 +926,14 @@
914926
"command": "dvc.views.experiments.showLogs",
915927
"when": "false"
916928
},
929+
{
930+
"command": "dvc.views.experimentsTree.stopExperiment",
931+
"when": "false"
932+
},
933+
{
934+
"command": "dvc.views.experiments.stopExperiment",
935+
"when": "false"
936+
},
917937
{
918938
"command": "dvc.views.experimentsFilterByTree.removeAllFilters",
919939
"when": "false"
@@ -1160,6 +1180,11 @@
11601180
"group": "inline",
11611181
"when": "view == dvc.views.experimentsFilterByTree && dvc.commands.available && viewItem != dvcRoot"
11621182
},
1183+
{
1184+
"command": "dvc.views.experimentsTree.stopExperiment",
1185+
"group": "inline@0",
1186+
"when": "view == dvc.views.experimentsTree && dvc.commands.available && viewItem == running"
1187+
},
11631188
{
11641189
"command": "dvc.views.experiments.applyExperiment",
11651190
"group": "inline@1",
@@ -1217,6 +1242,11 @@
12171242
"when": "view == dvc.views.experimentsColumnsTree",
12181243
"group": "navigation@2"
12191244
},
1245+
{
1246+
"command": "dvc.stopAllRunningExperiments",
1247+
"when": "view == dvc.views.experimentsTree && dvc.experiments.webview.active && dvc.experiment.running",
1248+
"group": "navigation@0"
1249+
},
12201250
{
12211251
"command": "dvc.runExperiment",
12221252
"when": "view == dvc.views.experimentsTree && !dvc.experiment.running.workspace && !dvc.experiment.checkpoints",
@@ -1237,11 +1267,6 @@
12371267
"when": "view == dvc.views.experimentsTree && !dvc.experiments.webview.active && dvc.experiment.running",
12381268
"group": "1_run@1"
12391269
},
1240-
{
1241-
"command": "dvc.stopAllRunningExperiments",
1242-
"when": "view == dvc.views.experimentsTree && dvc.experiments.webview.active && dvc.experiment.running",
1243-
"group": "navigation@1"
1244-
},
12451270
{
12461271
"command": "dvc.views.experimentsTree.selectExperiments",
12471272
"when": "view == dvc.views.experimentsTree && dvc.plots.webview.active",
@@ -1287,11 +1312,6 @@
12871312
"when": "view == dvc.views.experimentsTree",
12881313
"group": "3_queue@3"
12891314
},
1290-
{
1291-
"command": "dvc.stopQueuedExperiments",
1292-
"when": "view == dvc.views.experimentsTree",
1293-
"group": "3_queue@5"
1294-
},
12951315
{
12961316
"command": "dvc.addExperimentsTableSort",
12971317
"when": "view == dvc.views.experimentsSortByTree",

‎extension/src/commands/external.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,6 @@ export enum RegisteredCliCommands {
99
EXPERIMENT_RESUME = 'dvc.resumeCheckpointExperiment',
1010
EXPERIMENT_RUN = 'dvc.runExperiment',
1111
QUEUE_EXPERIMENT = 'dvc.queueExperiment',
12-
QUEUE_KILL = 'dvc.stopQueuedExperiments',
1312
QUEUE_START = 'dvc.startExperimentsQueue',
1413
QUEUE_STOP = 'dvc.stopExperimentsQueue',
1514

@@ -18,7 +17,6 @@ export enum RegisteredCliCommands {
1817
EXPERIMENT_VIEW_PUSH = 'dvc.views.experiments.pushExperiment',
1918
EXPERIMENT_VIEW_REMOVE = 'dvc.views.experiments.removeExperiment',
2019
EXPERIMENT_VIEW_SHOW_LOGS = 'dvc.views.experiments.showLogs',
21-
EXPERIMENT_VIEW_STOP = 'dvc.views.experiments.stopQueueExperiment',
2220

2321
EXPERIMENT_VIEW_QUEUE = 'dvc.views.experiments.queueExperiment',
2422
EXPERIMENT_VIEW_RESUME = 'dvc.views.experiments.resumeCheckpointExperiment',
@@ -62,7 +60,9 @@ export enum RegisteredCommands {
6260
EXPERIMENT_SORT_REMOVE = 'dvc.views.experimentsSortByTree.removeSort',
6361
EXPERIMENT_SORTS_REMOVE = 'dvc.removeExperimentsTableSorts',
6462
EXPERIMENT_SORTS_REMOVE_ALL = 'dvc.views.experimentsSortByTree.removeAllSorts',
63+
EXPERIMENT_STOP = 'dvc.stopExperiments',
6564
EXPERIMENT_TOGGLE = 'dvc.views.experiments.toggleStatus',
65+
EXPERIMENT_VIEW_STOP = 'dvc.views.experiments.stopExperiment',
6666
STOP_EXPERIMENTS = 'dvc.stopAllRunningExperiments',
6767

6868
PLOTS_PATH_TOGGLE = 'dvc.views.plotsPathsTree.toggleStatus',

‎extension/src/experiments/commands/register.ts‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,12 @@ const registerExperimentNameCommands = (
115115
({ dvcRoot, ids }: { dvcRoot: string; ids: string[] }) =>
116116
experiments.runCommand(AvailableCommands.EXP_REMOVE, dvcRoot, ...ids)
117117
)
118+
119+
internalCommands.registerExternalCommand(
120+
RegisteredCommands.EXPERIMENT_VIEW_STOP,
121+
({ dvcRoot, ids }: { dvcRoot: string; ids: string[] }) =>
122+
experiments.stopExperiments(dvcRoot, ...ids)
123+
)
118124
}
119125

120126
const registerExperimentInputCommands = (
@@ -200,9 +206,9 @@ const registerExperimentQuickPickCommands = (
200206
experiments.selectColumns(getDvcRootFromContext(context))
201207
)
202208

203-
internalCommands.registerExternalCliCommand(
204-
RegisteredCliCommands.QUEUE_KILL,
205-
() => experiments.selectQueueTasksToKill()
209+
internalCommands.registerExternalCommand(
210+
RegisteredCommands.EXPERIMENT_STOP,
211+
() => experiments.selectExperimentsToStop()
206212
)
207213

208214
internalCommands.registerExternalCliCommand(

‎extension/src/experiments/index.ts‎

Lines changed: 28 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ import { starredSort } from './model/sortBy/constants'
3030
import { pickSortsToRemove, pickSortToAdd } from './model/sortBy/quickPick'
3131
import { ColumnsModel } from './columns/model'
3232
import { ExperimentsData } from './data'
33+
import { stopWorkspaceExperiment } from './processExecution'
3334
import {
3435
Experiment,
3536
ColumnType,
@@ -378,11 +379,11 @@ export class Experiments extends BaseRepository<TableData> {
378379
)
379380
}
380381

381-
public pickQueueTasksToKill() {
382+
public pickRunningExperiments() {
382383
return pickExperiments(
383-
this.experiments.getRunningQueueTasks(),
384+
this.experiments.getRunningExperiments(),
384385
this.getFirstThreeColumnOrder(),
385-
Title.SELECT_QUEUE_KILL
386+
Title.SELECT_EXPERIMENTS_STOP
386387
)
387388
}
388389

@@ -486,6 +487,30 @@ export class Experiments extends BaseRepository<TableData> {
486487
return this.notifyChanged()
487488
}
488489

490+
public stopExperiments(ids: string[]) {
491+
const { runningInQueueIds, runningInWorkspaceId } =
492+
this.experiments.getStopDetails(ids)
493+
494+
const promises: Promise<string | void>[] = []
495+
496+
if (runningInQueueIds) {
497+
promises.push(
498+
this.internalCommands.executeCommand(
499+
AvailableCommands.QUEUE_KILL,
500+
this.dvcRoot,
501+
...runningInQueueIds
502+
)
503+
)
504+
}
505+
if (runningInWorkspaceId) {
506+
promises.push(stopWorkspaceExperiment(this.dvcRoot, runningInWorkspaceId))
507+
}
508+
509+
return Toast.showOutput(
510+
Promise.all(promises).then(output => output.filter(Boolean).join('\n'))
511+
)
512+
}
513+
489514
public hasRunningExperiment() {
490515
return this.experiments.hasRunningExperiment()
491516
}
@@ -562,12 +587,6 @@ export class Experiments extends BaseRepository<TableData> {
562587
() => this.getWebview(),
563588
() => this.notifyChanged(),
564589
() => this.selectColumns(),
565-
(dvcRoot: string, ...ids: string[]) =>
566-
this.internalCommands.executeCommand(
567-
AvailableCommands.QUEUE_KILL,
568-
dvcRoot,
569-
...ids
570-
),
571590
() =>
572591
this.internalCommands.executeCommand(
573592
AvailableCommands.STAGE_LIST,

‎extension/src/experiments/model/collect.ts‎

Lines changed: 45 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,8 @@ import {
1010
Experiment,
1111
CommitData,
1212
RunningExperiment,
13-
isQueued
13+
isQueued,
14+
isRunning
1415
} from '../webview/contract'
1516
import {
1617
EXPERIMENT_WORKSPACE_ID,
@@ -205,7 +206,6 @@ const getExecutor = (experiment: Experiment): Executor => {
205206
}
206207

207208
const collectExecutorInfo = (
208-
acc: ExperimentsAccumulator,
209209
experiment: Experiment,
210210
executor: ExecutorState
211211
): void => {
@@ -221,13 +221,19 @@ const collectExecutorInfo = (
221221
if (state && state !== ExperimentStatus.SUCCESS) {
222222
experiment.status = state
223223
}
224+
}
224225

225-
if (experiment.status === ExperimentStatus.RUNNING) {
226-
acc.runningExperiments.push({
227-
executor: getExecutor(experiment),
228-
id: experiment.id
229-
})
226+
const collectRunningExperiment = (
227+
acc: ExperimentsAccumulator,
228+
experiment: Experiment
229+
): void => {
230+
if (!isRunning(experiment.status)) {
231+
return
230232
}
233+
acc.runningExperiments.push({
234+
executor: getExecutor(experiment),
235+
id: experiment.id
236+
})
231237
}
232238

233239
const collectExpRange = (
@@ -261,7 +267,8 @@ const collectExpRange = (
261267
experiment.description = `[${name}]`
262268
}
263269

264-
collectExecutorInfo(acc, experiment, executor)
270+
collectExecutorInfo(experiment, executor)
271+
collectRunningExperiment(acc, experiment)
265272

266273
addToMapArray(acc.experimentsByCommit, baseline.id, experiment)
267274
}
@@ -402,3 +409,33 @@ export const collectOrderedCommitsAndExperiments = (
402409
}
403410
return acc
404411
}
412+
413+
export const collectRunningInQueue = (
414+
ids: Set<string>,
415+
runningExperiments: RunningExperiment[]
416+
): string[] | undefined => {
417+
const runningInQueueIds = new Set<string>()
418+
for (const { executor, id } of runningExperiments) {
419+
if (!ids.has(id)) {
420+
continue
421+
}
422+
if (executor === Executor.DVC_TASK) {
423+
runningInQueueIds.add(id)
424+
}
425+
}
426+
return runningInQueueIds.size > 0 ? [...runningInQueueIds] : undefined
427+
}
428+
429+
export const collectRunningInWorkspace = (
430+
ids: Set<string>,
431+
runningExperiments: RunningExperiment[]
432+
): string | undefined => {
433+
for (const { executor, id } of runningExperiments) {
434+
if (!ids.has(id)) {
435+
continue
436+
}
437+
if (executor === EXPERIMENT_WORKSPACE_ID) {
438+
return id
439+
}
440+
}
441+
}

‎extension/src/experiments/model/index.ts‎

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,9 @@ import { SortDefinition, sortExperiments } from './sortBy'
33
import { FilterDefinition, filterExperiment, getFilterId } from './filterBy'
44
import {
55
collectExperiments,
6-
collectOrderedCommitsAndExperiments
6+
collectOrderedCommitsAndExperiments,
7+
collectRunningInQueue,
8+
collectRunningInWorkspace
79
} from './collect'
810
import {
911
collectColoredStatus,
@@ -284,7 +286,11 @@ export class ExperimentsModel extends ModelWithPersistence {
284286
{
285287
...this.addDetails(this.workspace),
286288
hasChildren: false,
287-
type: ExperimentType.WORKSPACE
289+
type: this.running.some(
290+
({ executor }) => executor === Executor.WORKSPACE
291+
)
292+
? ExperimentType.RUNNING
293+
: ExperimentType.WORKSPACE
288294
},
289295
...this.commits.map(commit => {
290296
return {
@@ -334,12 +340,21 @@ export class ExperimentsModel extends ModelWithPersistence {
334340
)
335341
}
336342

337-
public getRunningQueueTasks() {
343+
public getRunningExperiments() {
338344
return this.getExperimentsAndQueued().filter(experiment =>
339-
isRunningInQueue(experiment)
345+
isRunning(experiment.status)
340346
)
341347
}
342348

349+
public getStopDetails(idsToStop: string[]) {
350+
const running = [...this.running]
351+
const ids = new Set(idsToStop)
352+
return {
353+
runningInQueueIds: collectRunningInQueue(ids, running),
354+
runningInWorkspaceId: collectRunningInWorkspace(ids, running)
355+
}
356+
}
357+
343358
public getRowData() {
344359
return [
345360
this.addDetails(this.workspace),

0 commit comments

Comments
 (0)