Merge branch 'main' into add-metric-param-plots

extension/src/experiments/index.ts

@@ -156,6 +156,7 @@ export class Experiments extends BaseRepository<TableData> {
         if (hadCheckpoints !== this.hasCheckpoints()) {
           this.checkpointsChanged.fire()
         }
+        void this.webviewMessages.changeHasConfig(true)
       })
     )
 

extension/src/experiments/webview/contract.ts

@@ -100,6 +100,7 @@ export type TableData = {
   hasColumns: boolean
   hasConfig: boolean
   hasRunningExperiment: boolean
+  hasValidDvcYaml: boolean
   rows: Row[]
   sorts: SortDefinition[]
   filteredCounts: FilteredCounts

extension/src/experiments/webview/messages.ts

@@ -43,6 +43,7 @@ export class WebviewMessages {
   private readonly hasStages: () => Promise<string>
 
   private hasConfig = false
+  private hasValidDvcYaml = true
 
   private readonly addStage: () => Promise<boolean>
 
@@ -74,8 +75,11 @@ export class WebviewMessages {
     this.addStage = addStage
   }
 
-  public async changeHasConfig() {
-    this.hasConfig = !!(await this.hasStages())
+  public async changeHasConfig(update?: boolean) {
+    const stages = await this.hasStages()
+    this.hasValidDvcYaml = stages !== undefined
+    this.hasConfig = !!stages
+    update && this.sendWebviewMessage()
   }
 
   public sendWebviewMessage() {
@@ -218,14 +222,14 @@ export class WebviewMessages {
       hasColumns: this.columns.hasNonDefaultColumns(),
       hasConfig: this.hasConfig,
       hasRunningExperiment: this.experiments.hasRunningExperiment(),
+      hasValidDvcYaml: this.hasValidDvcYaml,
       rows: this.experiments.getRowData(),
       sorts: this.experiments.getSorts()
     }
   }
 
   private async addConfiguration() {
     await this.addStage()
-    await this.changeHasConfig()
   }
 
   private async setMaxTableHeadDepth() {

extension/src/experiments/workspace.test.ts

@@ -14,6 +14,7 @@ import { OutputChannel } from '../vscode/outputChannel'
 import { Title } from '../vscode/title'
 import { Args } from '../cli/dvc/constants'
 import { findOrCreateDvcYamlFile, getFileExtension } from '../fileSystem'
+import { Toast } from '../vscode/toast'
 
 const mockedShowWebview = jest.fn()
 const mockedDisposable = jest.mocked(Disposable)
@@ -595,5 +596,38 @@ describe('Experiments', () => {
         mockedDvcRoot
       )
     })
+
+    it('should show a toast if the dvc.yaml file is invalid', async () => {
+      const showErrorSpy = jest.spyOn(Toast, 'showError')
+
+      mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot)
+      mockedListStages.mockResolvedValueOnce(undefined)
+
+      await workspaceExperiments.getCwdThenRun(mockedCommandId)
+
+      expect(showErrorSpy).toHaveBeenCalledWith(
+        'Cannot perform task. Your dvc.yaml file is invalid.'
+      )
+    })
+
+    it('should not ask to create a stage if the dvc.yaml file is invalid', async () => {
+      mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot)
+      mockedListStages.mockResolvedValueOnce(undefined)
+
+      await workspaceExperiments.getCwdThenRun(mockedCommandId)
+
+      expect(mockedGetValidInput).not.toHaveBeenCalled()
+    })
+
+    it('should not show a toast if the dvc.yaml file is valid', async () => {
+      const showErrorSpy = jest.spyOn(Toast, 'showError')
+
+      mockedQuickPickOne.mockResolvedValueOnce(mockedDvcRoot)
+      mockedListStages.mockResolvedValueOnce('train')
+
+      await workspaceExperiments.getCwdThenRun(mockedCommandId)
+
+      expect(showErrorSpy).not.toHaveBeenCalled()
+    })
   })
 })

extension/src/experiments/workspace.ts

@@ -444,18 +444,30 @@ export class WorkspaceExperiments extends BaseWorkspaceWebviews<
       cwd
     )
 
+    if (stages === undefined) {
+      await Toast.showError(
+        'Cannot perform task. Your dvc.yaml file is invalid.'
+      )
+      return false
+    }
+
     if (!stages) {
-      const stageName = await this.askForStageName()
-      if (!stageName) {
-        return false
-      }
+      return this.addPipeline(cwd)
+    }
+    return true
+  }
 
