Merge branch 'master' into fix-stub

Author: mattseddon

extension/package.json

@@ -1063,6 +1063,7 @@
     "chokidar": "^3.5.2",
     "execa": "^5.1.1",
     "fs-extra": "^10.0.0",
+    "js-yaml": "^4.1.0",
     "lodash.clonedeep": "^4.5.0",
     "lodash.get": "^4.4.2",
     "lodash.isempty": "^4.4.0",
@@ -1078,6 +1079,7 @@
     "@types/copy-webpack-plugin": "^8.0.0",
     "@types/fs-extra": "^9.0.13",
     "@types/jest": "^27.4.0",
+    "@types/js-yaml": "^4.0.5",
     "@types/lodash.clonedeep": "^4.5.6",
     "@types/lodash.get": "^4.4.6",
     "@types/lodash.isempty": "^4.4.6",

extension/src/experiments/checkpoints/collect.test.ts

@@ -0,0 +1,89 @@
+import { collectHasCheckpoints } from './collect'
+import { PartialDvcYaml } from '../../fileSystem'
+
+describe('collectHasCheckpoints', () => {
+  it('should correctly identify that the demo project has checkpoints', () => {
+    const hasCheckpoints = collectHasCheckpoints({
+      stages: {
+        train: {
+          cmd: 'python train.py',
+          deps: ['data/MNIST', 'train.py'],
+          live: { logs: { html: true, summary: true } },
+          outs: [{ 'model.pt': { checkpoint: true } }],
+          params: ['seed', 'lr', 'weight_decay'],
+          plots: [
+            'plots',
+            {
+              'predictions.json': {
+                cache: false,
+                template: 'confusion',
+                x: 'actual',
+                y: 'predicted'
+              }
+            }
+          ]
+        }
+      }
+    } as PartialDvcYaml)
+    expect(hasCheckpoints).toBe(true)
+  })
+
+  it('should correctly classify a dvc.yaml without checkpoint', () => {
+    const hasCheckpoints = collectHasCheckpoints({
+      stages: {
+        extract: {
+          cmd: 'tar -xzf data/images.tar.gz --directory data',
+          deps: ['data/images.tar.gz'],
+          outs: [{ 'data/images/': { cache: false } }]
+        },
+        train: {
+          cmd: 'python3 src/train.py',
+          deps: ['data/images/', 'src/train.py'],
+          live: { logs: { html: true, summary: true } },
+          metrics: [{ 'metrics.json': { cache: false } }],
+          outs: ['models/model.h5'],
+          params: ['model.conv_units', 'train.epochs'],
+          plots: [{ 'logs.csv': { cache: false } }]
+        }
+      }
+    } as PartialDvcYaml)
+
+    expect(hasCheckpoints).toBe(false)
+  })
+
+  it('should correctly classify a more complex dvc.yaml without checkpoint', () => {
+    const hasCheckpoints = collectHasCheckpoints({
+      stages: {
+        evaluate: {
+          cmd: 'python src/evaluate.py model.pkl data/features scores.json prc.json roc.json',
+          deps: ['data/features', 'model.pkl', 'src/evaluate.py'],
+          metrics: [{ 'scores.json': { cache: false } }],
+          plots: [
+            { 'prc.json': { cache: false, x: 'recall', y: 'precision' } },
+            { 'roc.json': { cache: false, x: 'fpr', y: 'tpr' } }
+          ]
+        },
+        featurize: {
+          cmd: 'python src/featurization.py data/prepared data/features',
+          deps: ['data/prepared', 'src/featurization.py'],
+          outs: ['data/features'],
+          params: ['featurize.max_features', 'featurize.ngrams']
+        },
+        prepare: {
+          cmd: 'python src/prepare.py data/data.xml',
+          deps: ['data/data.xml', 'src/prepare.py'],
+          outs: ['data/prepared'],
+          params: ['prepare.seed', 'prepare.split']
+        },
+        train: {
+          cmd: 'python src/train.py data/features model.pkl',
+          deps: ['data/features', 'src/train.py'],
+          outs: ['model.pkl'],
+          params: ['train.min_split', 'train.n_est', 'train.seed']
+        }
+      }
+    } as PartialDvcYaml)
+
+    expect(hasCheckpoints).toBe(false)
+  })
+})

extension/src/experiments/checkpoints/collect.ts

@@ -0,0 +1,13 @@
+import { PartialDvcYaml } from '../../fileSystem'
+
+export const collectHasCheckpoints = (yaml: PartialDvcYaml): boolean => {
+  return !!yaml.stages.train.outs.find(out => {
+    if (typeof out === 'string') {
+      return false
+    }
+
+    if (Object.values(out).find(file => file?.checkpoint)) {
+      return true
+    }
+  })
+}

