diff --git a/extension/src/experiments/index.ts b/extension/src/experiments/index.ts index 18c01d09a0..68a33a663c 100644 --- a/extension/src/experiments/index.ts +++ b/extension/src/experiments/index.ts @@ -43,7 +43,11 @@ import { createTypedAccumulator } from '../util/object' import { pickPaths } from '../path/selection/quickPick' import { Toast } from '../vscode/toast' import { ConfigKey, getConfigValue, setUserConfigValue } from '../vscode/config' -import { checkSignalFile, pollSignalFileForProcess } from '../fileSystem' +import { + checkSignalFile, + getEntryFromJsonFile, + pollSignalFileForProcess +} from '../fileSystem' import { DVCLIVE_ONLY_RUNNING_SIGNAL_FILE } from '../cli/dvc/constants' import { Response } from '../vscode/response' import { Pipeline } from '../pipeline' @@ -599,6 +603,19 @@ export class Experiments extends BaseRepository { return this.webviewMessages.sendWebviewMessage() } + public hasDvcLiveOnlyRunning() { + return this.experiments.hasDvcLiveOnlyRunning() + } + + public checkWorkspaceDuplicated(fetched: string[]) { + const updated = this.experiments.checkWorkspaceDuplicated(fetched) + if (!updated) { + return + } + + this.notifyChanged() + } + protected sendInitialWebviewData() { return this.webviewMessages.sendWebviewMessage() } @@ -683,9 +700,13 @@ export class Experiments extends BaseRepository { } private async checkSignalFile() { - const dvcLiveOnly = await checkSignalFile(this.dvcLiveOnlySignalFile) + const running = await checkSignalFile(this.dvcLiveOnlySignalFile) - if (dvcLiveOnly && !this.dvcLiveOnlyCleanupInitialized) { + if (!running) { + return { running } + } + + if (!this.dvcLiveOnlyCleanupInitialized) { this.dvcLiveOnlyCleanupInitialized = true void pollSignalFileForProcess(this.dvcLiveOnlySignalFile, () => { this.dvcLiveOnlyCleanupInitialized = false @@ -694,7 +715,9 @@ export class Experiments extends BaseRepository { } }) } - return dvcLiveOnly + const expName = getEntryFromJsonFile(this.dvcLiveOnlySignalFile, 'exp_name') + + return { expName, running } } private async informMaxSelected() { diff --git a/extension/src/experiments/model/index.test.ts b/extension/src/experiments/model/index.test.ts index a20c6999af..62cdd261b4 100644 --- a/extension/src/experiments/model/index.test.ts +++ b/extension/src/experiments/model/index.test.ts @@ -39,10 +39,10 @@ beforeEach(() => { const DEFAULT_DATA: [ string, - boolean, + { running: boolean; expName?: string }, { branch: string; sha: string }[], { [branch: string]: number } -] = ['', false, [], { main: 2000 }] +] = ['', { running: false }, [], { main: 2000 }] type TransformAndSetInputs = [ExpShowOutput, ...typeof DEFAULT_DATA] @@ -52,7 +52,7 @@ describe('ExperimentsModel', () => { model.transformAndSetLocal( outputFixture, gitLogFixture, - false, + { running: false }, rowOrderFixture, { main: 6 } ) @@ -65,7 +65,7 @@ describe('ExperimentsModel', () => { model.transformAndSetLocal( survivalOutputFixture, '', - false, + { running: false }, [ { branch: 'main', sha: '3d5adcb974bb2c85917a5d61a489b933adaa2b7f' }, { branch: 'main', sha: 'a49e03966a1f9f1299ec222ebc4bed8625d2c54d' }, @@ -107,7 +107,9 @@ describe('ExperimentsModel', () => { } ) - model.transformAndSetLocal(dvcLiveOnly, '', true, [], { main: 2000 }) + model.transformAndSetLocal(dvcLiveOnly, '', { running: true }, [], { + main: 2000 + }) // eslint-disable-next-line @typescript-eslint/no-explicit-any const runningWorkspace = (model as any).workspace expect(runningWorkspace?.executor).toStrictEqual(EXPERIMENT_WORKSPACE_ID) @@ -203,7 +205,7 @@ describe('ExperimentsModel', () => { model.transformAndSetLocal( deeplyNestedOutputFixture, '', - false, + { running: false }, [{ branch: 'main', sha: '53c3851f46955fa3e2b8f6e1c52999acc8c9ea77' }], { main: 10 } ) @@ -217,7 +219,7 @@ describe('ExperimentsModel', () => { model.transformAndSetLocal( dataTypesOutputFixture, '', - false, + { running: false }, [{ branch: 'main', sha: '53c3851f46955fa3e2b8f6e1c52999acc8c9ea77' }], { main: 10 } ) @@ -483,7 +485,7 @@ describe('ExperimentsModel', () => { const runningExperimentData: TransformAndSetInputs = [ outputFixture, gitLogFixture, - false, + { running: false }, [], { main: 2000 @@ -504,7 +506,7 @@ describe('ExperimentsModel', () => { } as ExpWithError ], gitLogFixture, - false, + { running: false }, [], { main: 2000 @@ -586,7 +588,7 @@ describe('ExperimentsModel', () => { error: { msg: errorMsg, type: 'caught error' } } ] - model.transformAndSetLocal(data, gitLogFixture, false, [], { + model.transformAndSetLocal(data, gitLogFixture, { running: false }, [], { main: 6 }) @@ -595,7 +597,7 @@ describe('ExperimentsModel', () => { model.transformAndSetLocal( outputFixture, gitLogFixture, - false, + { running: false }, rowOrderFixture, { main: 6 } ) diff --git a/extension/src/experiments/model/index.ts b/extension/src/experiments/model/index.ts index e16fdfeede..e573579c24 100644 --- a/extension/src/experiments/model/index.ts +++ b/extension/src/experiments/model/index.ts @@ -29,7 +29,12 @@ import { WORKSPACE_BRANCH } from '../webview/contract' import { reorderListSubset } from '../../util/array' -import { Executor, ExpShowOutput, ExecutorStatus } from '../../cli/dvc/contract' +import { + Executor, + ExpShowOutput, + ExecutorStatus, + EXPERIMENT_WORKSPACE_ID +} from '../../cli/dvc/contract' import { flattenMapValues } from '../../util/map' import { ModelWithPersistence } from '../../persistence/model' import { PersistenceKey } from '../../persistence/constants' @@ -81,6 +86,8 @@ export class ExperimentsModel extends ModelWithPersistence { private running: RunningExperiment[] = [] private startedRunning: Set = new Set() + private dvcLiveOnlyExpName: string | undefined + constructor(dvcRoot: string, workspaceState: Memento) { super(dvcRoot, workspaceState) @@ -126,7 +133,7 @@ export class ExperimentsModel extends ModelWithPersistence { public transformAndSetLocal( expShow: ExpShowOutput, gitLog: string, - dvcLiveOnly: boolean, + dvcLiveOnly: { running: boolean; expName?: string }, rowOrder: { branch: string; sha: string }[], availableNbCommits: { [branch: string]: number } ) { @@ -137,7 +144,11 @@ export class ExperimentsModel extends ModelWithPersistence { hasCheckpoints, runningExperiments, workspace - } = collectExperiments(expShow, gitLog, dvcLiveOnly) + } = collectExperiments(expShow, gitLog, dvcLiveOnly.running) + + if (dvcLiveOnly.expName) { + this.dvcLiveOnlyExpName = dvcLiveOnly.expName + } const { hasMoreCommits, isShowingMoreCommits } = collectAddRemoveCommitsDetails(availableNbCommits, (branch: string) => @@ -528,6 +539,28 @@ export class ExperimentsModel extends ModelWithPersistence { return this.availableBranchesToSelect } + public hasDvcLiveOnlyRunning() { + return !!this.dvcLiveOnlyExpName + } + + public checkWorkspaceDuplicated(fetched: string[]) { + if (!this.dvcLiveOnlyExpName) { + return false + } + + const newExperimentFetched = fetched.includes(this.dvcLiveOnlyExpName) + const workspaceSelectionDuplicated = + this.coloredStatus[EXPERIMENT_WORKSPACE_ID] === + this.coloredStatus[this.dvcLiveOnlyExpName] + + if (newExperimentFetched && workspaceSelectionDuplicated) { + this.coloredStatus[EXPERIMENT_WORKSPACE_ID] = UNSELECTED + this.dvcLiveOnlyExpName = undefined + this.persistStatus() + return true + } + } + private findIndexByPath(pathToRemove: string) { return this.currentSorts.findIndex(({ path }) => path === pathToRemove) } @@ -590,7 +623,8 @@ export class ExperimentsModel extends ModelWithPersistence { this.experimentsByCommit, this.coloredStatus, this.availableColors, - this.startedRunning + this.startedRunning, + this.dvcLiveOnlyExpName ) this.startedRunning = new Set() diff --git a/extension/src/experiments/model/status/collect.test.ts b/extension/src/experiments/model/status/collect.test.ts index c0580c33c5..f21db36dc1 100644 --- a/extension/src/experiments/model/status/collect.test.ts +++ b/extension/src/experiments/model/status/collect.test.ts @@ -27,7 +27,8 @@ describe('collectColoredStatus', () => { new Map(), {}, copyOriginalColors(), - new Set() + new Set(), + undefined ) expect(availableColors).toStrictEqual(colors) @@ -48,7 +49,8 @@ describe('collectColoredStatus', () => { new Map(), {}, copyOriginalColors(), - new Set() + new Set(), + undefined ) expect(availableColors).toStrictEqual(colors) @@ -74,7 +76,8 @@ describe('collectColoredStatus', () => { exp7: colors[6] }, [], - new Set(['exp8']) + new Set(['exp8']), + undefined ) expect(availableColors).toStrictEqual([]) @@ -107,7 +110,8 @@ describe('collectColoredStatus', () => { exp8: UNSELECTED }, copyOriginalColors().slice(3), - new Set(['exp1']) + new Set(['exp1']), + undefined ) expect(coloredStatus).toStrictEqual({ exp1: colors[0] }) @@ -129,7 +133,8 @@ describe('collectColoredStatus', () => { exp9: colors[1] }, unassignedColors, - new Set() + new Set(), + undefined ) expect(coloredStatus).toStrictEqual({ @@ -157,7 +162,8 @@ describe('collectColoredStatus', () => { new Map(), { exp9: colors[0] }, unassignColors, - new Set() + new Set(), + undefined ) expect(availableColors).toStrictEqual(unassignColors) @@ -190,7 +196,8 @@ describe('collectColoredStatus', () => { exp9: colors[5] }, copyOriginalColors().slice(6), - new Set(['exp1', 'exp2', 'exp3']) + new Set(['exp1', 'exp2', 'exp3']), + undefined ) expect(availableColors).toStrictEqual([]) @@ -228,7 +235,8 @@ describe('collectColoredStatus', () => { workspace: UNSELECTED }, colors, - new Set() + new Set(), + undefined ) expect(coloredStatus).toStrictEqual({ 'exp-1': selectedColor, @@ -240,4 +248,45 @@ describe('collectColoredStatus', () => { colors.filter(color => color !== selectedColor) ) }) + + it("should duplicate the workspace's color when a new experiment is provided that has the same name as found in the DVCLive only signal file", () => { + const colors = copyOriginalColors() + const selectedColor = colors[2] + const { availableColors, coloredStatus } = collectColoredStatus( + [ + { + executor: null, + executorStatus: ExecutorStatus.SUCCESS, + id: 'exp-1' + }, + { + id: EXPERIMENT_WORKSPACE_ID + }, + { id: 'main' }, + { + executor: null, + executorStatus: ExecutorStatus.SUCCESS, + id: 'exp-2' + } + ] as Experiment[], + new Map(), + { + 'exp-1': UNSELECTED, + workspace: selectedColor + }, + colors, + new Set(), + 'exp-2' + ) + expect(coloredStatus).toStrictEqual({ + [EXPERIMENT_WORKSPACE_ID]: selectedColor, + 'exp-1': UNSELECTED, + 'exp-2': selectedColor, + main: UNSELECTED + }) + + expect(availableColors).toStrictEqual( + colors.filter(color => color !== selectedColor) + ) + }) }) diff --git a/extension/src/experiments/model/status/collect.ts b/extension/src/experiments/model/status/collect.ts index 584ee5a52c..1baa940cf2 100644 --- a/extension/src/experiments/model/status/collect.ts +++ b/extension/src/experiments/model/status/collect.ts @@ -10,19 +10,42 @@ import { hasKey } from '../../../util/object' import { Experiment, isQueued, RunningExperiment } from '../../webview/contract' import { definedAndNonEmpty, reorderListSubset } from '../../../util/array' import { flattenMapValues } from '../../../util/map' -import { Executor } from '../../../cli/dvc/contract' +import { Executor, EXPERIMENT_WORKSPACE_ID } from '../../../cli/dvc/contract' const canAssign = ( coloredStatus: ColoredStatus, unassignedColors: Color[] ): boolean => canSelect(coloredStatus) && definedAndNonEmpty(unassignedColors) -const collectStatus = (acc: ColoredStatus, experiment: Experiment): void => { +const isWorkspaceSelected = (acc: ColoredStatus) => + !!acc[EXPERIMENT_WORKSPACE_ID] + +const isFinishedDvcLiveOnlyExp = ( + id: string, + dvcLiveOnlyExpName: string | undefined +): boolean => id === dvcLiveOnlyExpName + +const duplicateWorkspaceColor = (acc: ColoredStatus, id: string) => { + acc[id] = acc[EXPERIMENT_WORKSPACE_ID] +} + +const collectStatus = ( + acc: ColoredStatus, + experiment: Experiment, + dvcLiveOnlyExpName: string | undefined +): void => { const { id, executorStatus: status } = experiment if (!id || isQueued(status) || hasKey(acc, id)) { return } + if ( + isWorkspaceSelected(acc) && + isFinishedDvcLiveOnlyExp(id, dvcLiveOnlyExpName) + ) { + return duplicateWorkspaceColor(acc, id) + } + acc[id] = UNSELECTED } @@ -89,7 +112,8 @@ export const collectColoredStatus = ( experimentsByCommit: Map, previousStatus: ColoredStatus, unassignedColors: Color[], - startedRunning: Set + startedRunning: Set, + dvcLiveOnlyExpName: string | undefined ): { coloredStatus: ColoredStatus; availableColors: Color[] } => { const flatExperimentsByCommit = flattenMapValues(experimentsByCommit) let availableColors = unassignColors( @@ -115,7 +139,7 @@ export const collectColoredStatus = ( ...workspaceAndCommits, ...flatExperimentsByCommit ]) { - collectStatus(coloredStatus, experiment) + collectStatus(coloredStatus, experiment, dvcLiveOnlyExpName) } return { diff --git a/extension/src/fileSystem/index.test.ts b/extension/src/fileSystem/index.test.ts index 17644bb84e..7b811dc8ea 100644 --- a/extension/src/fileSystem/index.test.ts +++ b/extension/src/fileSystem/index.test.ts @@ -20,13 +20,15 @@ import { writeCsv, writeTsv, isPathInProject, - getPidFromFile + getPidFromFile, + getEntryFromJsonFile } from '.' import { dvcDemoPath } from '../test/util' import { DOT_DVC } from '../cli/dvc/constants' import { ScriptCommand } from '../pipeline' import { processExists } from '../process/execution' +jest.mock('../common/logger') jest.mock('../cli/dvc/reader') jest.mock('../process/execution') jest.mock('fs-extra', () => { @@ -518,3 +520,40 @@ describe('getPidFromFile', () => { expect(pid).toBe(3676) }) }) + +describe('getEntryFromJsonFile', () => { + it('should return undefined if the file does not exist', () => { + const undef = getEntryFromJsonFile('no-file', 'no-keys') + + expect(mockedReadFileSync).toHaveBeenCalledTimes(1) + expect(undef).toBeUndefined() + }) + + it('should return undefined for a file containing a string', () => { + mockedReadFileSync.mockReturnValueOnce(Buffer.from('string')) + const undef = getEntryFromJsonFile(__filename, 'no-keys') + + expect(mockedReadFileSync).toHaveBeenCalledTimes(1) + expect(undef).toBeUndefined() + }) + + it('should return undefined if the file does not contain the key', () => { + mockedReadFileSync.mockReturnValueOnce( + JSON.stringify({ 'other-key': '3676' }) + ) + mockedProcessExists.mockResolvedValueOnce(true) + const undef = getEntryFromJsonFile(__filename, 'key') + + expect(mockedReadFileSync).toHaveBeenCalledTimes(1) + expect(undef).toBeUndefined() + }) + + it('should return the value if the file contains the key', () => { + mockedReadFileSync.mockReturnValueOnce(JSON.stringify({ key: 'value' })) + mockedProcessExists.mockResolvedValueOnce(true) + const value = getEntryFromJsonFile(__filename, 'key') + + expect(mockedReadFileSync).toHaveBeenCalledTimes(1) + expect(value).toStrictEqual('value') + }) +}) diff --git a/extension/src/fileSystem/index.ts b/extension/src/fileSystem/index.ts index 314406100b..fe81f5901b 100644 --- a/extension/src/fileSystem/index.ts +++ b/extension/src/fileSystem/index.ts @@ -279,6 +279,20 @@ export const getPidFromFile = async ( return pid } +export const getEntryFromJsonFile = ( + path: string, + key: string +): string | undefined => { + const json = loadJson(path) + if (!json) { + return + } + + try { + return (json as { [key: string]: string })[key] + } catch {} +} + export const checkSignalFile = async (path: string): Promise => { return !!(await getPidFromFile(path)) } diff --git a/extension/src/plots/index.ts b/extension/src/plots/index.ts index 8dc77ac473..812b0add28 100644 --- a/extension/src/plots/index.ts +++ b/extension/src/plots/index.ts @@ -189,7 +189,24 @@ export class Plots extends BaseRepository { private async sendPlots() { await this.isReady() - return this.webviewMessages.sendWebviewMessage() + await this.webviewMessages.sendWebviewMessage() + return this.checkDvcLiveOnlyDuplicate() + } + + private checkDvcLiveOnlyDuplicate() { + if (!this.experiments.hasDvcLiveOnlyRunning()) { + return + } + + const fetchedRevs = [] + for (const { id, fetched } of this.plots.getSelectedRevisionDetails()) { + if (!fetched) { + continue + } + fetchedRevs.push(id) + } + + return this.experiments.checkWorkspaceDuplicated(fetchedRevs) } private createWebviewMessageHandler(