Exclude queued experiments from dvc root collection (#5833)

iterative · Nov 24, 2024 · 9497767 · 9497767
1 parent e8e153d
commit 9497767
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 7 deletions.
diff --git a/extension/src/cli/dvc/constants.ts b/extension/src/cli/dvc/constants.ts
@@ -12,7 +12,7 @@ export const TEMP_PLOTS_DIR = join(DOT_DVC, 'tmp', 'plots')
 
 export const FIELD_SEPARATOR = '::'
 
-const TEMP_EXP_DIR = join(DOT_DVC, 'tmp', 'exps')
+export const TEMP_EXP_DIR = join(DOT_DVC, 'tmp', 'exps')
 const TEMP_EXP_RUN_DIR = join(TEMP_EXP_DIR, 'run')
 export const DVCLIVE_ONLY_RUNNING_SIGNAL_FILE = join(
   TEMP_EXP_RUN_DIR,

diff --git a/extension/src/fileSystem/index.test.ts b/extension/src/fileSystem/index.test.ts
@@ -337,6 +337,19 @@ describe('findDvcRootPaths', () => {
       new Set([dvcDemoPath, mockFirstDvcRoot, mockSecondDvcRoot])
     )
   })
+
+  it('should exclude queued experiment directories', async () => {
+    const queuedExpDir = join(dvcDemoPath, DOT_DVC, 'tmp', 'exps', 'tmp_2xr')
+
+    mockedFindFiles.mockResolvedValue([
+      convertAsFindDvcConfigFile(dvcDemoPath),
+      convertAsFindDvcConfigFile(queuedExpDir)
+    ])
+
+    const dvcRoots = await findDvcRootPaths()
+
+    expect(dvcRoots).toStrictEqual(new Set([dvcDemoPath]))
+  })
 })
 
 describe('findAbsoluteDvcRootPath', () => {

diff --git a/extension/src/fileSystem/index.ts b/extension/src/fileSystem/index.ts
@@ -30,7 +30,7 @@ import { gitPath } from '../cli/git/constants'
 import { createValidInteger } from '../util/number'
 import { processExists } from '../process/execution'
 import { getFirstWorkspaceFolder } from '../vscode/workspaceFolders'
-import { DOT_DVC, FULLY_NESTED_DVC } from '../cli/dvc/constants'
+import { DOT_DVC, FULLY_NESTED_DVC, TEMP_EXP_DIR } from '../cli/dvc/constants'
 import { delay } from '../util/time'
 import { PlotConfigData, PlotConfigDataAxis } from '../pipeline/quickPick'
 
@@ -72,9 +72,13 @@ export const findDvcRootPaths = async (): Promise<Set<string>> => {
 
   const nested = await findFiles(FULLY_NESTED_DVC)
   if (definedAndNonEmpty(nested)) {
-    dvcRoots.push(
-      ...nested.map(nestedRoot => standardizePath(dirname(dirname(nestedRoot))))
-    )
+    for (const nestedRoot of nested) {
+      const dvcRoot = standardizePath(dirname(dirname(nestedRoot)))
+      if (dvcRoot.includes(TEMP_EXP_DIR)) {
+        continue
+      }
+      dvcRoots.push(dvcRoot)
+    }
   }
 
   return new Set(sortCollectedArray(dvcRoots))

diff --git a/extension/src/pipeline/data.ts b/extension/src/pipeline/data.ts
@@ -3,6 +3,7 @@ import { AvailableCommands, InternalCommands } from '../commands/internal'
 import { BaseData } from '../data'
 import { findFiles } from '../fileSystem/workspace'
 import { isPathInProject } from '../fileSystem'
+import { TEMP_EXP_DIR } from '../cli/dvc/constants'
 
 export class PipelineData extends BaseData<{
   dag: string
@@ -57,7 +58,8 @@ export class PipelineData extends BaseData<{
     return this.notifyChanged({ dag, stages })
   }
 
-  private findDvcYamls() {
-    return findFiles('**/dvc.yaml')
+  private async findDvcYamls() {
+    const dvcYamls = await findFiles('**/dvc.yaml')
+    return dvcYamls.filter(file => !file.includes(TEMP_EXP_DIR))
   }
 }