Watch metric files to update plots (#3321)

* collect and watch metrics files for plots updates

* test change

* remove original patch

* unit test collect function

* self review

Author: mattseddon

extension/src/plots/data/collect.test.ts

@@ -0,0 +1,42 @@
+import { collectMetricsFiles } from './collect'
+import expShowFixture from '../../test/fixtures/expShow/base/output'
+import { ExperimentsOutput } from '../../cli/dvc/contract'
+
+describe('collectMetricsFiles', () => {
+  it('should return the expected metrics files from the test fixture', () => {
+    expect(collectMetricsFiles(expShowFixture, [])).toStrictEqual([
+      'summary.json'
+    ])
+  })
+
+  it('should persist existing files', () => {
+    const existingFiles = ['file.json', 'file.py', 'file.tsv', 'file.txt']
+
+    expect(
+      collectMetricsFiles({ workspace: { baseline: {} } }, existingFiles)
+    ).toStrictEqual(existingFiles)
+  })
+
+  it('should not fail when given an error', () => {
+    const existingFile = ['metrics.json']
+
+    expect(
+      collectMetricsFiles(
+        {
+          workspace: {
+            baseline: { error: { msg: 'I broke', type: 'not important' } }
+          }
+        },
+        existingFile
+      )
+    ).toStrictEqual(existingFile)
+  })
+
+  it('should not fail when given empty output', () => {
+    const existingFiles: string[] = []
+
+    expect(
+      collectMetricsFiles({} as ExperimentsOutput, existingFiles)
+    ).toStrictEqual(existingFiles)
+  })
+})

extension/src/plots/data/collect.ts

@@ -1,4 +1,4 @@
-import { PlotsOutputOrError } from '../../cli/dvc/contract'
+import { ExperimentsOutput, PlotsOutputOrError } from '../../cli/dvc/contract'
 import { isDvcError } from '../../cli/dvc/reader'
 import { uniqueValues } from '../../util/array'
 import { isImagePlot, Plot, TemplatePlot } from '../webview/contract'
@@ -56,3 +56,14 @@ export const collectFiles = (
 
   return uniqueValues(acc)
 }
+
+export const collectMetricsFiles = (
+  data: ExperimentsOutput,
+  existingFiles: string[]
+): string[] =>
+  uniqueValues([
+    ...Object.keys({
+      ...data?.workspace?.baseline?.data?.metrics
+    }).filter(Boolean),
+    ...existingFiles
+  ]).sort()

extension/src/plots/data/index.ts

@@ -1,12 +1,13 @@
 import { EventEmitter } from 'vscode'
-import { collectFiles } from './collect'
+import { collectMetricsFiles, collectFiles } from './collect'
 import {
   EXPERIMENT_WORKSPACE_ID,
+  ExperimentsOutput,
   PlotsOutputOrError
 } from '../../cli/dvc/contract'
 import { AvailableCommands, InternalCommands } from '../../commands/internal'
 import { BaseData } from '../../data'
-import { flattenUnique, sameContents } from '../../util/array'
+import { flattenUnique, sameContents, uniqueValues } from '../../util/array'
 import { PlotsModel } from '../model'
 
 export class PlotsData extends BaseData<{
@@ -15,6 +16,8 @@ export class PlotsData extends BaseData<{
 }> {
   private readonly model: PlotsModel
 
+  private metricFiles: string[] = []
+
   constructor(
     dvcRoot: string,
     internalCommands: InternalCommands,
@@ -31,7 +34,7 @@ export class PlotsData extends BaseData<{
           process: () => this.update()
         }
       ],
-      ['dvc.yaml', 'dvc.lock', 'metrics.json']
+      ['dvc.yaml', 'dvc.lock']
     )
     this.model = model
   }
@@ -62,6 +65,17 @@ export class PlotsData extends BaseData<{
     return this.processManager.run('update')
   }
 
+  public setMetricFiles(data: ExperimentsOutput) {
+    const metricsFiles = collectMetricsFiles(data, this.metricFiles)
+    if (!sameContents(metricsFiles, this.metricFiles)) {
+      this.metricFiles = metricsFiles
+      this.collectedFiles = uniqueValues([
+        ...this.collectedFiles,
+        ...metricsFiles
+      ])
+    }
+  }
+
   protected collectFiles({ data }: { data: PlotsOutputOrError }) {
     this.collectedFiles = collectFiles(data, this.collectedFiles)
   }

extension/src/plots/index.ts

@@ -28,7 +28,6 @@ export class Plots extends BaseRepository<TPlotsData> {
 
   private readonly pathsChanged = this.dispose.track(new EventEmitter<void>())
 
-  private readonly experiments: Experiments
   private readonly plots: PlotsModel
   private readonly paths: PathsModel
   private readonly data: PlotsData
@@ -45,8 +44,6 @@ export class Plots extends BaseRepository<TPlotsData> {
   ) {
     super(dvcRoot, webviewIcon)
 
-    this.experiments = experiments
-
     this.plots = this.dispose.track(
       new PlotsModel(this.dvcRoot, experiments, workspaceState)
     )
@@ -57,7 +54,7 @@ export class Plots extends BaseRepository<TPlotsData> {
     this.webviewMessages = this.createWebviewMessageHandler(
       this.paths,
       this.plots,
-      this.experiments
+      experiments
     )
 
     this.data = this.dispose.track(
@@ -174,6 +171,7 @@ export class Plots extends BaseRepository<TPlotsData> {
         if (data) {
           this.dispose.untrack(waitForInitialExpData)
           waitForInitialExpData.dispose()
+          this.data.setMetricFiles(data)
           this.setupExperimentsListener(experiments)
           void this.initializeData(data)
         }
@@ -185,7 +183,10 @@ export class Plots extends BaseRepository<TPlotsData> {
     this.dispose.track(
       experiments.onDidChangeExperiments(async data => {
         if (data) {
-          await this.plots.transformAndSetExperiments(data)
+          await Promise.all([
+            this.plots.transformAndSetExperiments(data),
+            this.data.setMetricFiles(data)
+          ])
         }
 
         this.notifyChanged()

extension/src/test/suite/plots/data/index.test.ts

@@ -134,7 +134,14 @@ suite('Plots Data Test Suite', () => {
     it('should collect files and watch them for updates', async () => {
       const mockNow = getMockNow()
       const parentDirectory = 'training'
-      const watchedFile = join(parentDirectory, 'metrics.json')
+      const metricsFile = join(parentDirectory, 'metrics.json')
+      const collectedFile = join(
+        'training',
+        'plots',
+        'metrics',
+        'train',
+        'acc.tsv'
+      )
 
       const mockExecuteCommand = (command: CommandId) => {
         if (command === AvailableCommands.PLOTS_DIFF) {
@@ -146,28 +153,12 @@ suite('Plots Data Test Suite', () => {
                     {
                       dvc_data_version_info: {
                         field: join('train', 'acc'),
-                        filename: join(
-                          'training',
-                          'plots',
-                          'metrics',
-                          'train',
-                          'acc.tsv'
-                        ),
+                        filename: collectedFile,
                         revision: EXPERIMENT_WORKSPACE_ID
                       },
                       dvc_inferred_y_value: '0.2707333333333333',
                       step: '0',
                       [join('train', 'acc')]: '0.2707333333333333'
-                    },
-                    {
-                      dvc_data_version_info: {
-                        field: join('test', 'acc'),
-                        filename: watchedFile,
-                        revision: EXPERIMENT_WORKSPACE_ID
-                      },
-                      dvc_inferred_y_value: '0.2712',
-                      step: '0',
-                      [join('test', 'acc')]: '0.2712'
                     }
                   ]
                 },
@@ -196,6 +187,17 @@ suite('Plots Data Test Suite', () => {
 
       void data.update()
       await data.isReady()
+      data.setMetricFiles({
+        workspace: {
+          baseline: { data: { metrics: { [metricsFile]: { data: {} } } } }
+        }
+      })
+
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      expect((data as any).collectedFiles).to.deep.equal([
+        collectedFile,
+        metricsFile
+      ])
 
       bypassProcessManagerDebounce(mockNow)
 
@@ -204,7 +206,7 @@ suite('Plots Data Test Suite', () => {
         disposable.track(data.onDidUpdate(() => resolve(undefined)))
       )
 
-      await fireWatcher(join(dvcDemoPath, watchedFile))
+      await fireWatcher(join(dvcDemoPath, metricsFile))
       await dataUpdatedEvent
 
       expect(