Eliminate checkpoint race condition when updating live plots

extension/src/experiments/index.ts

@@ -9,6 +9,7 @@ import {
 } from './model/filterBy/quickPick'
 import { pickSortsToRemove, pickSortToAdd } from './model/sortBy/quickPick'
 import { MetricsAndParamsModel } from './metricsAndParams/model'
+import { CheckpointsModel } from './checkpoints/model'
 import { ExperimentsData } from './data'
 import { TableData } from './webview/contract'
 import { ResourceLocator } from '../resourceLocator'
@@ -22,17 +23,20 @@ import {
   MessageFromWebviewType
 } from '../webview/contract'
 import { Logger } from '../common/logger'
+import { FileSystemData } from '../fileSystem/data'
 
 export class Experiments extends BaseRepository<TableData> {
   public readonly onDidChangeExperiments: Event<ExperimentsOutput | void>
   public readonly onDidChangeMetricsOrParams: Event<void>
 
   public readonly viewKey = ViewKey.EXPERIMENTS
 
-  private readonly data: ExperimentsData
+  private readonly cliData: ExperimentsData
+  private readonly fileSystemData: FileSystemData
 
   private readonly experiments: ExperimentsModel
   private readonly metricsAndParams: MetricsAndParamsModel
+  private readonly checkpoints: CheckpointsModel
 
   private readonly experimentsChanged = this.dispose.track(
     new EventEmitter<ExperimentsOutput | void>()
@@ -48,7 +52,8 @@ export class Experiments extends BaseRepository<TableData> {
     updatesPaused: EventEmitter<boolean>,
     resourceLocator: ResourceLocator,
     workspaceState: Memento,
-    data?: ExperimentsData
+    cliData?: ExperimentsData,
+    fileSystemData?: FileSystemData
   ) {
     super(dvcRoot, resourceLocator.beaker)
 
@@ -63,11 +68,22 @@ export class Experiments extends BaseRepository<TableData> {
       new MetricsAndParamsModel(dvcRoot, workspaceState)
     )
 
-    this.data = this.dispose.track(
-      data || new ExperimentsData(dvcRoot, internalCommands, updatesPaused)
+    this.checkpoints = this.dispose.track(new CheckpointsModel())
+
+    this.cliData = this.dispose.track(
+      cliData || new ExperimentsData(dvcRoot, internalCommands, updatesPaused)
     )
 
-    this.dispose.track(this.data.onDidUpdate(data => this.setState(data)))
+    this.fileSystemData = this.dispose.track(
+      fileSystemData || new FileSystemData(dvcRoot)
+    )
+
+    this.dispose.track(this.cliData.onDidUpdate(data => this.setState(data)))
+    this.dispose.track(
+      this.fileSystemData.onDidUpdate(data =>
+        this.checkpoints.transformAndSet(data)
+      )
+    )
 
     this.handleMessageFromWebview()
 
@@ -81,11 +97,11 @@ export class Experiments extends BaseRepository<TableData> {
   }
 
   public update() {
-    return this.data.managedUpdate()
+    return this.cliData.managedUpdate()
   }
 
   public forceUpdate() {
-    return this.data.forceUpdate()
+    return this.cliData.forceUpdate()
   }
 
   public async setState(data: ExperimentsOutput) {
@@ -97,6 +113,10 @@ export class Experiments extends BaseRepository<TableData> {
     return this.notifyChanged(data)
   }
 
+  public hasCheckpoints() {
+    return this.checkpoints.hasCheckpoints()
+  }
+
   public getChildMetricsOrParams(path?: string) {
     return this.metricsAndParams.getChildren(path)
   }

extension/src/plots/model/collect.test.ts

@@ -81,9 +81,9 @@ describe('collectBranchRevision', () => {
 })
 
 describe('collectMutableRevisions', () => {
-  it('should return all of the running non-checkpoint experiments from the test fixture', () => {
-    const mutable = collectMutableRevisions(expShowFixture)
-    expect(mutable).toEqual(['workspace'])
+  it('should always return an empty array when checkpoints are present', () => {
+    const mutable = collectMutableRevisions(expShowFixture, true)
+    expect(mutable).toEqual([])
   })
 
   it('should return all running experiments when there are no checkpoints', () => {
@@ -131,7 +131,7 @@ describe('collectMutableRevisions', () => {
       }
     }
 
-    const mutable = collectMutableRevisions(experimentsRunningInTemp)
+    const mutable = collectMutableRevisions(experimentsRunningInTemp, false)
     expect(mutable).toEqual(['6ee95de', 'ebaa07e'])
   })
 })

extension/src/plots/model/collect.ts

@@ -289,7 +289,14 @@ const collectMutableFromExperiment = (
   })
 }
 
-export const collectMutableRevisions = (data: ExperimentsOutput): string[] => {
+export const collectMutableRevisions = (
+  data: ExperimentsOutput,
+  hasCheckpoints: boolean
+): string[] => {
+  if (hasCheckpoints) {
+    return []
+  }
+
   const acc: string[] = []
 
   if (data.workspace.baseline.data?.running) {

extension/src/plots/model/index.ts

@@ -93,7 +93,7 @@ export class PlotsModel {
         collectLivePlotsData(data),
         collectRevisions(data),
         collectBranchRevision(data),
-        collectMutableRevisions(data)
+        collectMutableRevisions(data, this.experiments.hasCheckpoints())
       ])
 
     const { branchNames, revisionsByTip, revisionsByBranch } = revisions

extension/src/test/suite/experiments/util.ts

@@ -7,6 +7,7 @@ import expShowFixture from '../../fixtures/expShow/output'
 import { buildMockMemento, dvcDemoPath } from '../../util'
 import { buildDependencies, buildMockData } from '../util'
 import { ExperimentsData } from '../../../experiments/data'
+import { FileSystemData } from '../../../fileSystem/data'
 
 export const buildExperiments = (
   disposer: Disposer,
@@ -29,7 +30,8 @@ export const buildExperiments = (
       updatesPaused,
       resourceLocator,
       buildMockMemento(),
-      buildMockData<ExperimentsData>()
+      buildMockData<ExperimentsData>(),
+      buildMockData<FileSystemData>()
     )
   )