extension/src/experiments/checkpoints/model.test.ts

@@ -0,0 +1,59 @@
+import { join } from 'path'
+import { CheckpointsModel } from './model'
+import { dvcDemoPath } from '../../test/util'
+
+describe('CheckpointsModel', () => {
+  it('should keep a record of yaml files that have checkpoints', () => {
+    const experimentCheckpointsModel = new CheckpointsModel()
+    expect(experimentCheckpointsModel.hasCheckpoints()).toBe(false)
+
+    const rootYamlHasCheckpoints = {
+      path: join(dvcDemoPath, 'dvc.yaml'),
+      yaml: {
+        stages: {
+          train: {
+            outs: [{ 'model.pt': { checkpoint: true } }]
+          }
+        }
+      }
+    }
+
+    const extraYamlHasCheckpoints = {
+      path: join(dvcDemoPath, 'extra', 'dvc.yaml'),
+      yaml: {
+        stages: {
+          train: {
+            outs: [{ 'extra-model.pt': { checkpoint: true } }]
+          }
+        }
+      }
+    }
+
+    const rootYamlNoCheckpoints = {
+      path: join(dvcDemoPath, 'dvc.yaml'),
+      yaml: {
+        stages: { train: { outs: ['model.pt'] } }
+      }
+    }
+
+    const extraYamlNoCheckpoints = {
+      path: join(dvcDemoPath, 'extra', 'dvc.yaml'),
+      yaml: {
+        stages: { train: { outs: ['extra-model.pt'] } }
+      }
+    }
+
+    experimentCheckpointsModel.transformAndSet(rootYamlHasCheckpoints)
+    expect(experimentCheckpointsModel.hasCheckpoints()).toBe(true)
+
+    experimentCheckpointsModel.transformAndSet(extraYamlHasCheckpoints)
+    expect(experimentCheckpointsModel.hasCheckpoints()).toBe(true)
+
+    experimentCheckpointsModel.transformAndSet(rootYamlNoCheckpoints)
+    expect(experimentCheckpointsModel.hasCheckpoints()).toBe(true)
+
+    experimentCheckpointsModel.transformAndSet(extraYamlNoCheckpoints)
+
+    expect(experimentCheckpointsModel.hasCheckpoints()).toBe(false)
+  })
+})

extension/src/experiments/checkpoints/model.ts

@@ -0,0 +1,30 @@
+import { Disposable } from '@hediet/std/disposable'
+import { collectHasCheckpoints } from './collect'
+import { PartialDvcYaml } from '../../fileSystem'
+import { definedAndNonEmpty, uniqueValues } from '../../util/array'
+
+export class CheckpointsModel {
+  public dispose = Disposable.fn()
+
+  private yamlWithCheckpoints: string[] = []
+
+  public hasCheckpoints() {
+    return definedAndNonEmpty(this.yamlWithCheckpoints)
+  }
+
+  public transformAndSet(data: { path: string; yaml: PartialDvcYaml }) {
+    const { path, yaml } = data
+    const hasCheckpoints = collectHasCheckpoints(yaml)
+
+    if (hasCheckpoints) {
+      this.yamlWithCheckpoints = uniqueValues([
+        ...this.yamlWithCheckpoints,
+        path
+      ])
+    } else {
+      this.yamlWithCheckpoints = this.yamlWithCheckpoints.filter(
+        file => file !== path
+      )
+    }
+  }
+}

extension/src/experiments/index.ts

@@ -9,6 +9,7 @@ import {
 } from './model/filterBy/quickPick'
 import { pickSortsToRemove, pickSortToAdd } from './model/sortBy/quickPick'
 import { MetricsAndParamsModel } from './metricsAndParams/model'
+import { CheckpointsModel } from './checkpoints/model'
 import { ExperimentsData } from './data'
 import { TableData } from './webview/contract'
 import { ResourceLocator } from '../resourceLocator'
@@ -22,17 +23,20 @@ import {
   MessageFromWebviewType
 } from '../webview/contract'
 import { Logger } from '../common/logger'