-      const { trainingScript, command } = await this.askForTrainingScript()
-      if (!trainingScript) {
-        return false
-      }
-      void findOrCreateDvcYamlFile(cwd, trainingScript, stageName, command)
+  private async addPipeline(cwd: string) {
+    const stageName = await this.askForStageName()
+    if (!stageName) {
+      return false
+    }
+
+    const { trainingScript, command } = await this.askForTrainingScript()
+    if (!trainingScript) {
+      return false
     }
+    void findOrCreateDvcYamlFile(cwd, trainingScript, stageName, command)
     return true
   }
 

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

@@ -10,6 +10,7 @@ const tableDataFixture: TableData = {
   hasCheckpoints: true,
   hasConfig: true,
   hasRunningExperiment: true,
+  hasValidDvcYaml: true,
   hasColumns: true,
   sorts: [],
   changes: [],

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

@@ -11,6 +11,7 @@ export const data: TableData = {
   hasCheckpoints: false,
   hasConfig: true,
   hasRunningExperiment: false,
+  hasValidDvcYaml: true,
   sorts: [],
   columns,
   hasColumns: true,

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

@@ -14,6 +14,7 @@ const data: TableData = {
   hasCheckpoints: false,
   hasConfig: true,
   hasRunningExperiment: false,
+  hasValidDvcYaml: true,
   sorts: [
     {
       path: 'params:params.yaml:nested1.doubled',

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

@@ -11,6 +11,7 @@ const data: TableData = {
   hasConfig: true,
   hasRunningExperiment: true,
   hasColumns: true,
+  hasValidDvcYaml: true,
   sorts: [],
   changes: [],
   columnOrder: [],

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

@@ -153,6 +153,7 @@ suite('Experiments Test Suite', () => {
         hasColumns: true,
         hasConfig: true,
         hasRunningExperiment: true,
+        hasValidDvcYaml: true,
         rows: rowsFixture,
         sorts: []
       }
@@ -208,6 +209,7 @@ suite('Experiments Test Suite', () => {
         hasColumns: true,
         hasConfig: false,
         hasRunningExperiment: true,
+        hasValidDvcYaml: true,
         rows: rowsFixture,
         sorts: []
       }
@@ -236,6 +238,7 @@ suite('Experiments Test Suite', () => {
         hasColumns: true,
         hasConfig: true,
         hasRunningExperiment: true,
+        hasValidDvcYaml: true,
         rows: rowsFixture,
         sorts: []
       }
@@ -968,6 +971,7 @@ suite('Experiments Test Suite', () => {
         hasColumns: true,
         hasConfig: true,
         hasRunningExperiment: true,
+        hasValidDvcYaml: true,
         rows: rowsFixture,
         sorts: []
       }

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

@@ -127,6 +127,7 @@ suite('Experiments Filter By Tree Test Suite', () => {
         hasColumns: true,
         hasConfig: true,
         hasRunningExperiment: true,
+        hasValidDvcYaml: true,
         rows: filteredRows,
         sorts: []
       }
@@ -159,6 +160,7 @@ suite('Experiments Filter By Tree Test Suite', () => {
         hasColumns: true,
         hasConfig: true,
         hasRunningExperiment: true,
+        hasValidDvcYaml: true,
         rows: [workspace, main],
         sorts: []
       }
@@ -426,6 +428,7 @@ suite('Experiments Filter By Tree Test Suite', () => {
         hasColumns: true,
         hasConfig: true,
         hasRunningExperiment: true,
+        hasValidDvcYaml: true,
         rows: filteredRows,
         sorts: []
       }

webview/src/experiments/components/AddStage.tsx

@@ -5,16 +5,28 @@ import { IconButton } from '../../shared/components/button/IconButton'
 import { Add } from '../../shared/components/icons'
 import { sendMessage } from '../../shared/vscode'
 
-export const AddStage: React.FC = () => (
+interface AddStageProps {
+  hasValidDvcYaml: boolean
+}
+
+export const AddStage: React.FC<AddStageProps> = ({ hasValidDvcYaml }) => (
   <div className={styles.addConfigButton}>
     <p>Easily and efficiently reproduce your experiments </p>
     <IconButton
       icon={Add}
       onClick={() =>
+        hasValidDvcYaml &&
         sendMessage({ type: MessageFromWebviewType.ADD_CONFIGURATION })
       }
       text="Add a Pipeline Stage"
+      disabled={!hasValidDvcYaml}
     />
+    {!hasValidDvcYaml && (
+      <p className={styles.error}>
+        Your dvc.yaml file should contain valid yaml before adding any pipeline
+        stages.
+      </p>
+    )}
     <p>
       <a href="https://dvc.org/doc/user-guide/project-structure/dvcyaml-files#stages">
         Learn more

webview/src/experiments/components/App.test.tsx

@@ -1495,5 +1495,27 @@ describe('App', () => {
         type: MessageFromWebviewType.ADD_CONFIGURATION
       })
     })
+
+    it('should disable the button and add an error message if the dvc.yaml file contains invalid yaml', async () => {
+      renderTable()
+      setTableData({
+        ...tableDataFixture,
+        hasConfig: false,
+        hasValidDvcYaml: false
+      })
+      const addPipelineButton = await screen.findByText('Add a Pipeline Stage')
+
+      fireEvent.click(addPipelineButton)
+
+      expect(mockPostMessage).not.toHaveBeenCalledWith({
+        type: MessageFromWebviewType.ADD_CONFIGURATION
+      })
+
+      expect(
+        screen.getByText(
+          'Your dvc.yaml file should contain valid yaml before adding any pipeline stages.'
+        )
+      ).toBeInTheDocument()
+    })
   })
 })

webview/src/experiments/components/Experiments.tsx

@@ -133,6 +133,7 @@ export const ExperimentsTable: React.FC = () => {
     columnWidths,
     hasColumns,
     hasConfig,
+    hasValidDvcYaml,
     rows: data
   } = useSelector((state: ExperimentsState) => state.tableData)
 
@@ -195,11 +196,10 @@ export const ExperimentsTable: React.FC = () => {
       />
     )
   }
-
   return (
     <RowSelectionProvider>
       <Table instance={instance} onColumnOrderChange={setColumnOrder} />
-      {!hasConfig && <AddStage />}
+      {!hasConfig && <AddStage hasValidDvcYaml={hasValidDvcYaml} />}
     </RowSelectionProvider>
   )
 }

webview/src/experiments/components/table/tableDataSlice.ts

@@ -24,6 +24,7 @@ export const tableDataInitialState: TableDataState = {
   hasConfig: false,
   hasData: false,
   hasRunningExperiment: false,
+  hasValidDvcYaml: true,
   rows: [],
   sorts: []
 }

webview/src/shared/components/button/Button.tsx

@@ -8,20 +8,23 @@ export type ButtonProps = {
   text: string
   isNested?: boolean
   children?: React.ReactNode
+  disabled?: boolean
 }
 
 export const Button: React.FC<ButtonProps> = ({
   appearance,
   children,
   onClick,
   text,
-  isNested
+  isNested,
+  disabled
 }: ButtonProps) => {
   return (
     <VSCodeButton
       appearance={appearance}
       onClick={onClick}
       className={isNested && styles.secondaryButton}
+      disabled={disabled}
     >
       {text}
       {children}

webview/src/shared/components/button/IconButton.tsx

@@ -9,14 +9,16 @@ export const IconButton: React.FC<IconButtonProps> = ({
   onClick,
   icon,
   isNested,
-  text
+  text,
+  disabled
 }: IconButtonProps) => {
   return (
     <Button
       appearance={appearance}
       isNested={isNested}
       onClick={onClick}
       text={text}
+      disabled={disabled}
     >
       <span slot="start">
         <Icon icon={icon} width={18} height={18} />

webview/src/stories/Table.stories.tsx

@@ -46,6 +46,7 @@ const tableData: TableDataState = {
   hasConfig: true,
   hasData: true,
   hasRunningExperiment: true,
+  hasValidDvcYaml: true,
   rows: addCommitDataToMainBranch(rowsFixture).map(row => ({
     ...row,
     subRows: row.subRows?.map(experiment => ({

webview/src/test/sort.ts

@@ -50,6 +50,7 @@ export const tableData: TableData = {
   hasColumns: true,
   hasConfig: true,
   hasRunningExperiment: false,
+  hasValidDvcYaml: true,
   rows: [
     {
       id: EXPERIMENT_WORKSPACE_ID,