From c85a46cddaf7e32f39b7c46a6511ee86ead265c5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Peter=20Rowlands=20=28=EB=B3=80=EA=B8=B0=ED=98=B8=29?= Date: Fri, 6 Nov 2020 14:44:21 +0900 Subject: [PATCH] experiments: add checkpoint output type (#4845) * output: add checkpoint output type * stage: handle checkpoint outs * run/repro: handle checkpoint outs * experiments: use checkpoint output type in tests * tests: update run command tests --- dvc/command/run.py | 12 ++++++- dvc/output/__init__.py | 33 +++++++++++++++-- dvc/output/base.py | 8 ++++- dvc/repo/reproduce.py | 12 +++++-- dvc/repo/run.py | 1 + dvc/schema.py | 6 +++- dvc/stage/__init__.py | 18 +++++++--- dvc/stage/serialize.py | 3 ++ dvc/stage/utils.py | 3 ++ tests/func/experiments/test_experiments.py | 42 +++++----------------- tests/unit/command/test_run.py | 5 +++ 11 files changed, 97 insertions(+), 46 deletions(-) diff --git a/dvc/command/run.py b/dvc/command/run.py index 735fc044e..f5be09dba 100644 --- a/dvc/command/run.py +++ b/dvc/command/run.py @@ -21,6 +21,7 @@ class CmdRun(CmdBase): self.args.plots_no_cache, self.args.outs_persist, self.args.outs_persist_no_cache, + self.args.checkpoints, self.args.params, self.args.command, ] @@ -45,12 +46,13 @@ class CmdRun(CmdBase): params=self.args.params, fname=self.args.file, wdir=self.args.wdir, - no_exec=self.args.no_exec, + no_exec=(self.args.no_exec or bool(self.args.checkpoints)), force=self.args.force, run_cache=not self.args.no_run_cache, no_commit=self.args.no_commit, outs_persist=self.args.outs_persist, outs_persist_no_cache=self.args.outs_persist_no_cache, + checkpoints=self.args.checkpoints, always_changed=self.args.always_changed, name=self.args.name, single_stage=self.args.single_stage, @@ -212,6 +214,14 @@ def add_parser(subparsers, parent_parser): "removed upon repro (do not put into DVC cache).", metavar="", ) + run_parser.add_argument( + "-c", + "--checkpoints", + action="append", + default=[], + help=argparse.SUPPRESS, + metavar="", + ).complete = completion.FILE run_parser.add_argument( "--always-changed", action="store_true", diff --git a/dvc/output/__init__.py b/dvc/output/__init__.py index 389860b46..0ab8920f2 100644 --- a/dvc/output/__init__.py +++ b/dvc/output/__init__.py @@ -60,12 +60,20 @@ SCHEMA[BaseOutput.PARAM_CACHE] = bool SCHEMA[BaseOutput.PARAM_METRIC] = BaseOutput.METRIC_SCHEMA SCHEMA[BaseOutput.PARAM_PLOT] = bool SCHEMA[BaseOutput.PARAM_PERSIST] = bool +SCHEMA[BaseOutput.PARAM_CHECKPOINT] = bool SCHEMA[HashInfo.PARAM_SIZE] = int SCHEMA[HashInfo.PARAM_NFILES] = int def _get( - stage, p, info=None, cache=True, metric=False, plot=False, persist=False + stage, + p, + info=None, + cache=True, + metric=False, + plot=False, + persist=False, + checkpoint=False, ): parsed = urlparse(p) @@ -80,6 +88,7 @@ def _get( metric=metric, plot=plot, persist=persist, + checkpoint=checkpoint, ) for o in OUTS: @@ -93,6 +102,7 @@ def _get( metric=metric, plot=plot, persist=persist, + checkpoint=checkpoint, ) return LocalOutput( stage, @@ -103,6 +113,7 @@ def _get( metric=metric, plot=plot, persist=persist, + checkpoint=checkpoint, ) @@ -114,6 +125,7 @@ def loadd_from(stage, d_list): metric = d.pop(BaseOutput.PARAM_METRIC, False) plot = d.pop(BaseOutput.PARAM_PLOT, False) persist = d.pop(BaseOutput.PARAM_PERSIST, False) + checkpoint = d.pop(BaseOutput.PARAM_CHECKPOINT, False) ret.append( _get( stage, @@ -123,13 +135,20 @@ def loadd_from(stage, d_list): metric=metric, plot=plot, persist=persist, + checkpoint=checkpoint, ) ) return ret def loads_from( - stage, s_list, use_cache=True, metric=False, plot=False, persist=False + stage, + s_list, + use_cache=True, + metric=False, + plot=False, + persist=False, + checkpoint=False, ): return [ _get( @@ -140,6 +159,7 @@ def loads_from( metric=metric, plot=plot, persist=persist, + checkpoint=checkpoint, ) for s in s_list ] @@ -183,5 +203,12 @@ def load_from_pipeline(stage, s_list, typ="outs"): from dvc.schema import PLOT_PROPS plt_d, flags = _split_dict(flags, keys=PLOT_PROPS.keys()) - extra = project(flags, ["cache", "persist"]) + extra = project( + flags, + [ + BaseOutput.PARAM_CACHE, + BaseOutput.PARAM_PERSIST, + BaseOutput.PARAM_CHECKPOINT, + ], + ) yield _get(stage, path, {}, plot=plt_d or plot, metric=metric, **extra) diff --git a/dvc/output/base.py b/dvc/output/base.py index b70936f57..d8f87f6b0 100644 --- a/dvc/output/base.py +++ b/dvc/output/base.py @@ -61,6 +61,7 @@ class BaseOutput: PARAM_PATH = "path" PARAM_CACHE = "cache" + PARAM_CHECKPOINT = "checkpoint" PARAM_METRIC = "metric" PARAM_METRIC_TYPE = "type" PARAM_METRIC_XPATH = "xpath" @@ -100,6 +101,7 @@ class BaseOutput: metric=False, plot=False, persist=False, + checkpoint=False, ): self._validate_output_path(path, stage) # This output (and dependency) objects have too many paths/urls @@ -124,6 +126,7 @@ class BaseOutput: self.metric = False if self.IS_DEPENDENCY else metric self.plot = False if self.IS_DEPENDENCY else plot self.persist = persist + self.checkpoint = checkpoint self.path_info = self._parse_path(tree, path) if self.use_cache and self.cache is None: @@ -315,6 +318,9 @@ class BaseOutput: if self.persist: ret[self.PARAM_PERSIST] = self.persist + if self.checkpoint: + ret[self.PARAM_CHECKPOINT] = self.checkpoint + return ret def verify_metric(self): @@ -350,7 +356,7 @@ class BaseOutput: **kwargs, ) except CheckoutError: - if allow_missing: + if allow_missing or self.checkpoint: return None raise diff --git a/dvc/repo/reproduce.py b/dvc/repo/reproduce.py index 87b35b626..e7a61f612 100644 --- a/dvc/repo/reproduce.py +++ b/dvc/repo/reproduce.py @@ -13,11 +13,19 @@ logger = logging.getLogger(__name__) def _reproduce_stage(stage, **kwargs): def _run_callback(repro_callback): _dump_stage(stage) + logger.debug(f"{repro_callback} ([{stage}])") repro_callback([stage]) checkpoint_func = kwargs.pop("checkpoint_func", None) - if checkpoint_func: - kwargs["checkpoint_func"] = partial(_run_callback, checkpoint_func) + if stage.is_checkpoint: + if checkpoint_func: + kwargs["checkpoint_func"] = partial(_run_callback, checkpoint_func) + else: + logger.warning( + "Checkpoint stages are not fully supported in 'dvc repro'. " + "Checkpoint stages should be reproduced with 'dvc exp run' " + "or 'dvc exp resume'." + ) if stage.frozen and not stage.is_import: logger.warning( diff --git a/dvc/repo/run.py b/dvc/repo/run.py index c61c8142a..c7e4e49dd 100644 --- a/dvc/repo/run.py +++ b/dvc/repo/run.py @@ -48,6 +48,7 @@ def _get_file_path(kwargs): kwargs.get("plots_no_cache", []), kwargs.get("outs_persist", []), kwargs.get("outs_persist_no_cache", []), + kwargs.get("checkpoints", []), ) ) diff --git a/dvc/schema.py b/dvc/schema.py index 7c470edf2..b58d30dcd 100644 --- a/dvc/schema.py +++ b/dvc/schema.py @@ -34,7 +34,11 @@ LOCK_FILE_STAGE_SCHEMA = { LOCKFILE_SCHEMA = {str: LOCK_FILE_STAGE_SCHEMA} OUT_PSTAGE_DETAILED_SCHEMA = { - str: {BaseOutput.PARAM_CACHE: bool, BaseOutput.PARAM_PERSIST: bool} + str: { + BaseOutput.PARAM_CACHE: bool, + BaseOutput.PARAM_PERSIST: bool, + BaseOutput.PARAM_CHECKPOINT: bool, + } } PLOT_PROPS = { diff --git a/dvc/stage/__init__.py b/dvc/stage/__init__.py index 4b7b8328e..5ac54a5a4 100644 --- a/dvc/stage/__init__.py +++ b/dvc/stage/__init__.py @@ -216,6 +216,14 @@ class Stage(params.StageParams): return isinstance(self.deps[0], dependency.RepoDependency) + @property + def is_checkpoint(self): + """ + A stage containing checkpoint outs is always considered as changed + since the checkpoint out is a circular dependency. + """ + return any(out.checkpoint for out in self.outs) + def changed_deps(self): if self.frozen: return False @@ -229,7 +237,7 @@ class Stage(params.StageParams): ) return True - if self.always_changed: + if self.always_changed or self.is_checkpoint: return True return self._changed_deps() @@ -284,7 +292,7 @@ class Stage(params.StageParams): def remove_outs(self, ignore_remove=False, force=False): """Used mainly for `dvc remove --outs` and :func:`Stage.reproduce`.""" for out in self.outs: - if out.persist and not force: + if (out.persist or out.checkpoint) and not force: out.unprotect() continue @@ -402,7 +410,7 @@ class Stage(params.StageParams): try: out.save() except OutputDoesNotExistError: - if not allow_missing: + if not (allow_missing or out.checkpoint): raise def ignore_outs(self): @@ -432,7 +440,7 @@ class Stage(params.StageParams): try: out.commit() except OutputDoesNotExistError: - if not allow_missing: + if not (allow_missing or out.checkpoint): raise except CacheLinkError: link_failures.append(out.path_info) @@ -527,7 +535,7 @@ class Stage(params.StageParams): ret.append({"changed outs": outs_status}) def _status_always_changed(self, ret): - if self.is_callback or self.always_changed: + if self.is_callback or self.always_changed or self.is_checkpoint: ret.append("always changed") def _status_stage(self, ret): diff --git a/dvc/stage/serialize.py b/dvc/stage/serialize.py index 46c2320ba..fef9edb30 100644 --- a/dvc/stage/serialize.py +++ b/dvc/stage/serialize.py @@ -26,6 +26,7 @@ PARAM_CACHE = BaseOutput.PARAM_CACHE PARAM_METRIC = BaseOutput.PARAM_METRIC PARAM_PLOT = BaseOutput.PARAM_PLOT PARAM_PERSIST = BaseOutput.PARAM_PERSIST +PARAM_CHECKPOINT = BaseOutput.PARAM_CHECKPOINT DEFAULT_PARAMS_FILE = ParamsDependency.DEFAULT_PARAMS_FILE @@ -37,6 +38,8 @@ sort_by_path = partial(sorted, key=attrgetter("def_path")) def _get_flags(out): if not out.use_cache: yield PARAM_CACHE, False + if out.checkpoint: + yield PARAM_CHECKPOINT, True if out.persist: yield PARAM_PERSIST, True if out.plot and isinstance(out.plot, dict): diff --git a/dvc/stage/utils.py b/dvc/stage/utils.py index ff94d83e4..134305293 100644 --- a/dvc/stage/utils.py +++ b/dvc/stage/utils.py @@ -53,6 +53,7 @@ def fill_stage_outputs(stage, **kwargs): "plots", "outs_no_cache", "outs", + "checkpoints", ] stage.outs = [] @@ -64,6 +65,7 @@ def fill_stage_outputs(stage, **kwargs): persist="persist" in key, metric="metrics" in key, plot="plots" in key, + checkpoint="checkpoints" in key, ) @@ -174,6 +176,7 @@ def compute_md5(stage): stage.PARAM_FROZEN, BaseOutput.PARAM_METRIC, BaseOutput.PARAM_PERSIST, + BaseOutput.PARAM_CHECKPOINT, HashInfo.PARAM_SIZE, HashInfo.PARAM_NFILES, ], diff --git a/tests/func/experiments/test_experiments.py b/tests/func/experiments/test_experiments.py index 8d57aab91..9beed5ea1 100644 --- a/tests/func/experiments/test_experiments.py +++ b/tests/func/experiments/test_experiments.py @@ -320,19 +320,11 @@ def test_new_checkpoint(tmp_dir, scm, dvc, mocker): cmd="python checkpoint.py foo 5 params.yaml metrics.yaml", metrics_no_cache=["metrics.yaml"], params=["foo"], - outs_persist=["foo"], - always_changed=True, + checkpoints=["foo"], + no_exec=True, name="checkpoint-file", ) - scm.add( - [ - "dvc.yaml", - "dvc.lock", - "checkpoint.py", - "params.yaml", - "metrics.yaml", - ] - ) + scm.add(["dvc.yaml", "checkpoint.py", "params.yaml"]) scm.commit("init") new_mock = mocker.spy(dvc.experiments, "new") @@ -353,19 +345,11 @@ def test_continue_checkpoint(tmp_dir, scm, dvc, mocker, last): cmd="python checkpoint.py foo 5 params.yaml metrics.yaml", metrics_no_cache=["metrics.yaml"], params=["foo"], - outs_persist=["foo"], - always_changed=True, + checkpoints=["foo"], + no_exec=True, name="checkpoint-file", ) - scm.add( - [ - "dvc.yaml", - "dvc.lock", - "checkpoint.py", - "params.yaml", - "metrics.yaml", - ] - ) + scm.add(["dvc.yaml", "checkpoint.py", "params.yaml"]) scm.commit("init") results = dvc.experiments.run( @@ -395,19 +379,11 @@ def test_reset_checkpoint(tmp_dir, scm, dvc, mocker): cmd="python checkpoint.py foo 5 params.yaml metrics.yaml", metrics_no_cache=["metrics.yaml"], params=["foo"], - outs_persist=["foo"], - always_changed=True, + checkpoints=["foo"], + no_exec=True, name="checkpoint-file", ) - scm.add( - [ - "dvc.yaml", - "dvc.lock", - "checkpoint.py", - "params.yaml", - "metrics.yaml", - ] - ) + scm.add(["dvc.yaml", "checkpoint.py", "params.yaml"]) scm.commit("init") dvc.experiments.run(stage.addressing, checkpoint=True) diff --git a/tests/unit/command/test_run.py b/tests/unit/command/test_run.py index ba718e17b..e0f82d1e3 100644 --- a/tests/unit/command/test_run.py +++ b/tests/unit/command/test_run.py @@ -34,6 +34,8 @@ def test_run(mocker, dvc): "outs-persist", "--outs-persist-no-cache", "outs-persist-no-cache", + "--checkpoints", + "checkpoints", "--always-changed", "--params", "file:param1,param2", @@ -60,6 +62,7 @@ def test_run(mocker, dvc): plots_no_cache=["plots-no-cache"], outs_persist=["outs-persist"], outs_persist_no_cache=["outs-persist-no-cache"], + checkpoints=["checkpoints"], params=["file:param1,param2", "param3"], fname="file", wdir="wdir", @@ -90,6 +93,7 @@ def test_run_args_from_cli(mocker, dvc): plots_no_cache=[], outs_persist=[], outs_persist_no_cache=[], + checkpoints=[], params=[], fname=None, wdir=None, @@ -120,6 +124,7 @@ def test_run_args_with_spaces(mocker, dvc): plots_no_cache=[], outs_persist=[], outs_persist_no_cache=[], + checkpoints=[], params=[], fname=None, wdir=None,