Handle non-standard experiment pipeline configurations

extension/src/cli/dvc/reader.ts

@@ -38,7 +38,7 @@ export const autoRegisteredCommands = {
   GLOBAL_VERSION: 'globalVersion',
   PLOTS_DIFF: 'plotsDiff',
   ROOT: 'root',
-  STAGE_LIST: 'listStages',
+  STAGE_LIST: 'stageList',
   VERSION: 'version'
 } as const
 
@@ -144,7 +144,7 @@ export class DvcReader extends DvcCli {
     return this.executeProcess(options)
   }
 
-  public async listStages(cwd: string): Promise<string | undefined> {
+  public async stageList(cwd: string): Promise<string | undefined> {
     try {
       return await this.executeDvcProcess(cwd, Command.STAGE, SubCommand.LIST)
     } catch {}

extension/src/data/index.ts

@@ -20,10 +20,9 @@ export abstract class BaseData<
   T extends
     | { data: PlotsOutputOrError; revs: string[] }
     | ExperimentsOutput
-    | string
+    | { dag: string; stages: { [pipeline: string]: string | undefined } }
 > extends DeferredDisposable {
   public readonly onDidUpdate: Event<T>
-  public readonly onDidChangeDvcYaml: Event<void>
 
   protected readonly dvcRoot: string
   protected readonly processManager: ProcessManager
@@ -36,10 +35,6 @@ export abstract class BaseData<
     new EventEmitter()
   )
 
-  private readonly dvcYamlChanged: EventEmitter<void> = this.dispose.track(
-    new EventEmitter<void>()
-  )
-
   constructor(
     dvcRoot: string,
     internalCommands: InternalCommands,
@@ -56,8 +51,6 @@ export abstract class BaseData<
     this.onDidUpdate = this.updated.event
     this.staticFiles = staticFiles
 
-    this.onDidChangeDvcYaml = this.dvcYamlChanged.event
-
     this.watchFiles()
 
     this.waitForInitialData()
@@ -96,13 +89,9 @@ export abstract class BaseData<
         ) {
           void this.managedUpdate(path)
         }
-
-        if (path.endsWith('dvc.yaml')) {
-          this.dvcYamlChanged.fire()
-        }
       }
     )
   }
 
-  abstract managedUpdate(path?: string): Promise<unknown>
+  abstract managedUpdate(path?: string): Promise<void>
 }

extension/src/experiments/index.ts

@@ -44,6 +44,7 @@ import { ConfigKey, getConfigValue, setUserConfigValue } from '../vscode/config'
 import { checkSignalFile, pollSignalFileForProcess } from '../fileSystem'
 import { DVCLIVE_ONLY_RUNNING_SIGNAL_FILE } from '../cli/dvc/constants'
 import { Response } from '../vscode/response'
