Merge branch 'main' into add-connect-to-studio

extension/src/experiments/index.ts

@@ -100,12 +100,15 @@ export class Experiments extends BaseRepository<TableData> {
   private dvcLiveOnlyCleanupInitialized = false
   private dvcLiveOnlySignalFile: string
 
+  private readonly addStage: () => Promise<boolean>
+
   constructor(
     dvcRoot: string,
     internalCommands: InternalCommands,
     updatesPaused: EventEmitter<boolean>,
     resourceLocator: ResourceLocator,
     workspaceState: Memento,
+    addStage: () => Promise<boolean>,
     cliData?: ExperimentsData,
     fileSystemData?: FileSystemData
   ) {
@@ -117,6 +120,7 @@ export class Experiments extends BaseRepository<TableData> {
     )
 
     this.internalCommands = internalCommands
+    this.addStage = addStage
 
     this.onDidChangeIsParamsFileFocused = this.paramsFileFocused.event
     this.onDidChangeExperiments = this.experimentsChanged.event
@@ -571,7 +575,13 @@ export class Experiments extends BaseRepository<TableData> {
           AvailableCommands.QUEUE_KILL,
           dvcRoot,
           ...ids
-        )
+        ),
+      () =>
+        this.internalCommands.executeCommand(
+          AvailableCommands.STAGE_LIST,
+          this.dvcRoot
+        ),
+      () => this.addStage()
     )
 
     this.dispose.track(

extension/src/experiments/webview/contract.ts

@@ -98,6 +98,7 @@ export type TableData = {
   columnWidths: Record<string, number>
   hasCheckpoints: boolean
   hasColumns: boolean
+  hasConfig: boolean
   hasRunningExperiment: boolean
   rows: Row[]
   sorts: SortDefinition[]

extension/src/experiments/webview/messages.ts

@@ -40,6 +40,12 @@ export class WebviewMessages {
     ...ids: string[]
   ) => Promise<string | undefined>
 
+  private readonly hasStages: () => Promise<string>
+
+  private hasConfig = false
+
+  private readonly addStage: () => Promise<boolean>
+
   constructor(
     dvcRoot: string,
     experiments: ExperimentsModel,
@@ -51,7 +57,9 @@ export class WebviewMessages {
     stopQueuedExperiments: (
       dvcRoot: string,
       ...ids: string[]
-    ) => Promise<string | undefined>
+    ) => Promise<string | undefined>,
+    hasStages: () => Promise<string>,
+    addStage: () => Promise<boolean>
   ) {
     this.dvcRoot = dvcRoot
     this.experiments = experiments
@@ -61,6 +69,13 @@ export class WebviewMessages {
     this.notifyChanged = notifyChanged
     this.selectColumns = selectColumns
     this.stopQueuedExperiments = stopQueuedExperiments
+    this.hasStages = hasStages
+    void this.changeHasConfig()
+    this.addStage = addStage
+  }
+
+  public async changeHasConfig() {
+    this.hasConfig = !!(await this.hasStages())
   }
 
   public sendWebviewMessage() {
@@ -166,6 +181,10 @@ export class WebviewMessages {
         return this.stopExperiments(message.payload)
       }
 
+      case MessageFromWebviewType.ADD_CONFIGURATION: {
+        return this.addConfiguration()
+      }
+
       default:
         Logger.error(`Unexpected message: ${JSON.stringify(message)}`)
     }
@@ -183,12 +202,18 @@ export class WebviewMessages {
       filters: this.experiments.getFilterPaths(),
       hasCheckpoints: this.checkpoints.hasCheckpoints(),
       hasColumns: this.columns.hasNonDefaultColumns(),
+      hasConfig: this.hasConfig,
       hasRunningExperiment: this.experiments.hasRunningExperiment(),
       rows: this.experiments.getRowData(),
       sorts: this.experiments.getSorts()
     }
   }
 
+  private async addConfiguration() {
+    await this.addStage()
+    await this.changeHasConfig()
+  }
+
   private async setMaxTableHeadDepth() {
     const newValue = await getPositiveIntegerInput(
       Title.SET_EXPERIMENTS_HEADER_HEIGHT,

extension/src/experiments/workspace.ts

@@ -364,7 +364,8 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews<
         this.internalCommands,
         updatesPaused,
         resourceLocator,
-        this.workspaceState
+        this.workspaceState,
+        () => this.checkOrAddPipeline(dvcRoot)
       )
     )
 

extension/src/test/fixtures/expShow/base/tableData.ts

@@ -8,6 +8,7 @@ const tableDataFixture: TableData = {
   columns: columnsFixture,
   filters: [],
   hasCheckpoints: true,
+  hasConfig: true,
   hasRunningExperiment: true,
   hasColumns: true,
   sorts: [],

extension/src/test/fixtures/expShow/dataTypes/tableData.ts

@@ -9,6 +9,7 @@ export const data: TableData = {
   filteredCounts: { experiments: 0, checkpoints: 0 },
   filters: [],
   hasCheckpoints: false,
+  hasConfig: true,
   hasRunningExperiment: false,
   sorts: [],
   columns,

extension/src/test/fixtures/expShow/deeplyNested/tableData.ts

@@ -12,6 +12,7 @@ const data: TableData = {
     'params:params.yaml:nested1%2Enested2%2Enested3.nested4.nested5b.doubled'
   ],
   hasCheckpoints: false,
+  hasConfig: true,
   hasRunningExperiment: false,
   sorts: [
     {

extension/src/test/fixtures/expShow/survival/tableData.ts

@@ -8,6 +8,7 @@ const data: TableData = {
   columns: columnsFixture,
   filters: [],
   hasCheckpoints: true,
+  hasConfig: true,
   hasRunningExperiment: true,
   hasColumns: true,
   sorts: [],

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

@@ -76,6 +76,7 @@ import { AvailableCommands } from '../../../commands/internal'
 import { Setup } from '../../../setup'
 import * as FileSystem from '../../../fileSystem'
 import * as ProcessExecution from '../../../processExecution'
+import { DvcReader } from '../../../cli/dvc/reader'
 
 const { openFileInEditor } = FileSystem
 
@@ -127,6 +128,8 @@ suite('Experiments Test Suite', () => {
 
   describe('showWebview', () => {
     it('should be able to make the experiment webview visible', async () => {
+      stub(DvcReader.prototype, 'listStages').resolves('train')
+
       const { experiments, messageSpy } = buildExperiments(
         disposable,
         expShowFixture
@@ -143,6 +146,7 @@ suite('Experiments Test Suite', () => {
         filters: [],
         hasCheckpoints: true,
         hasColumns: true,
+        hasConfig: true,
         hasRunningExperiment: true,
         rows: rowsFixture,
         sorts: []
@@ -177,6 +181,62 @@ suite('Experiments Test Suite', () => {
 
       expect(windowSpy).not.to.have.been.called
     }).timeout(WEBVIEW_TEST_TIMEOUT)
+
+    it('should set hasConfig to false if there are no stages', async () => {
+      stub(DvcReader.prototype, 'listStages').resolves('')
+
+      const { experiments, messageSpy } = buildExperiments(
+        disposable,
+        expShowFixture
+      )
+
+      await experiments.showWebview()
+
+      const expectedTableData: TableData = {
+        changes: workspaceChangesFixture,
+        columnOrder: columnsOrderFixture,
+        columnWidths: {},
+        columns: columnsFixture,
+        filteredCounts: { checkpoints: 0, experiments: 0 },
+        filters: [],
+        hasCheckpoints: true,
+        hasColumns: true,
+        hasConfig: false,
+        hasRunningExperiment: true,
+        rows: rowsFixture,
+        sorts: []
+      }
+
+      expect(messageSpy).to.be.calledWithExactly(expectedTableData)
+    }).timeout(WEBVIEW_TEST_TIMEOUT)
+
+    it('should set hasConfig to true if there are stages', async () => {
+      stub(DvcReader.prototype, 'listStages').resolves('train')
+
+      const { experiments, messageSpy } = buildExperiments(
+        disposable,
+        expShowFixture
+      )
+
+      await experiments.showWebview()
+
+      const expectedTableData: TableData = {
+        changes: workspaceChangesFixture,
+        columnOrder: columnsOrderFixture,
+        columnWidths: {},
+        columns: columnsFixture,
+        filteredCounts: { checkpoints: 0, experiments: 0 },
+        filters: [],
+        hasCheckpoints: true,
+        hasColumns: true,
+        hasConfig: true,
+        hasRunningExperiment: true,
+        rows: rowsFixture,
+        sorts: []
+      }
+
+      expect(messageSpy).to.be.calledWithExactly(expectedTableData)
+    }).timeout(WEBVIEW_TEST_TIMEOUT)
   })
 
   describe('handleMessageFromWebview', () => {
@@ -193,6 +253,7 @@ suite('Experiments Test Suite', () => {
         experimentsModel,
         internalCommands,
         dvcExecutor,
+        mockCheckOrAddPipeline,
         messageSpy
       } = buildExperiments(disposable, expShowFixture)
       const mockExecuteCommand = stub(
@@ -209,6 +270,7 @@ suite('Experiments Test Suite', () => {
         experiments,
         experimentsModel,
         messageSpy,
+        mockCheckOrAddPipeline,
         mockExecuteCommand
       }
     }
@@ -788,6 +850,8 @@ suite('Experiments Test Suite', () => {
     }).timeout(WEBVIEW_TEST_TIMEOUT)
 
     it('should be able to handle a message to select columns', async () => {
+      stub(DvcReader.prototype, 'listStages').resolves('train')
+
       const { columnsModel, experiments, messageSpy } =
         setupExperimentsAndMockCommands()
 
@@ -826,6 +890,7 @@ suite('Experiments Test Suite', () => {
         filters: [],
         hasCheckpoints: true,
         hasColumns: true,
+        hasConfig: true,
         hasRunningExperiment: true,
         rows: rowsFixture,
         sorts: []
@@ -1120,6 +1185,23 @@ suite('Experiments Test Suite', () => {
       expect(mockProcessExists).to.be.calledWithExactly(mockPid)
       expect(mockStopProcesses).to.be.calledWithExactly([mockPid])
     }).timeout(WEBVIEW_TEST_TIMEOUT)
+
+    it('should handle a message to add a configuration', async () => {
+      stub(DvcReader.prototype, 'listStages').resolves('')
+
+      const { experiments, mockCheckOrAddPipeline, messageSpy } =
+        setupExperimentsAndMockCommands()
+
+      const webview = await experiments.showWebview()
+      messageSpy.resetHistory()
+      const mockMessageReceived = getMessageReceivedEmitter(webview)
+
+      mockMessageReceived.fire({
+        type: MessageFromWebviewType.ADD_CONFIGURATION
+      })
+
+      expect(mockCheckOrAddPipeline).to.be.calledOnce
+    }).timeout(WEBVIEW_TEST_TIMEOUT)
   })
 
   describe('Sorting', () => {
@@ -1148,6 +1230,7 @@ suite('Experiments Test Suite', () => {
           updatesPaused,
           resourceLocator,
           buildMockMemento(),
+          () => Promise.resolve(true),
           buildMockData<ExperimentsData>(),
           buildMockData<FileSystemData>()
         )
@@ -1364,6 +1447,7 @@ suite('Experiments Test Suite', () => {
           {} as EventEmitter<boolean>,
           {} as ResourceLocator,
           mockMemento,
+          () => Promise.resolve(true),
           buildMockData<ExperimentsData>(),
           buildMockData<FileSystemData>()
         )
@@ -1541,6 +1625,7 @@ suite('Experiments Test Suite', () => {
           {} as EventEmitter<boolean>,
           {} as ResourceLocator,
           mockMemento,
+          () => Promise.resolve(true),
           buildMockData<ExperimentsData>(),
           buildMockData<FileSystemData>()
         )

extension/src/test/suite/experiments/model/filterBy/tree.test.ts

@@ -41,6 +41,7 @@ import {
   FilterItem
 } from '../../../../../experiments/model/filterBy/tree'
 import { starredFilter } from '../../../../../experiments/model/filterBy/constants'
+import { DvcReader } from '../../../../../cli/dvc/reader'
 
 suite('Experiments Filter By Tree Test Suite', () => {
   const disposable = Disposable.fn()
@@ -62,6 +63,8 @@ suite('Experiments Filter By Tree Test Suite', () => {
     })
 
     it('should be able to update the table data by adding and removing a filter', async () => {
+      stub(DvcReader.prototype, 'listStages').resolves('train')
+
       const { experiments, messageSpy } = buildExperiments(disposable)
 
       await experiments.isReady()
@@ -124,6 +127,7 @@ suite('Experiments Filter By Tree Test Suite', () => {
         filters: [accuracyPath],
         hasCheckpoints: true,
         hasColumns: true,
+        hasConfig: true,
         hasRunningExperiment: true,
         rows: filteredRows,
         sorts: []
@@ -155,6 +159,7 @@ suite('Experiments Filter By Tree Test Suite', () => {
         filters: [],
         hasCheckpoints: true,
         hasColumns: true,
+        hasConfig: true,
         hasRunningExperiment: true,
         rows: [workspace, main],
         sorts: []
@@ -457,6 +462,7 @@ suite('Experiments Filter By Tree Test Suite', () => {
     })
 
     it('should be able to filter to starred experiments', async () => {
+      stub(DvcReader.prototype, 'listStages').resolves('train')
       const { experiments, messageSpy } = buildExperiments(disposable)
 
       await experiments.isReady()
@@ -485,6 +491,7 @@ suite('Experiments Filter By Tree Test Suite', () => {
         filters: ['starred'],
         hasCheckpoints: true,
         hasColumns: true,
+        hasConfig: true,
         hasRunningExperiment: true,
         rows: filteredRows,
         sorts: []

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

@@ -65,6 +65,7 @@ export const buildExperiments = (
   const mockExperimentsData = buildMockData<ExperimentsData>(
     mockUpdateExperimentsData
   )
+  const mockCheckOrAddPipeline = stub()
 
   const experiments = disposer.track(
     new Experiments(
@@ -73,6 +74,7 @@ export const buildExperiments = (
       updatesPaused,
       resourceLocator,
       buildMockMemento(),
+      mockCheckOrAddPipeline,
       mockExperimentsData,
       buildMockData<FileSystemData>()
     )
@@ -94,6 +96,7 @@ export const buildExperiments = (
     gitReader,
     internalCommands,
     messageSpy,
+    mockCheckOrAddPipeline,
     mockCheckSignalFile,
     mockExperimentShow,
     mockGetCommitMessages,