Prevent plotting of running experiments

extension/src/experiments/commands/register.ts

@@ -257,7 +257,7 @@ const registerExperimentQuickPickCommands = (
   internalCommands.registerExternalCommand(
     RegisteredCommands.EXPERIMENT_SELECT,
     (context: Context) =>
-      experiments.selectExperiments(getDvcRootFromContext(context))
+      experiments.selectExperimentsToPlot(getDvcRootFromContext(context))
   )
 
   internalCommands.registerExternalCommand(

extension/src/experiments/index.ts

@@ -30,13 +30,18 @@ import { starredSort } from './model/sortBy/constants'
 import { pickSortsToRemove, pickSortToAdd } from './model/sortBy/quickPick'
 import { ColumnsModel } from './columns/model'
 import { ExperimentsData } from './data'
-import { Experiment, ColumnType, TableData } from './webview/contract'
+import {
+  Experiment,
+  ColumnType,
+  TableData,
+  isRunning
+} from './webview/contract'
 import { WebviewMessages } from './webview/messages'
 import { DecorationProvider } from './model/decorationProvider'
 import { starredFilter } from './model/filterBy/constants'
 import { ResourceLocator } from '../resourceLocator'
 import { AvailableCommands, InternalCommands } from '../commands/internal'
-import { EXPERIMENT_WORKSPACE_ID, ExpShowOutput } from '../cli/dvc/contract'
+import { ExpShowOutput } from '../cli/dvc/contract'
 import { ViewKey } from '../webview/constants'
 import { BaseRepository } from '../webview/repository'
 import { Title } from '../vscode/title'
@@ -217,15 +222,6 @@ export class Experiments extends BaseRepository<TableData> {
   public toggleExperimentStatus(
     id: string
   ): Color | typeof UNSELECTED | undefined {
-    if (
-      this.experiments.isRunningInWorkspace(id) &&
-      !this.experiments.isSelected(id)
-    ) {
-      return this.experiments.isSelected(EXPERIMENT_WORKSPACE_ID)
-        ? undefined
-        : this.toggleExperimentStatus(EXPERIMENT_WORKSPACE_ID)
-    }
-
     const selected = this.experiments.isSelected(id)
     if (!selected && !this.experiments.canSelect()) {
       return
@@ -353,8 +349,10 @@ export class Experiments extends BaseRepository<TableData> {
     return this.experiments.getWorkspaceCommitsAndExperiments()
   }
 
-  public async selectExperiments() {
-    const experiments = this.experiments.getWorkspaceCommitsAndExperiments()
+  public async selectExperimentsToPlot() {
+    const experiments = this.experiments
+      .getWorkspaceCommitsAndExperiments()
+      .filter(({ status }) => !isRunning(status))
 
     const selected = await pickExperimentsToPlot(
       experiments,

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

@@ -18,7 +18,8 @@ import survivalRowsFixture from '../../test/fixtures/expShow/survival/rows'
 
 import {
   ExperimentStatus,
-  EXPERIMENT_WORKSPACE_ID
+  EXPERIMENT_WORKSPACE_ID,
+  Executor
 } from '../../cli/dvc/contract'
 import { PersistenceKey } from '../../persistence/constants'
 
@@ -220,6 +221,68 @@ describe('ExperimentsModel', () => {
     expect(experimentsModel.getSelectedRevisions()).toHaveLength(7)
   })
 
+  it('should swap an experiment running in the workspace for the workspace and not select an experiment running in the queue', () => {
+    const params = {
+      params: {
+        'params.yaml': {
+          data: {
+            lr: 0.000000000000001
+          }
+        }
+      }
+    }
+    const data = generateTestExpShowOutput(
+      {},
+      {
+        experiments: [
+          {
+            data: params,
+            executor: {
+              local: null,
+              name: Executor.WORKSPACE,
+              state: ExperimentStatus.RUNNING
+            },
+            name: 'unselectable-nuffy',
+            rev: EXPERIMENT_WORKSPACE_ID
+          },
+          {},
+          {},
+          {},
+          {
+            data: params,
+            executor: {
+              local: null,
+              name: Executor.DVC_TASK,
+              state: ExperimentStatus.RUNNING
+            },
+            name: 'unselectable-task',
+            rev: EXPERIMENT_WORKSPACE_ID
+          }
+        ],
+        rev: 'testBranch'
+      }
+    )
+
+    const model = new ExperimentsModel('', buildMockMemento())
+    model.transformAndSet(data, false, '')
+
+    expect(model.getSelectedRevisions().map(({ id }) => id)).toStrictEqual([
+      EXPERIMENT_WORKSPACE_ID
+    ])
+
+    model.setSelected([])
+    expect(model.getSelectedRevisions().map(({ id }) => id)).toStrictEqual([])
+
+    model.setSelected(model.getCombinedList())
+    expect(model.getSelectedRevisions().map(({ id }) => id)).toStrictEqual([
+      EXPERIMENT_WORKSPACE_ID,
+      'testBranch',
+      'exp-2',
+      'exp-3',
+      'exp-4'
+    ])
+  })
+
   it('should fetch commit params', () => {
     const model = new ExperimentsModel('', buildMockMemento())
     model.transformAndSet(outputFixture, false, '')

extension/src/experiments/model/index.ts

@@ -5,26 +5,26 @@ import { collectExperiments } from './collect'
 import {
   collectColoredStatus,
   collectFinishedRunningExperiments,
-  collectSelected,
+  collectSelectable,
+  collectSelectedColors,
   collectStartedRunningExperiments
 } from './status/collect'
 import { Color, copyOriginalColors } from './status/colors'
-import {
-  canSelect,
-  limitToMaxSelected,
-  ColoredStatus,
-  tooManySelected,
-  UNSELECTED
-} from './status'
+import { canSelect, ColoredStatus, UNSELECTED } from './status'
 import { collectFlatExperimentParams } from './modify/collect'
 import {
   Experiment,
   isQueued,
+  isRunning,
   isRunningInQueue,
   RunningExperiment
 } from '../webview/contract'
 import { definedAndNonEmpty, reorderListSubset } from '../../util/array'
-import { EXPERIMENT_WORKSPACE_ID, ExpShowOutput } from '../../cli/dvc/contract'
+import {
+  EXPERIMENT_WORKSPACE_ID,
+  Executor,
+  ExpShowOutput
+} from '../../cli/dvc/contract'
 import { flattenMapValues } from '../../util/map'
 import { ModelWithPersistence } from '../../persistence/model'
 import { PersistenceKey } from '../../persistence/constants'
@@ -136,14 +136,16 @@ export class ExperimentsModel extends ModelWithPersistence {
   }
 
   public toggleStatus(id: string) {
-    if (
-      isQueued(
-        this.getExperimentsAndQueued().find(
-          ({ id: queuedId }) => queuedId === id
-        )?.status
-      )
-    ) {
-      return
+    const experiment = this.getExperimentsAndQueued().find(
+      ({ id: expId }) => expId === id
+    )
+
+    if (experiment && isRunning(experiment.status)) {
+      return this.preventSelectionOfRunningExperiment(experiment)
+    }
+
+    if (isQueued(experiment?.status)) {
+      return UNSELECTED
     }
 
     const current = this.coloredStatus[id]
@@ -166,17 +168,6 @@ export class ExperimentsModel extends ModelWithPersistence {
     return this.running.length > 0
   }
 
-  public isRunningInWorkspace(id: string) {
-    if (id === EXPERIMENT_WORKSPACE_ID) {
-      return false
-    }
-
-    return this.running.some(
-      ({ id: runningId, executor }) =>
-        executor === EXPERIMENT_WORKSPACE_ID && runningId === id
-    )
-  }
-
   public hasCheckpoints() {
     return this.checkpoints
   }
@@ -255,12 +246,12 @@ export class ExperimentsModel extends ModelWithPersistence {
   }
 
   public setSelected(selectedExperiments: Experiment[]) {
-    if (tooManySelected(selectedExperiments)) {
-      selectedExperiments = limitToMaxSelected(selectedExperiments)
-    }
+    const possibleToSelect = collectSelectable(selectedExperiments, {
+      ...this.workspace
+    })
 
-    const { availableColors, coloredStatus } = collectSelected(
-      selectedExperiments.filter(({ status }) => !isQueued(status)),
+    const { availableColors, coloredStatus } = collectSelectedColors(
+      possibleToSelect,
       this.getWorkspaceCommitsAndExperiments(),
       this.coloredStatus,
       this.availableColors
@@ -505,6 +496,24 @@ export class ExperimentsModel extends ModelWithPersistence {
     )
   }
 
+  private preventSelectionOfRunningExperiment(
+    experiment: Experiment
+  ): Color | undefined | typeof UNSELECTED {
+    if (isRunningInQueue(experiment)) {
+      return UNSELECTED
+    }
+
+    const { executor, id } = experiment
+    if (
+      executor === Executor.WORKSPACE &&
+      id !== EXPERIMENT_WORKSPACE_ID &&
+      !this.isSelected(id) &&
+      !this.isSelected(EXPERIMENT_WORKSPACE_ID)
+    ) {
+      return this.toggleStatus(EXPERIMENT_WORKSPACE_ID)
+    }
+  }
+
   private persistSorts() {
     return this.persist(PersistenceKey.EXPERIMENTS_SORT_BY, this.currentSorts)
   }

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

@@ -1,15 +1,22 @@
-import { canSelect, ColoredStatus, UNSELECTED } from '.'
+import {
+  canSelect,
+  ColoredStatus,
+  limitToMaxSelected,
+  tooManySelected,
+  UNSELECTED
+} from '.'
 import { Color, copyOriginalColors } from './colors'
 import { hasKey } from '../../../util/object'
 import {
   Experiment,
   isQueued,
   isRunning,
+  isRunningInQueue,
   RunningExperiment
 } from '../../webview/contract'
 import { definedAndNonEmpty, reorderListSubset } from '../../../util/array'
 import { flattenMapValues } from '../../../util/map'
-import { EXPERIMENT_WORKSPACE_ID } from '../../../cli/dvc/contract'
+import { Executor, EXPERIMENT_WORKSPACE_ID } from '../../../cli/dvc/contract'
 
 const canAssign = (
   coloredStatus: ColoredStatus,
@@ -187,7 +194,50 @@ const assignSelected = (
   return { availableColors, coloredStatus }
 }
 
-export const collectSelected = (
+const cannotSelect = (
+  ids: Set<string>,
+  { executor, id, status }: Experiment
+): boolean =>
+  isQueued(status) || isRunningInQueue({ executor, status }) || ids.has(id)
+
+const shouldSelectWorkspace = ({ executor, status }: Experiment): boolean =>
+  executor === Executor.WORKSPACE && isRunning(status)
+
+const collectWorkspace = (
+  acc: Experiment[],
+  collectedIds: Set<string>,
+  workspace: Experiment
+) => {
+  if (!collectedIds.has(EXPERIMENT_WORKSPACE_ID)) {
+    acc.push(workspace)
+  }
+  collectedIds.add(EXPERIMENT_WORKSPACE_ID)
+}
+
+export const collectSelectable = (
+  selectedExperiments: Experiment[],
+  workspace: Experiment
+): Experiment[] => {
+  const acc: Experiment[] = []
+  const collectedIds = new Set<string>()
+
+  for (const experiment of selectedExperiments) {
+    if (cannotSelect(collectedIds, experiment)) {
+      continue
+    }
+
+    if (shouldSelectWorkspace(experiment)) {
+      collectWorkspace(acc, collectedIds, workspace)
+      continue
+    }
+    acc.push(experiment)
+    collectedIds.add(experiment.id)
+  }
+
+  return tooManySelected(acc) ? limitToMaxSelected(acc) : acc
+}
+
+export const collectSelectedColors = (
   selectedExperiments: Experiment[],
   experiments: Experiment[],
   previousStatus: ColoredStatus,
@@ -217,7 +267,10 @@ export const collectStartedRunningExperiments = (
   const acc = new Set<string>()
 
   for (const { id: runningId, executor } of nowRunning) {
-    if (previouslyRunning.some(({ id }) => id === runningId)) {
+    if (
+      executor === Executor.DVC_TASK ||
+      previouslyRunning.some(({ id }) => id === runningId)
+    ) {
       continue
     }
     acc.add(

extension/src/experiments/workspace.ts

@@ -142,12 +142,12 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews<
     return this.getRepository(dvcRoot).removeSorts()
   }
 
-  public async selectExperiments(overrideRoot?: string) {
+  public async selectExperimentsToPlot(overrideRoot?: string) {
     const dvcRoot = await this.getDvcRoot(overrideRoot)
     if (!dvcRoot) {
       return
     }
-    return this.getRepository(dvcRoot).selectExperiments()
+    return this.getRepository(dvcRoot).selectExperimentsToPlot()
   }
 
   public async selectColumns(overrideRoot?: string) {

extension/src/plots/webview/messages.ts

@@ -220,7 +220,7 @@ export class WebviewMessages {
   }
 
   private selectExperimentsFromWebview() {
-    void this.experiments.selectExperiments()
+    void this.experiments.selectExperimentsToPlot()
     sendTelemetryEvent(
       EventName.VIEWS_PLOTS_SELECT_EXPERIMENTS,
       undefined,

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

@@ -150,7 +150,7 @@ const data: Commit[] = [
             'c3961d777cfbd7727f9fde4851896006'
           )
         },
-        displayColor: colorsList[1],
+        displayColor: undefined,
         description: '[exp-e7a67]',
         executor: Executor.DVC_TASK,
         id: 'exp-e7a67',
@@ -179,7 +179,7 @@ const data: Commit[] = [
           }
         },
         status: ExperimentStatus.RUNNING,
-        selected: true,
+        selected: false,
         sha: '4fb124aebddb2adf1545030907687fa9a4c80e70',
         starred: false,
         Created: '2020-12-29T15:31:52'