From e80267df420499a44f302443b7f95adf03ae91e6 Mon Sep 17 00:00:00 2001
From: Matt Seddon <mattseddon@hotmail.com>
Date: Wed, 17 Aug 2022 16:29:22 +1000
Subject: [PATCH] Do not remove watcher from previously collected metric file
 (exp show live updates)

---
 extension/src/data/index.ts                   |  2 +-
 .../src/experiments/data/collect.test.ts      | 38 ++++++++++++++++---
 extension/src/experiments/data/collect.ts     | 18 +++++----
 extension/src/experiments/data/index.ts       |  2 +-
 4 files changed, 46 insertions(+), 14 deletions(-)

diff --git a/extension/src/data/index.ts b/extension/src/data/index.ts
index 62af3504f4..108f36b0aa 100644
--- a/extension/src/data/index.ts
+++ b/extension/src/data/index.ts
@@ -16,8 +16,8 @@ export abstract class BaseData<
   protected readonly dvcRoot: string
   protected readonly processManager: ProcessManager
   protected readonly internalCommands: InternalCommands
+  protected collectedFiles: string[] = []
 
-  private collectedFiles: string[] = []
   private readonly staticFiles: string[]
 
   private readonly updated: EventEmitter<T> = this.dispose.track(
diff --git a/extension/src/experiments/data/collect.test.ts b/extension/src/experiments/data/collect.test.ts
index 588ae24853..0d8d37211a 100644
--- a/extension/src/experiments/data/collect.test.ts
+++ b/extension/src/experiments/data/collect.test.ts
@@ -5,7 +5,7 @@ import expShowFixture from '../../test/fixtures/expShow/output'
 
 describe('collectFiles', () => {
   it('should collect all of the available files from the test fixture', () => {
-    expect(collectFiles(expShowFixture)).toStrictEqual([
+    expect(collectFiles(expShowFixture, [])).toStrictEqual([
       'params.yaml',
       join('nested', 'params.yaml'),
       'summary.json'
@@ -21,7 +21,7 @@ describe('collectFiles', () => {
       }
     }
 
-    expect(collectFiles(workspace)).toStrictEqual([])
+    expect(collectFiles(workspace, [])).toStrictEqual([])
   })
 
   it('should handle a missing params key', () => {
@@ -37,7 +37,7 @@ describe('collectFiles', () => {
       }
     }
 
-    expect(collectFiles(workspace)).toStrictEqual(['logs.json'])
+    expect(collectFiles(workspace, [])).toStrictEqual(['logs.json'])
   })
 
   it('should handle a missing metrics key', () => {
@@ -53,7 +53,7 @@ describe('collectFiles', () => {
       }
     }
 
-    expect(collectFiles(workspace)).toStrictEqual(['params.yaml'])
+    expect(collectFiles(workspace, [])).toStrictEqual(['params.yaml'])
   })
 
   it('should collect all of the available files from a more complex example', () => {
@@ -76,7 +76,7 @@ describe('collectFiles', () => {
       }
     } as ExperimentsOutput
 
-    expect(collectFiles(workspace).sort()).toStrictEqual([
+    expect(collectFiles(workspace, []).sort()).toStrictEqual([
       'further/nested/params.yaml',
       'logs.json',
       'metrics.json',
@@ -85,4 +85,32 @@ describe('collectFiles', () => {
       'summary.json'
     ])
   })
+
+  it('should not remove a previously collected file if it is deleted (removal breaks live updates logged by dvclive)', () => {
+    const workspace = {
+      workspace: {
+        baseline: {
+          data: {
+            executor: 'workspace',
+            metrics: {},
+            params: {
+              'params.yaml': {
+                data: {
+                  epochs: 100
+                }
+              }
+            },
+            queued: false,
+            running: true,
+            timestamp: null
+          }
+        }
+      }
+    } as ExperimentsOutput
+
+    expect(collectFiles(workspace, ['dvclive.json'])).toStrictEqual([
+      'params.yaml',
+      'dvclive.json'
+    ])
+  })
 })
diff --git a/extension/src/experiments/data/collect.ts b/extension/src/experiments/data/collect.ts
index 30add5506a..8a1a84e5d0 100644
--- a/extension/src/experiments/data/collect.ts
+++ b/extension/src/experiments/data/collect.ts
@@ -1,12 +1,16 @@
 import { ExperimentsOutput } from '../../cli/reader'
 
-export const collectFiles = (data: ExperimentsOutput): string[] => {
-  const files = new Set<string>(
-    Object.keys({
-      ...data?.workspace?.baseline?.data?.params,
-      ...data?.workspace?.baseline?.data?.metrics
-    }).filter(Boolean)
-  )
+export const collectFiles = (
+  data: ExperimentsOutput,
+  existingFiles: string[]
+): string[] => {
+  const files = new Set<string>([
+    ...Object.keys({
+      ...data?.workspace.baseline?.data?.params,
+      ...data?.workspace.baseline?.data?.metrics
+    }).filter(Boolean),
+    ...existingFiles
+  ])
 
   return [...files]
 }
diff --git a/extension/src/experiments/data/index.ts b/extension/src/experiments/data/index.ts
index 0c5d0e59e7..fd9746b665 100644
--- a/extension/src/experiments/data/index.ts
+++ b/extension/src/experiments/data/index.ts
@@ -48,7 +48,7 @@ export class ExperimentsData extends BaseData<ExperimentsOutput> {
   }
 
   public collectFiles(data: ExperimentsOutput) {
-    return collectFiles(data)
+    return collectFiles(data, this.collectedFiles)
   }
 
   public managedUpdate(path?: string) {