diff --git a/extension/src/experiments/index.ts b/extension/src/experiments/index.ts index 602795b39b..f15019b2e3 100644 --- a/extension/src/experiments/index.ts +++ b/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,6 +23,7 @@ import { MessageFromWebviewType } from '../webview/contract' import { Logger } from '../common/logger' +import { FileSystemData } from '../fileSystem/data' export class Experiments extends BaseRepository { public readonly onDidChangeExperiments: Event @@ -29,10 +31,12 @@ export class Experiments extends BaseRepository { 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() @@ -48,7 +52,8 @@ export class Experiments extends BaseRepository { updatesPaused: EventEmitter, resourceLocator: ResourceLocator, workspaceState: Memento, - data?: ExperimentsData + cliData?: ExperimentsData, + fileSystemData?: FileSystemData ) { super(dvcRoot, resourceLocator.beaker) @@ -63,11 +68,22 @@ export class Experiments extends BaseRepository { 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 { } 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 { return this.notifyChanged(data) } + public hasCheckpoints() { + return this.checkpoints.hasCheckpoints() + } + public getChildMetricsOrParams(path?: string) { return this.metricsAndParams.getChildren(path) } diff --git a/extension/src/plots/model/collect.test.ts b/extension/src/plots/model/collect.test.ts index ec186a211a..a1b0196122 100644 --- a/extension/src/plots/model/collect.test.ts +++ b/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']) }) }) diff --git a/extension/src/plots/model/collect.ts b/extension/src/plots/model/collect.ts index b775483c72..bcf6c50da8 100644 --- a/extension/src/plots/model/collect.ts +++ b/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) { diff --git a/extension/src/plots/model/index.ts b/extension/src/plots/model/index.ts index ecb18ab708..50564f8eb4 100644 --- a/extension/src/plots/model/index.ts +++ b/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 diff --git a/extension/src/test/suite/experiments/util.ts b/extension/src/test/suite/experiments/util.ts index 07e1ee3e89..51001cbe61 100644 --- a/extension/src/test/suite/experiments/util.ts +++ b/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() + buildMockData(), + buildMockData() ) )