From 949776792ad29f84804b0d4ef85c292a057e732d Mon Sep 17 00:00:00 2001
From: Matt Seddon <37993418+mattseddon@users.noreply.github.com>
Date: Sun, 24 Nov 2024 14:18:18 +1100
Subject: [PATCH] Exclude queued experiments from dvc root collection (#5833)

---
 extension/src/cli/dvc/constants.ts     |  2 +-
 extension/src/fileSystem/index.test.ts | 13 +++++++++++++
 extension/src/fileSystem/index.ts      | 12 ++++++++----
 extension/src/pipeline/data.ts         |  6 ++++--
 4 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/extension/src/cli/dvc/constants.ts b/extension/src/cli/dvc/constants.ts
index 6005988999..174fda2cd3 100644
--- 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
index ecd475a045..4a30138d1b 100644
--- 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
index bad54a5c28..2a19fc957d 100644
--- 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
index 98ab550846..c8428b8edf 100644
--- 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))
   }
 }