+import { Pipeline } from '../pipeline'
 
 export const ExperimentsScale = {
   ...omit(ColumnType, 'TIMESTAMP'),
@@ -64,8 +65,9 @@ export class Experiments extends BaseRepository<TableData> {
 
   public readonly viewKey = ViewKey.EXPERIMENTS
 
-  private readonly data: ExperimentsData
+  private readonly pipeline: Pipeline
 
+  private readonly data: ExperimentsData
   private readonly experiments: ExperimentsModel
   private readonly columns: ColumnsModel
 
@@ -96,17 +98,16 @@ export class Experiments extends BaseRepository<TableData> {
   private dvcLiveOnlyCleanupInitialized = false
   private dvcLiveOnlySignalFile: string
 
-  private readonly addStage: () => Promise<boolean>
   private readonly selectBranches: (
     branchesSelected: string[]
   ) => Promise<string[] | undefined>
 
   constructor(
     dvcRoot: string,
     internalCommands: InternalCommands,
+    pipeline: Pipeline,
     resourceLocator: ResourceLocator,
     workspaceState: Memento,
-    addStage: () => Promise<boolean>,
     selectBranches: (
       branchesSelected: string[]
     ) => Promise<string[] | undefined>,
@@ -120,7 +121,7 @@ export class Experiments extends BaseRepository<TableData> {
     )
 
     this.internalCommands = internalCommands
-    this.addStage = addStage
+    this.pipeline = pipeline
     this.selectBranches = selectBranches
 
     this.onDidChangeIsExperimentsFileFocused = this.experimentsFileFocused.event
@@ -146,11 +147,6 @@ export class Experiments extends BaseRepository<TableData> {
     )
 
     this.dispose.track(this.data.onDidUpdate(data => this.setState(data)))
-    this.dispose.track(
-      this.data.onDidChangeDvcYaml(() =>
-        this.webviewMessages.changeHasConfig(true)
-      )
-    )
 
     this.dispose.track(
       workspace.onDidChangeConfiguration((event: ConfigurationChangeEvent) => {
@@ -435,7 +431,8 @@ export class Experiments extends BaseRepository<TableData> {
   }
 
   public async modifyWorkspaceParamsAndRun(
-    commandId: ModifiedExperimentAndRunCommandId
+    commandId: ModifiedExperimentAndRunCommandId,
+    cwd: string
   ) {
     const paramsToModify = await this.pickAndModifyParams()
     if (!paramsToModify) {
@@ -444,13 +441,13 @@ export class Experiments extends BaseRepository<TableData> {
 
     await this.internalCommands.executeCommand<string>(
       commandId,
-      this.dvcRoot,
+      cwd,
       ...paramsToModify
     )
     return this.notifyChanged()
   }
 
-  public async modifyWorkspaceParamsAndQueue() {
+  public async modifyWorkspaceParamsAndQueue(cwd: string) {
     const paramsToModify = await this.pickAndModifyParams()
     if (!paramsToModify) {
       return
@@ -459,7 +456,7 @@ export class Experiments extends BaseRepository<TableData> {
     await Toast.showOutput(
       this.internalCommands.executeCommand<string>(
         AvailableCommands.EXP_QUEUE,
-        this.dvcRoot,
+        cwd,
         ...paramsToModify
       )
     )
@@ -521,16 +518,26 @@ export class Experiments extends BaseRepository<TableData> {
     return this.columns.getRelativeMetricsFiles()
   }
 
+  public getPipelineCwd() {
+    return this.pipeline.getCwd()
+  }
+
   protected sendInitialWebviewData() {
     return this.webviewMessages.sendWebviewMessage()
   }
 
   private setupInitialData() {
     const waitForInitialData = this.dispose.track(
-      this.onDidChangeExperiments(() => {
+      this.onDidChangeExperiments(async () => {
+        await this.pipeline.isReady()
         this.deferred.resolve()
         this.dispose.untrack(waitForInitialData)
         waitForInitialData.dispose()
+        this.dispose.track(
+          this.pipeline.onDidUpdate(() =>
+            this.webviewMessages.sendWebviewMessage()
+          )
+        )
       })
     )
   }
@@ -555,15 +562,10 @@ export class Experiments extends BaseRepository<TableData> {
       this.dvcRoot,
       this.experiments,
       this.columns,
+      this.pipeline,
       () => this.getWebview(),
       () => this.notifyChanged(),
       () => this.selectColumns(),
-      () =>
-        this.internalCommands.executeCommand(
-          AvailableCommands.STAGE_LIST,
-          this.dvcRoot
-        ),
-      () => this.addStage(),
       (branchesSelected: string[]) => this.selectBranches(branchesSelected),
       () => this.data.update()
     )

extension/src/experiments/webview/contract.ts

@@ -104,7 +104,6 @@ export type TableData = {
   hasConfig: boolean
   hasMoreCommits: Record<string, boolean>
   hasRunningWorkspaceExperiment: boolean
-  hasValidDvcYaml: boolean
   isShowingMoreCommits: Record<string, boolean>
   rows: Commit[]
   selectedForPlotsCount: number

extension/src/experiments/webview/messages.ts

@@ -20,25 +20,20 @@ import { SortDefinition } from '../model/sortBy'
 import { getPositiveIntegerInput } from '../../vscode/inputBox'
 import { Title } from '../../vscode/title'
 import { ConfigKey, setConfigValue } from '../../vscode/config'
-import { hasDvcYamlFile } from '../../fileSystem'
 import { NUM_OF_COMMITS_TO_INCREASE } from '../../cli/dvc/constants'
+import { Pipeline } from '../../pipeline'
 
 export class WebviewMessages {
   private readonly dvcRoot: string
 
   private readonly experiments: ExperimentsModel
   private readonly columns: ColumnsModel
+  private readonly pipeline: Pipeline
 
   private readonly getWebview: () => BaseWebview<TableData> | undefined
   private readonly notifyChanged: () => void
   private readonly selectColumns: () => Promise<void>
 
-  private readonly hasStages: () => Promise<string | undefined>
-
-  private hasConfig = false
-  private hasValidDvcYaml = true
-
-  private readonly addStage: () => Promise<boolean>
   private readonly selectBranches: (
     branchesSelected: string[]
   ) => Promise<string[] | undefined>
@@ -49,11 +44,10 @@ export class WebviewMessages {
     dvcRoot: string,
     experiments: ExperimentsModel,
     columns: ColumnsModel,
+    pipeline: Pipeline,
     getWebview: () => BaseWebview<TableData> | undefined,
     notifyChanged: () => void,
     selectColumns: () => Promise<void>,
-    hasStages: () => Promise<string>,
-    addStage: () => Promise<boolean>,
     selectBranches: (
       branchesSelected: string[]
     ) => Promise<string[] | undefined>,
@@ -62,22 +56,12 @@ export class WebviewMessages {
     this.dvcRoot = dvcRoot
     this.experiments = experiments
     this.columns = columns
+    this.pipeline = pipeline
     this.getWebview = getWebview
     this.notifyChanged = notifyChanged
     this.selectColumns = selectColumns
-    this.hasStages = hasStages
-    this.addStage = addStage
     this.selectBranches = selectBranches
     this.update = update
-
-    void this.changeHasConfig()
-  }
-
-  public async changeHasConfig(update?: boolean) {
-    const stages = await this.hasStages()
-    this.hasValidDvcYaml = !hasDvcYamlFile(this.dvcRoot) || stages !== undefined
-    this.hasConfig = !!stages
-    update && this.sendWebviewMessage()
   }
 
   public sendWebviewMessage() {
@@ -267,20 +251,19 @@ export class WebviewMessages {
         this.experiments.getAvailableBranchesToShow().length > 0,
       hasCheckpoints: this.experiments.hasCheckpoints(),
       hasColumns: this.columns.hasNonDefaultColumns(),
-      hasConfig: this.hasConfig,
+      hasConfig: !!(this.pipeline.hasStage() || this.pipeline.hasPipeline()),
       hasMoreCommits: this.experiments.getHasMoreCommits(),
       hasRunningWorkspaceExperiment:
         this.experiments.hasRunningWorkspaceExperiment(),
-      hasValidDvcYaml: this.hasValidDvcYaml,
       isShowingMoreCommits: this.experiments.getIsShowingMoreCommits(),
       rows: this.experiments.getRowData(),
       selectedForPlotsCount: this.experiments.getSelectedRevisions().length,
       sorts: this.experiments.getSorts()
     }
   }
 
-  private async addConfiguration() {
-    await this.addStage()
+  private addConfiguration() {
+    return this.pipeline.checkOrAddPipeline()
   }
 
   private async setMaxTableHeadDepth() {