diff --git a/extension/src/cli/dvc/constants.ts b/extension/src/cli/dvc/constants.ts index ddea40a8b1..32b5047fdd 100644 --- a/extension/src/cli/dvc/constants.ts +++ b/extension/src/cli/dvc/constants.ts @@ -43,6 +43,7 @@ export enum SubCommand { export enum Flag { ALL_COMMITS = '-A', + FOLLOW = '-f', FORCE = '-f', GRANULAR = '--granular', JOBS = '-j', @@ -67,6 +68,7 @@ export enum ExperimentSubCommand { export enum QueueSubCommand { KILL = 'kill', + LOGS = 'logs', START = 'start', STOP = 'stop' } diff --git a/extension/src/cli/dvc/runner.ts b/extension/src/cli/dvc/runner.ts index e3344c442b..25d99a94fc 100644 --- a/extension/src/cli/dvc/runner.ts +++ b/extension/src/cli/dvc/runner.ts @@ -98,7 +98,11 @@ export class DvcRunner extends Disposable implements ICli { ) this.pseudoTerminal = this.dispose.track( - new PseudoTerminal(this.processOutput, this.processTerminated) + new PseudoTerminal( + this.processOutput, + this.processTerminated, + 'DVC: exp run' + ) ) } diff --git a/extension/src/cli/dvc/viewer.ts b/extension/src/cli/dvc/viewer.ts new file mode 100644 index 0000000000..b475241911 --- /dev/null +++ b/extension/src/cli/dvc/viewer.ts @@ -0,0 +1,114 @@ +import { EventEmitter, Event } from 'vscode' +import { Args, Command, Flag, QueueSubCommand } from './constants' +import { getOptions } from './options' +import { CliResult, CliStarted, ICli, typeCheckCommands } from '..' +import { Config } from '../../config' +import { Disposable } from '../../class/dispose' +import { ViewableCliProcess } from '../viewable' + +export const autoRegisteredCommands = { + QUEUE_LOGS: 'queueLogs' +} as const + +export class DvcViewer extends Disposable implements ICli { + public readonly autoRegisteredCommands = typeCheckCommands( + autoRegisteredCommands, + this + ) + + public readonly processCompleted: EventEmitter + public readonly onDidCompleteProcess: Event + + public readonly processStarted: EventEmitter + public readonly onDidStartProcess: Event + + private processes: { + [id: string]: ViewableCliProcess + } + + private readonly config: Config + + constructor(config: Config) { + super() + + this.config = config + + this.processes = {} + + this.processCompleted = this.dispose.track(new EventEmitter()) + this.onDidCompleteProcess = this.processCompleted.event + + this.processStarted = this.dispose.track(new EventEmitter()) + this.onDidStartProcess = this.processStarted.event + } + + public run(name: string, cwd: string, ...args: Args) { + const viewableProcess = this.getRunningProcess(cwd, ...args) + if (viewableProcess) { + return viewableProcess.show() + } + + return this.createProcess(name, cwd, args) + } + + public queueLogs(cwd: string, expName: string) { + return this.run( + expName, + cwd, + Command.QUEUE, + QueueSubCommand.LOGS, + expName, + Flag.FOLLOW + ) + } + + private createProcess(name: string, cwd: string, args: Args) { + const viewableProcess = this.viewProcess(name, cwd, args) + + this.setRunningProcess(viewableProcess, cwd, ...args) + + const listener = this.dispose.track( + viewableProcess.onDidDispose(() => { + delete this.processes[this.getId(cwd, ...args)] + this.dispose.untrack(listener) + listener.dispose() + }) + ) + } + + private viewProcess(name: string, cwd: string, args: Args) { + return this.dispose.track( + new ViewableCliProcess( + `DVC: ${name}`, + this.getOptions(cwd, args), + this.processStarted, + this.processCompleted + ) + ) + } + + private getOptions(cwd: string, args: Args) { + return getOptions( + this.config.getPythonBinPath(), + this.config.getCliPath(), + cwd, + ...args + ) + } + + private getRunningProcess(cwd: string, ...args: Args) { + return this.processes[this.getId(cwd, ...args)] + } + + private setRunningProcess( + viewableProcess: ViewableCliProcess, + cwd: string, + ...args: Args + ) { + this.processes[this.getId(cwd, ...args)] = viewableProcess + } + + private getId(cwd: string, ...args: Args) { + return [cwd, ...args].join(':') + } +} diff --git a/extension/src/cli/viewable.ts b/extension/src/cli/viewable.ts index 662baefb9f..fc3d47291c 100644 --- a/extension/src/cli/viewable.ts +++ b/extension/src/cli/viewable.ts @@ -18,7 +18,7 @@ export class ViewableCliProcess extends DeferredDisposable { private readonly pseudoTerminal: PseudoTerminal constructor( - termName: string, + id: string, options: ProcessOptions, processStarted: EventEmitter, processCompleted: EventEmitter @@ -29,7 +29,7 @@ export class ViewableCliProcess extends DeferredDisposable { const onDidCloseTerminal = terminalClosed.event this.pseudoTerminal = this.dispose.track( - new PseudoTerminal(processOutput, terminalClosed, termName) + new PseudoTerminal(processOutput, terminalClosed, id) ) this.pseudoTerminal.setBlocked(true) diff --git a/extension/src/commands/internal.ts b/extension/src/commands/internal.ts index 9b5392025d..b788643fb1 100644 --- a/extension/src/commands/internal.ts +++ b/extension/src/commands/internal.ts @@ -2,9 +2,10 @@ import { commands } from 'vscode' import { RegisteredCliCommands, RegisteredCommands } from './external' import { ICli } from '../cli' import { Args } from '../cli/constants' -import { autoRegisteredCommands as CliExecutorCommands } from '../cli/dvc/executor' -import { autoRegisteredCommands as CliReaderCommands } from '../cli/dvc/reader' -import { autoRegisteredCommands as dvcRunnerCommands } from '../cli/dvc/runner' +import { autoRegisteredCommands as DvcExecutorCommands } from '../cli/dvc/executor' +import { autoRegisteredCommands as DvcReaderCommands } from '../cli/dvc/reader' +import { autoRegisteredCommands as DvcRunnerCommands } from '../cli/dvc/runner' +import { autoRegisteredCommands as DvcViewerCommands } from '../cli/dvc/viewer' import { autoRegisteredCommands as GitExecutorCommands } from '../cli/git/executor' import { autoRegisteredCommands as GitReaderCommands } from '../cli/git/reader' import { sendTelemetryEvent, sendTelemetryEventAndThrow } from '../telemetry' @@ -18,14 +19,16 @@ type Command = (...args: Args) => unknown | Promise export const AvailableCommands = Object.assign( { EXP_PUSH: 'expPush' } as const, - CliExecutorCommands, - CliReaderCommands, - dvcRunnerCommands, + DvcExecutorCommands, + DvcReaderCommands, + DvcRunnerCommands, + DvcViewerCommands, GitExecutorCommands, GitReaderCommands -) as typeof CliExecutorCommands & - typeof CliReaderCommands & - typeof dvcRunnerCommands & +) as typeof DvcExecutorCommands & + typeof DvcReaderCommands & + typeof DvcRunnerCommands & + typeof DvcViewerCommands & typeof GitExecutorCommands & typeof GitReaderCommands & { EXP_PUSH: 'expPush' } export type CommandId = diff --git a/extension/src/extension.ts b/extension/src/extension.ts index 945f908064..dd992ea190 100644 --- a/extension/src/extension.ts +++ b/extension/src/extension.ts @@ -55,6 +55,7 @@ import { Flag } from './cli/dvc/constants' import { LanguageClient } from './languageClient' import { collectRunningExperimentPids } from './experiments/processExecution/collect' import { registerPatchCommand } from './patch' +import { DvcViewer } from './cli/dvc/viewer' export class Extension extends Disposable { protected readonly internalCommands: InternalCommands @@ -68,6 +69,7 @@ export class Extension extends Disposable { private readonly dvcExecutor: DvcExecutor private readonly dvcReader: DvcReader private readonly dvcRunner: DvcRunner + private readonly dvcViewer: DvcViewer private readonly gitExecutor: GitExecutor private readonly gitReader: GitReader @@ -95,6 +97,7 @@ export class Extension extends Disposable { this.dvcExecutor = this.dispose.track(new DvcExecutor(config)) this.dvcReader = this.dispose.track(new DvcReader(config)) this.dvcRunner = this.dispose.track(new DvcRunner(config)) + this.dvcViewer = this.dispose.track(new DvcViewer(config)) this.gitExecutor = this.dispose.track(new GitExecutor()) this.gitReader = this.dispose.track(new GitReader()) @@ -103,6 +106,7 @@ export class Extension extends Disposable { this.dvcExecutor, this.dvcReader, this.dvcRunner, + this.dvcViewer, this.gitExecutor, this.gitReader ] diff --git a/extension/src/test/suite/cli/dvc/viewer.test.ts b/extension/src/test/suite/cli/dvc/viewer.test.ts new file mode 100644 index 0000000000..fc57d0210d --- /dev/null +++ b/extension/src/test/suite/cli/dvc/viewer.test.ts @@ -0,0 +1,53 @@ +import { afterEach, beforeEach, describe, it, suite } from 'mocha' +import { expect } from 'chai' +import { restore, spy } from 'sinon' +import { Disposable } from '../../../../extension' +import { Config } from '../../../../config' +import { DvcViewer } from '../../../../cli/dvc/viewer' +import { dvcDemoPath } from '../../../util' +import { ViewableCliProcess } from '../../../../cli/viewable' + +suite('DVC Viewer Test Suite', () => { + const disposable = Disposable.fn() + + beforeEach(() => { + restore() + }) + + afterEach(() => { + disposable.dispose() + }) + + describe('DvcViewer', () => { + it('should only be able to run a command once', async () => { + const mockConfig = { + getCliPath: () => 'sleep', + getPythonBinPath: () => undefined + } as Config + const dvcViewer = disposable.track(new DvcViewer(mockConfig)) + + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const getProcess = (dvcViewer: any, id: string): ViewableCliProcess => + dvcViewer.processes[id] + + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const createProcessSpy = spy(dvcViewer as any, 'createProcess') + + await dvcViewer.run('command', dvcDemoPath, '10000') + + expect(createProcessSpy).to.be.called + const viewableProcess = getProcess( + dvcViewer, + [dvcDemoPath, '10000'].join(':') + ) + const showProcessSpy = spy(viewableProcess, 'show') + + createProcessSpy.resetHistory() + + await dvcViewer.run('command', dvcDemoPath, '10000') + + expect(createProcessSpy).not.to.be.called + expect(showProcessSpy).to.be.calledOnce + }) + }) +}) diff --git a/extension/src/vscode/pseudoTerminal.ts b/extension/src/vscode/pseudoTerminal.ts index 10ae52547c..a0e12eabcd 100644 --- a/extension/src/vscode/pseudoTerminal.ts +++ b/extension/src/vscode/pseudoTerminal.ts @@ -17,7 +17,7 @@ export class PseudoTerminal extends Disposable { constructor( processOutput: EventEmitter, processTerminated: EventEmitter, - termName = 'DVC' + termName: string ) { super()