+import { FileSystemData } from '../fileSystem/data'
 
 export class Experiments extends BaseRepository<TableData> {
   public readonly onDidChangeExperiments: Event<ExperimentsOutput | void>
   public readonly onDidChangeMetricsOrParams: Event<void>
 
   public readonly viewKey = ViewKey.EXPERIMENTS
 
-  private readonly data: ExperimentsData
+  private readonly cliData: ExperimentsData
+  private readonly fileSystemData: FileSystemData
 
   private readonly experiments: ExperimentsModel
   private readonly metricsAndParams: MetricsAndParamsModel
+  private readonly checkpoints: CheckpointsModel
 
   private readonly experimentsChanged = this.dispose.track(
     new EventEmitter<ExperimentsOutput | void>()
@@ -48,7 +52,8 @@ export class Experiments extends BaseRepository<TableData> {
     updatesPaused: EventEmitter<boolean>,
     resourceLocator: ResourceLocator,
     workspaceState: Memento,
-    data?: ExperimentsData
+    cliData?: ExperimentsData,
+    fileSystemData?: FileSystemData
   ) {
     super(dvcRoot, resourceLocator.beaker)
 
@@ -63,11 +68,22 @@ export class Experiments extends BaseRepository<TableData> {
       new MetricsAndParamsModel(dvcRoot, workspaceState)
     )
 
-    this.data = this.dispose.track(
-      data || new ExperimentsData(dvcRoot, internalCommands, updatesPaused)
+    this.checkpoints = this.dispose.track(new CheckpointsModel())
+
+    this.cliData = this.dispose.track(
+      cliData || new ExperimentsData(dvcRoot, internalCommands, updatesPaused)
     )
 
-    this.dispose.track(this.data.onDidUpdate(data => this.setState(data)))
+    this.fileSystemData = this.dispose.track(
+      fileSystemData || new FileSystemData(dvcRoot)
+    )
+
+    this.dispose.track(this.cliData.onDidUpdate(data => this.setState(data)))
+    this.dispose.track(
+      this.fileSystemData.onDidUpdate(data =>
+        this.checkpoints.transformAndSet(data)
+      )
+    )
 
     this.handleMessageFromWebview()
 
@@ -81,11 +97,11 @@ export class Experiments extends BaseRepository<TableData> {
   }
 
   public update() {
-    return this.data.managedUpdate()
+    return this.cliData.managedUpdate()
   }
 
   public forceUpdate() {
-    return this.data.forceUpdate()
+    return this.cliData.forceUpdate()
   }
 
   public async setState(data: ExperimentsOutput) {
@@ -97,6 +113,10 @@ export class Experiments extends BaseRepository<TableData> {
     return this.notifyChanged(data)
   }
 
+  public hasCheckpoints() {
+    return this.checkpoints.hasCheckpoints()
+  }
+
   public getChildMetricsOrParams(path?: string) {
     return this.metricsAndParams.getChildren(path)
   }

extension/src/fileSystem/data/index.ts

@@ -0,0 +1,62 @@
+import { Event, EventEmitter } from 'vscode'
+import { Disposable } from '@hediet/std/disposable'
+import { Deferred } from '@hediet/std/synchronization'
+import { isSameOrChild, loadYaml, PartialDvcYaml } from '..'
+import { findFiles } from '../workspace'
+import { join } from '../../test/util/path'
+import { createFileSystemWatcher } from '../watcher'
+
+export class FileSystemData {
+  public readonly dispose = Disposable.fn()
+
+  public readonly onDidUpdate: Event<{ path: string; yaml: PartialDvcYaml }>
+
+  private readonly dvcRoot: string
+
+  private readonly updated = this.dispose.track(
+    new EventEmitter<{ path: string; yaml: PartialDvcYaml }>()
+  )
+
+  private readonly deferred = new Deferred()
+  private readonly initialized = this.deferred.promise
+
+  constructor(dvcRoot: string) {
+    this.dvcRoot = dvcRoot
+    this.onDidUpdate = this.updated.event
+
+    this.watchDvcYaml()
+    this.initialize()
+  }
+
+  public isReady() {
+    return this.initialized
+  }
+
+  private async initialize() {
+    const files = await findFiles(join('**', 'dvc.yaml'))
+    const filesInRepo = files.filter(file => isSameOrChild(this.dvcRoot, file))
+
+    filesInRepo.map(path => {
+      const yaml = loadYaml<PartialDvcYaml>(path)
+      if (yaml) {
+        this.updated.fire({ path, yaml })
+      }
+    })
+
+    this.deferred.resolve()
+  }
+
+  private watchDvcYaml() {
+    this.dispose.track(
+      createFileSystemWatcher(join(this.dvcRoot, '**', 'dvc.yaml'), path => {
+        if (!path) {
+          return
+        }
+        const yaml = loadYaml<PartialDvcYaml>(path)
+        if (yaml) {
+          this.updated.fire({ path, yaml })
+        }
+      })
+    )
+  }
+}