diff --git a/extension/src/cli/dvc/constants.ts b/extension/src/cli/dvc/constants.ts index 023aece588..2e107713ca 100644 --- a/extension/src/cli/dvc/constants.ts +++ b/extension/src/cli/dvc/constants.ts @@ -67,6 +67,7 @@ export enum ExperimentSubCommand { APPLY = 'apply', BRANCH = 'branch', GARBAGE_COLLECT = 'gc', + PUSH = 'push', REMOVE = 'remove', RUN = 'run' } diff --git a/extension/src/cli/dvc/executor.test.ts b/extension/src/cli/dvc/executor.test.ts index a20b0dbcdf..642013876f 100644 --- a/extension/src/cli/dvc/executor.test.ts +++ b/extension/src/cli/dvc/executor.test.ts @@ -306,6 +306,24 @@ describe('CliExecutor', () => { }) }) + describe('experimentPush', () => { + it('should call createProcess with the correct parameters to push an existing experiment to the remote', async () => { + const cwd = __dirname + const stdout = '' + mockedCreateProcess.mockReturnValueOnce(getMockedProcess(stdout)) + + const output = await dvcExecutor.experimentPush(cwd, 'toric-sail') + expect(output).toStrictEqual(stdout) + + expect(mockedCreateProcess).toHaveBeenCalledWith({ + args: ['exp', 'push', 'origin', 'toric-sail'], + cwd, + env: mockedEnv, + executable: 'dvc' + }) + }) + }) + describe('experimentRemove', () => { it('should call createProcess with the correct parameters to remove an existing experiment from the workspace', async () => { const cwd = __dirname diff --git a/extension/src/cli/dvc/executor.ts b/extension/src/cli/dvc/executor.ts index d551dff589..ecc6c06213 100644 --- a/extension/src/cli/dvc/executor.ts +++ b/extension/src/cli/dvc/executor.ts @@ -13,6 +13,7 @@ import { addStudioAccessToken } from './options' import { CliResult, CliStarted, typeCheckCommands } from '..' import { ContextKey, setContextValue } from '../../vscode/context' import { Config } from '../../config' +import { DEFAULT_REMOTE } from '../git/constants' export const autoRegisteredCommands = { ADD: 'add', @@ -22,6 +23,7 @@ export const autoRegisteredCommands = { EXPERIMENT_APPLY: 'experimentApply', EXPERIMENT_BRANCH: 'experimentBranch', EXPERIMENT_GARBAGE_COLLECT: 'experimentGarbageCollect', + EXPERIMENT_PUSH: 'experimentPush', EXPERIMENT_QUEUE: 'experimentRunQueue', EXPERIMENT_REMOVE: 'experimentRemove', EXPERIMENT_REMOVE_QUEUE: 'experimentRemoveQueue', @@ -109,6 +111,15 @@ export class DvcExecutor extends DvcCli { ) } + public experimentPush(cwd: string, id: string) { + return this.executeExperimentProcess( + cwd, + ExperimentSubCommand.PUSH, + DEFAULT_REMOTE, + id + ) + } + public experimentRemove(cwd: string, ...experimentNames: string[]) { return this.executeExperimentProcess( cwd, diff --git a/extension/src/commands/external.ts b/extension/src/commands/external.ts index 2f2dbd4ed9..9ab290075f 100644 --- a/extension/src/commands/external.ts +++ b/extension/src/commands/external.ts @@ -16,6 +16,7 @@ export enum RegisteredCliCommands { EXPERIMENT_VIEW_APPLY = 'dvc.views.experiments.applyExperiment', EXPERIMENT_VIEW_BRANCH = 'dvc.views.experiments.branchExperiment', + EXPERIMENT_VIEW_PUSH = 'dvc.views.experiments.pushExperiment', EXPERIMENT_VIEW_REMOVE = 'dvc.views.experiments.removeExperiment', EXPERIMENT_VIEW_SHARE_AS_BRANCH = 'dvc.views.experiments.shareExperimentAsBranch', EXPERIMENT_VIEW_SHARE_AS_COMMIT = 'dvc.views.experiments.shareExperimentAsCommit', diff --git a/extension/src/experiments/commands/index.ts b/extension/src/experiments/commands/index.ts index 52906fee2f..972fb0bc4d 100644 --- a/extension/src/experiments/commands/index.ts +++ b/extension/src/experiments/commands/index.ts @@ -115,10 +115,22 @@ export const getShareExperimentToStudioCommand = return commands.executeCommand(RegisteredCommands.SETUP_SHOW) } - return internalCommands.executeCommand( - AvailableCommands.EXP_PUSH, - studioAccessToken, - dvcRoot, - id - ) + return Toast.showProgress('Pushing Experiment', async progress => { + progress.report({ increment: 0 }) + + progress.report({ increment: 25, message: 'Pushing experiment...' }) + + await Toast.runCommandAndIncrementProgress( + () => + internalCommands.executeCommand( + AvailableCommands.EXPERIMENT_PUSH, + dvcRoot, + id + ), + progress, + 75 + ) + + return Toast.delayProgressClosing() + }) } diff --git a/extension/src/extension.ts b/extension/src/extension.ts index 9fd8799e12..3d2ad86491 100644 --- a/extension/src/extension.ts +++ b/extension/src/extension.ts @@ -52,7 +52,6 @@ import { stopProcesses } from './process/execution' 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' import { registerSetupCommands } from './setup/register' import { Status } from './status' @@ -200,7 +199,6 @@ export class Extension extends Disposable { ) ) - registerPatchCommand(this.internalCommands) registerExperimentCommands( this.experiments, this.internalCommands, diff --git a/extension/src/patch.ts b/extension/src/patch.ts deleted file mode 100644 index bf83ecdb36..0000000000 --- a/extension/src/patch.ts +++ /dev/null @@ -1,209 +0,0 @@ -import { commands } from 'vscode' -import fetch, { Response as FetchResponse } from 'node-fetch' -import { - ExpData, - ExpRange, - ExpShowOutput, - ExpState, - ValueTree, - MetricsOrParams, - experimentHasError, - fileHasError -} from './cli/dvc/contract' -import { AvailableCommands, InternalCommands } from './commands/internal' -import { Args, ExperimentFlag } from './cli/dvc/constants' -import { Response as UserResponse } from './vscode/response' -import { Toast } from './vscode/toast' -import { Modal } from './vscode/modal' -import { RegisteredCommands } from './commands/external' - -export const STUDIO_ENDPOINT = 'https://studio.iterative.ai/api/live' - -type ExperimentDetails = { - baselineSha: string - sha: string - metrics: MetricsOrParams | undefined | null - name: string - params: ValueTree | null -} - -type RequestBody = { - client: 'vscode' - repo_url: string - name: string - baseline_sha: string - experiment_rev: string - metrics: MetricsOrParams - params: ValueTree - type: 'done' -} - -const collectExperiment = (data: ExpData) => { - const { metrics, params } = data - - const paramsAcc: ValueTree = {} - for (const [file, fileParams] of Object.entries(params || {})) { - if (fileHasError(fileParams)) { - continue - } - const data = fileParams?.data - if (data) { - paramsAcc[file] = data - } - } - - return { metrics, params: paramsAcc } -} - -const findExperimentByName = ( - name: string, - expState: ExpState & { experiments?: ExpRange[] | null } - // eslint-disable-next-line sonarjs/cognitive-complexity -) => { - if (!expState.experiments) { - return - } - - for (const experiment of expState.experiments) { - if (experiment.name !== name) { - continue - } - - const [exp] = experiment.revs - if (experimentHasError(exp)) { - return - } - - const { data, rev } = exp - - if (data) { - const { metrics, params } = collectExperiment(data) - - return { - baselineSha: expState.rev, - metrics, - name, - params, - sha: rev - } - } - } -} - -const collectExperimentDetails = ( - name: string, - expData: ExpShowOutput -): ExperimentDetails | undefined => { - for (const expState of expData.slice(1)) { - const details = findExperimentByName(name, expState) - if (details) { - return details - } - } -} - -const sendPostRequest = ( - studioAccessToken: string, - body: RequestBody -): Promise => - fetch(STUDIO_ENDPOINT, { - body: JSON.stringify(body), - headers: { - Authorization: `token ${studioAccessToken}`, - 'Content-type': 'application/json' - }, - method: 'POST' - }) - -const showUnauthorized = async () => { - const response = await Modal.errorWithOptions( - 'The current Studio access token is invalid. Please add a new token.', - UserResponse.SHOW - ) - if (response === UserResponse.SHOW) { - return commands.executeCommand(RegisteredCommands.SETUP_SHOW) - } -} - -const shareWithProgress = ( - experimentDetails: ExperimentDetails, - repoUrl: string, - studioAccessToken: string -): Thenable => - Toast.showProgress('Sharing Experiment', async progress => { - const { metrics, params, baselineSha, sha, name } = experimentDetails - - progress.report({ - increment: 0, - message: 'Initializing experiment...' - }) - - progress.report({ increment: 50, message: 'Sending data...' }) - const response = await sendPostRequest(studioAccessToken, { - baseline_sha: baselineSha, - client: 'vscode', - experiment_rev: sha, - metrics: metrics || {}, - name, - params: params || {}, - repo_url: repoUrl, - type: 'done' - }) - - progress.report({ increment: 25, message: 'Response received...' }) - - if (response.status === 401) { - progress.report({ increment: 25, message: 'Access unauthorized' }) - void showUnauthorized() - return Toast.delayProgressClosing() - } - - progress.report({ increment: 25, message: 'Done' }) - - return Toast.delayProgressClosing() - }) - -const pushExperiment = async ( - internalCommands: InternalCommands, - dvcRoot: string, - name: string, - studioAccessToken: string -) => { - const [repoUrl, expShowOutput] = await Promise.all([ - internalCommands.executeCommand( - AvailableCommands.GIT_GET_REMOTE_URL, - dvcRoot - ), - internalCommands.executeCommand( - AvailableCommands.EXP_SHOW, - dvcRoot, - ExperimentFlag.NO_FETCH - ) - ]) - - const experimentDetails = collectExperimentDetails(name, expShowOutput) - - if (!repoUrl) { - return Toast.showError( - 'Failed to share experiment, unable to generate Git repo URL' - ) - } - - if (!experimentDetails) { - return Toast.showError( - 'Failed to share experiment, unable to locate the required data' - ) - } - - return shareWithProgress(experimentDetails, repoUrl, studioAccessToken) -} - -export const registerPatchCommand = (internalCommands: InternalCommands) => - internalCommands.registerCommand( - AvailableCommands.EXP_PUSH, - (...args: Args) => { - const [studioAccessToken, dvcRoot, name] = args - - return pushExperiment(internalCommands, dvcRoot, name, studioAccessToken) - } - ) diff --git a/extension/src/telemetry/constants.ts b/extension/src/telemetry/constants.ts index 8b63d74e2d..02dddc9ea9 100644 --- a/extension/src/telemetry/constants.ts +++ b/extension/src/telemetry/constants.ts @@ -144,6 +144,7 @@ export interface IEventNamePropertyMapping { [EventName.EXPERIMENT_TOGGLE]: undefined [EventName.EXPERIMENT_VIEW_APPLY]: undefined [EventName.EXPERIMENT_VIEW_BRANCH]: undefined + [EventName.EXPERIMENT_VIEW_PUSH]: undefined [EventName.EXPERIMENT_VIEW_REMOVE]: undefined [EventName.EXPERIMENT_VIEW_SHARE_AS_BRANCH]: undefined [EventName.EXPERIMENT_VIEW_SHARE_AS_COMMIT]: undefined diff --git a/extension/src/test/suite/experiments/index.test.ts b/extension/src/test/suite/experiments/index.test.ts index 0cab2cc389..955ed2355b 100644 --- a/extension/src/test/suite/experiments/index.test.ts +++ b/extension/src/test/suite/experiments/index.test.ts @@ -650,7 +650,7 @@ suite('Experiments Test Suite', () => { 'getStudioAccessToken' ) - const tokenAccessed = new Promise(resolve => + const tokenNotFound = new Promise(resolve => mockGetStudioAccessToken.callsFake(() => { resolve(undefined) return '' @@ -662,11 +662,38 @@ suite('Experiments Test Suite', () => { type: MessageFromWebviewType.SHARE_EXPERIMENT_TO_STUDIO }) - await tokenAccessed + await tokenNotFound expect(executeCommandSpy).to.be.calledWithExactly( RegisteredCommands.SETUP_SHOW ) + + mockGetStudioAccessToken.resetBehavior() + + const tokenFound = new Promise(resolve => + mockGetStudioAccessToken.callsFake(() => { + resolve(undefined) + return 'isat_token' + }) + ) + const mockExperimentPush = stub(DvcExecutor.prototype, 'experimentPush') + const commandExecuted = new Promise(resolve => + mockExperimentPush.callsFake(() => { + resolve(undefined) + return Promise.resolve( + `Pushed experiment ${mockExpId} to Git remote 'origin'` + ) + }) + ) + + mockMessageReceived.fire({ + payload: mockExpId, + type: MessageFromWebviewType.SHARE_EXPERIMENT_TO_STUDIO + }) + + await Promise.all([tokenFound, commandExecuted]) + + expect(mockExperimentPush).to.be.calledWithExactly(dvcDemoPath, mockExpId) }).timeout(WEBVIEW_TEST_TIMEOUT) it('should handle a message to share an experiment as a new branch', async () => { diff --git a/extension/src/test/suite/patch.test.ts b/extension/src/test/suite/patch.test.ts deleted file mode 100644 index cb17ab728a..0000000000 --- a/extension/src/test/suite/patch.test.ts +++ /dev/null @@ -1,220 +0,0 @@ -import { join } from 'path' -import { afterEach, beforeEach, describe, it, suite } from 'mocha' -import { restore, spy, stub } from 'sinon' -import { expect } from 'chai' -import * as Fetch from 'node-fetch' -import { commands } from 'vscode' -import { - buildInternalCommands, - bypassProgressCloseDelay, - closeAllEditors -} from './util' -import { Disposable } from '../../extension' -import { STUDIO_ENDPOINT, registerPatchCommand } from '../../patch' -import { AvailableCommands } from '../../commands/internal' -import expShowFixture from '../fixtures/expShow/base/output' -import { dvcDemoPath } from '../util' -import { Toast } from '../../vscode/toast' -import { Modal } from '../../vscode/modal' -import { Response } from '../../vscode/response' -import { RegisteredCommands } from '../../commands/external' -import { ExpData } from '../../cli/dvc/contract' - -suite('Patch Test Suite', () => { - const disposable = Disposable.fn() - - beforeEach(() => { - restore() - }) - - afterEach(function () { - disposable.dispose() - return closeAllEditors() - }) - - describe('exp push patch', () => { - it('should share an experiment to Studio', async () => { - bypassProgressCloseDelay() - const mockFetch = stub(Fetch, 'default').resolves({} as Fetch.Response) - const mockStudioAccessToken = 'isat_12123123123123123' - const mockRepoUrl = 'https://github.com/iterative/vscode-dvc-demo' - - const id = 'exp-e7a67' - - const { internalCommands, gitReader, dvcReader } = - buildInternalCommands(disposable) - - const mockGetRemoteUrl = stub(gitReader, 'getRemoteUrl').resolves( - mockRepoUrl - ) - const mockExpShow = stub(dvcReader, 'expShow').resolves(expShowFixture) - - registerPatchCommand(internalCommands) - - await internalCommands.executeCommand( - AvailableCommands.EXP_PUSH, - mockStudioAccessToken, - dvcDemoPath, - id - ) - - expect(mockGetRemoteUrl).to.be.calledOnce - expect(mockExpShow).to.be.calledOnce - expect(mockFetch).to.be.calledOnce - - const { metrics, params } = ( - expShowFixture - .find(({ rev }) => rev === '53c3851f46955fa3e2b8f6e1c52999acc8c9ea77') - ?.experiments?.find( - ({ revs }) => - revs[0].rev === '4fb124aebddb2adf1545030907687fa9a4c80e70' - )?.revs?.[0] as { data: ExpData } - ).data - - const body = mockFetch.lastCall.args[1]?.body - - expect(JSON.parse((body as string) || '{}')).to.deep.equal({ - baseline_sha: '53c3851f46955fa3e2b8f6e1c52999acc8c9ea77', - client: 'vscode', - experiment_rev: '4fb124aebddb2adf1545030907687fa9a4c80e70', - metrics, - name: id, - params: { - 'params.yaml': (params?.['params.yaml'] as { data: unknown })?.data, - [join('nested', 'params.yaml')]: ( - params?.[join('nested', 'params.yaml')] as { data: unknown } - )?.data - }, - repo_url: mockRepoUrl, - type: 'done' - }) - - expect(mockFetch).to.be.calledWithExactly(STUDIO_ENDPOINT, { - body: JSON.stringify({ - baseline_sha: '53c3851f46955fa3e2b8f6e1c52999acc8c9ea77', - client: 'vscode', - experiment_rev: '4fb124aebddb2adf1545030907687fa9a4c80e70', - metrics, - name: id, - params: { - 'params.yaml': (params?.['params.yaml'] as { data: unknown })?.data, - [join('nested', 'params.yaml')]: ( - params?.[join('nested', 'params.yaml')] as { data: unknown } - )?.data - }, - repo_url: mockRepoUrl, - type: 'done' - }), - headers: { - Authorization: `token ${mockStudioAccessToken}`, - 'Content-type': 'application/json' - }, - method: 'POST' - }) - }) - - it('should show an error modal if Studio returns a 401 response', async () => { - bypassProgressCloseDelay() - const mockFetch = stub(Fetch, 'default').resolves({ - status: 401 - } as Fetch.Response) - const mockStudioAccessToken = 'isat_12123123123123123' - const mockRepoUrl = 'https://github.com/iterative/vscode-dvc-demo' - - const executeCommandSpy = spy(commands, 'executeCommand') - - const { internalCommands, gitReader, dvcReader } = - buildInternalCommands(disposable) - - const mockGetRemoteUrl = stub(gitReader, 'getRemoteUrl').resolves( - mockRepoUrl - ) - const mockExpShow = stub(dvcReader, 'expShow').resolves(expShowFixture) - - const mockErrorWithOptions = stub(Modal, 'errorWithOptions').resolves( - Response.SHOW - ) - - const id = 'exp-e7a67' - - registerPatchCommand(internalCommands) - - await internalCommands.executeCommand( - AvailableCommands.EXP_PUSH, - mockStudioAccessToken, - dvcDemoPath, - id - ) - - expect(mockGetRemoteUrl).to.be.calledOnce - expect(mockExpShow).to.be.calledOnce - expect(mockFetch).to.be.calledOnce - expect(mockErrorWithOptions).to.be.calledOnce - expect(executeCommandSpy).to.be.calledWithExactly( - RegisteredCommands.SETUP_SHOW - ) - - const { metrics, params } = ( - expShowFixture - .find(({ rev }) => rev === '53c3851f46955fa3e2b8f6e1c52999acc8c9ea77') - ?.experiments?.find( - ({ revs }) => - revs[0].rev === '4fb124aebddb2adf1545030907687fa9a4c80e70' - )?.revs?.[0] as { data: ExpData } - ).data - - expect(mockFetch).to.be.calledWithExactly(STUDIO_ENDPOINT, { - body: JSON.stringify({ - baseline_sha: '53c3851f46955fa3e2b8f6e1c52999acc8c9ea77', - client: 'vscode', - experiment_rev: '4fb124aebddb2adf1545030907687fa9a4c80e70', - metrics, - name: id, - params: { - 'params.yaml': (params?.['params.yaml'] as { data: unknown })?.data, - [join('nested', 'params.yaml')]: ( - params?.[join('nested', 'params.yaml')] as { data: unknown } - )?.data - }, - repo_url: mockRepoUrl, - type: 'done' - }), - headers: { - Authorization: `token ${mockStudioAccessToken}`, - 'Content-type': 'application/json' - }, - method: 'POST' - }) - }) - - it('should show an error message if the experiment cannot be found', async () => { - const mockShowError = stub(Toast, 'showError').resolves(undefined) - - const mockStudioAccessToken = 'isat_not-needed' - const mockRepoUrl = 'https://github.com/iterative/vscode-dvc-demo' - - const { internalCommands, gitReader, dvcReader } = - buildInternalCommands(disposable) - - const mockGetRemoteUrl = stub(gitReader, 'getRemoteUrl').resolves( - mockRepoUrl - ) - const mockExpShow = stub(dvcReader, 'expShow').resolves(expShowFixture) - - registerPatchCommand(internalCommands) - - await internalCommands.executeCommand( - AvailableCommands.EXP_PUSH, - mockStudioAccessToken, - dvcDemoPath, - 'not-an-experiment' - ) - - expect(mockGetRemoteUrl).to.be.calledOnce - expect(mockGetRemoteUrl).to.be.calledWith(dvcDemoPath) - expect(mockExpShow).to.be.calledOnce - expect(mockExpShow).to.be.calledWithExactly(dvcDemoPath, '--no-fetch') - expect(mockShowError).to.be.calledOnce - }) - }) -})