Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add experiment checkpoints model #1281

Merged
merged 3 commits into from
Feb 3, 2022
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
add experiment checkpoints model
mattseddon committed Feb 3, 2022
commit 9f4f65a0d2ca548cad67ce582c2f7ea2f7c3576d
89 changes: 89 additions & 0 deletions extension/src/experiments/checkpoints/collect.test.ts
Original file line number Diff line number Diff line change
@@ -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)
})
})
13 changes: 13 additions & 0 deletions extension/src/experiments/checkpoints/collect.ts
Original file line number Diff line number Diff line change
@@ -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
}
})
}
59 changes: 59 additions & 0 deletions extension/src/experiments/checkpoints/model.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,59 @@
import { join } from 'path'
import { ExperimentCheckpointsModel } from './model'
import { dvcDemoPath } from '../../test/util'

describe('ExperimentCheckpointsModel', () => {
it('should keep a record of yaml files that have checkpoints', () => {
const experimentCheckpointsModel = new ExperimentCheckpointsModel()
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)
})
})
30 changes: 30 additions & 0 deletions extension/src/experiments/checkpoints/model.ts
Original file line number Diff line number Diff line change
@@ -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 ExperimentCheckpointsModel {
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
)
}
}
}
5 changes: 5 additions & 0 deletions extension/src/fileSystem/index.ts
Original file line number Diff line number Diff line change
@@ -54,6 +54,11 @@ export const isSameOrChild = (root: string, path: string) => {
return !rel.startsWith('..')
}

export type PartialDvcYaml = {
stages: {
train: { outs: (string | Record<string, { checkpoint?: boolean }>)[] }
}
}
export const isAnyDvcYaml = (path?: string): boolean =>
!!(
path &&