Use exp show data along with runner to set dvc experiment running context value (#1848)

Author: mattseddon

extension/package.json

@@ -10,7 +10,6 @@ import { getOptions } from './options'
 import { Config } from '../config'
 import { PseudoTerminal } from '../vscode/pseudoTerminal'
 import { createProcess, Process } from '../processExecution'
-import { setContextValue } from '../vscode/context'
 import { StopWatch } from '../util/time'
 import { sendErrorTelemetryEvent, sendTelemetryEvent } from '../telemetry'
 import { EventName } from '../telemetry/constants'
@@ -70,7 +69,6 @@ export class CliRunner extends Disposable implements ICli {
     this.dispose.track(
       this.onDidCompleteProcess(() => {
         this.pseudoTerminal.setBlocked(false)
-        CliRunner.setRunningContext(false)
         this.processOutput.fire(
           '\r\nTerminal will be reused by DVC, press any key to close it\r\n\n'
         )
@@ -101,9 +99,6 @@ export class CliRunner extends Disposable implements ICli {
     )
   }
 
-  private static setRunningContext = (isRunning: boolean) =>
-    setContextValue('dvc.runner.running', isRunning)
-
   public runExperiment(dvcRoot: string, ...args: Args) {
     return this.run(
       dvcRoot,
@@ -207,7 +202,6 @@ export class CliRunner extends Disposable implements ICli {
   }
 
   private startProcess(cwd: string, args: Args) {
-    CliRunner.setRunningContext(true)
     this.pseudoTerminal.setBlocked(true)
     this.processOutput.fire(`Running: dvc ${args.join(' ')}\r\n\n`)
     this.currentProcess = this.createProcess({

extension/src/cli/runner.ts

@@ -0,0 +1,49 @@
+import { Disposable } from './class/dispose'
+import { CliRunner } from './cli/runner'
+import { WorkspaceExperiments } from './experiments/workspace'
+import { setContextValue } from './vscode/context'
+
+export class Context extends Disposable {
+  private readonly experiments: WorkspaceExperiments
+  private readonly cliRunner: CliRunner
+
+  constructor(experiments: WorkspaceExperiments, cliRunner: CliRunner) {
+    super()
+
+    this.experiments = experiments
+    this.cliRunner = cliRunner
+
+    this.dispose.track(
+      this.cliRunner.onDidStartProcess(() => {
+        this.setIsExperimentRunning()
+      })
+    )
+
+    this.dispose.track(
+      this.cliRunner.onDidCompleteProcess(({ cwd }) =>
+        this.experiments.getRepository(cwd).update()
+      )
+    )
+
+    this.dispose.track(
+      this.experiments.onDidChangeExperiments(() => {
+        this.setIsExperimentRunning()
+      })
+    )
+  }
+
+  private setIsExperimentRunning() {
+    if (this.cliRunner.isExperimentRunning()) {
+      setContextValue('dvc.experiment.running', true)
+      return
+    }
+
+    let hasRunning = false
+    for (const dvcRoot of this.experiments.getDvcRoots()) {
+      if (this.experiments.getRepository(dvcRoot).hasRunningExperiment()) {
+        hasRunning = true
+      }
+    }
+    setContextValue('dvc.experiment.running', hasRunning)
+  }
+}

extension/src/context.ts

@@ -377,6 +377,10 @@ export class Experiments extends BaseRepository<TableData> {
     }
   }
 
+  public hasRunningExperiment() {
+    return this.experiments.hasRunningExperiment()
+  }
+
   private hideTableColumn(path: string) {
     this.toggleColumnStatus(path)
     sendTelemetryEvent(
@@ -438,6 +442,7 @@ export class Experiments extends BaseRepository<TableData> {
       columns: this.columns.getSelected(),
       hasCheckpoints: this.hasCheckpoints(),
       hasColumns: this.columns.hasColumns(),
+      hasRunningExperiment: this.experiments.hasRunningExperiment(),
       rows: this.experiments.getRowData(),
       sorts: this.experiments.getSorts()
     }

extension/src/experiments/index.ts

@@ -5,10 +5,12 @@ export class ExperimentsAccumulator {
   public branches: Experiment[] = []
   public checkpointsByTip: Map<string, Experiment[]> = new Map()
   public experimentsByBranch: Map<string, Experiment[]> = new Map()
+  public hasRunning: boolean
 
   constructor(workspace: Experiment | undefined) {
     if (workspace) {
       this.workspace = workspace
     }
+    this.hasRunning = !!workspace?.running
   }
 }

extension/src/experiments/model/accumulator.ts

@@ -192,6 +192,15 @@ const transformExperimentOrCheckpointData = (
   }
 }
 
+const collectHasRunningExperiment = (
+  acc: ExperimentsAccumulator,
+  experiment: Experiment
+) => {
+  if (experiment.running) {
+    acc.hasRunning = true
+  }
+}
+
 const collectExperimentOrCheckpoint = (
   acc: ExperimentsAccumulator,
   experiment: Experiment,
@@ -227,6 +236,7 @@ const collectFromExperimentsObject = (
     }
 
     collectExperimentOrCheckpoint(acc, experiment, branchName, checkpointTipId)
+    collectHasRunningExperiment(acc, experiment)
   }
 }
 
@@ -248,6 +258,7 @@ const collectFromBranchesObject = (
 
     if (branch) {
       collectFromExperimentsObject(acc, experimentsObject, sha, branch.label)
+      collectHasRunningExperiment(acc, branch)
 
       acc.branches.push(branch)
     }

extension/src/experiments/model/collect.ts

@@ -56,6 +56,7 @@ export class ExperimentsModel extends ModelWithPersistence {
   private useFiltersForSelection = false
 
   private currentSorts: SortDefinition[]
+  private hasRunning = false
 
   constructor(dvcRoot: string, workspaceState: Memento) {
     super(dvcRoot, workspaceState)
@@ -85,13 +86,19 @@ export class ExperimentsModel extends ModelWithPersistence {
   }
 
   public transformAndSet(data: ExperimentsOutput) {
-    const { workspace, branches, experimentsByBranch, checkpointsByTip } =
-      collectExperiments(data)
+    const {
+      workspace,
+      branches,
+      experimentsByBranch,
+      checkpointsByTip,
+      hasRunning
+    } = collectExperiments(data)
 
     this.workspace = workspace
     this.branches = branches
     this.experimentsByBranch = experimentsByBranch
     this.checkpointsByTip = checkpointsByTip
+    this.hasRunning = hasRunning
 
     this.setColoredStatus()
   }
@@ -117,6 +124,10 @@ export class ExperimentsModel extends ModelWithPersistence {
     return this.coloredStatus[id]
   }
 
+  public hasRunningExperiment() {
+    return this.hasRunning
+  }
+
   public canSelect() {
     return canSelect(this.coloredStatus)
   }

extension/src/experiments/model/index.ts

@@ -57,6 +57,7 @@ export type TableData = {
   columnWidths: Record<string, number>
   hasCheckpoints: boolean
   hasColumns: boolean
+  hasRunningExperiment: boolean
   rows: Row[]
   sorts: SortDefinition[]
 }

extension/src/experiments/webview/contract.ts

@@ -19,6 +19,8 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews<
     new EventEmitter<void>()
   )
 
+  public readonly onDidChangeExperiments = this.experimentsChanged.event
+
   public readonly columnsChanged = this.dispose.track(new EventEmitter<void>())
 
   public readonly updatesPaused: EventEmitter<boolean>

extension/src/experiments/workspace.ts

@@ -1,8 +1,9 @@
 import { commands, env, Event, EventEmitter, ExtensionContext } from 'vscode'
-import { Config } from './config'
 import { CliExecutor } from './cli/executor'
 import { CliRunner } from './cli/runner'
 import { CliReader } from './cli/reader'
+import { Config } from './config'
+import { Context } from './context'
 import { isVersionCompatible } from './cli/version'
 import { isPythonExtensionInstalled } from './extensions/python'
 import { WorkspaceExperiments } from './experiments/workspace'
@@ -133,11 +134,7 @@ export class Extension extends Disposable implements IExtension {
       new WorkspaceRepositories(this.internalCommands)
     )
 
-    this.dispose.track(
-      this.cliRunner.onDidCompleteProcess(({ cwd }) =>
-        this.experiments.getRepository(cwd).update()
-      )
-    )
+    this.dispose.track(new Context(this.experiments, this.cliRunner))
 
     this.dispose.track(
       new ExperimentsColumnsTree(