From b9f547637ccd9b61e5336113f2a612c3e2eb9811 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: Mon, 19 Oct 2020 15:25:08 +0900 Subject: [PATCH] checkpoints: add `--reset` option (#4731) * experiments: add --reset option to reset/restart checkpoint experiments * add CheckpointExistsError internal exception * update tests --- dvc/command/experiments.py | 22 +++++++- dvc/repo/experiments/__init__.py | 61 +++++++++++++++++----- dvc/repo/experiments/executor.py | 8 ++- tests/func/experiments/test_experiments.py | 41 +++++++++++++++ tests/unit/command/test_experiments.py | 1 + 5 files changed, 118 insertions(+), 15 deletions(-) diff --git a/dvc/command/experiments.py b/dvc/command/experiments.py index bdb9226f5..930ae3668 100644 --- a/dvc/command/experiments.py +++ b/dvc/command/experiments.py @@ -413,6 +413,14 @@ class CmdExperimentsRun(CmdRepro): elif not self.args.targets: self.args.targets = self.default_targets + if ( + self.args.checkpoint_reset + and self.args.checkpoint_continue is not None + ): + raise InvalidArgumentError( + "--continue and --reset cannot be used together" + ) + ret = 0 for target in self.args.targets: try: @@ -425,8 +433,10 @@ class CmdExperimentsRun(CmdRepro): checkpoint=( self.args.checkpoint or self.args.checkpoint_continue is not None + or self.args.checkpoint_reset ), checkpoint_continue=self.args.checkpoint_continue, + checkpoint_reset=self.args.checkpoint_reset, **self._repro_kwargs, ) except DvcException: @@ -674,7 +684,17 @@ def add_parser(subparsers, parent_parser): default=None, dest="checkpoint_continue", help=( - "Continue from the specified checkpoint experiment" + "Continue from the specified checkpoint experiment " + "(implies --checkpoint)." + ), + ) + experiments_run_parser.add_argument( + "--reset", + action="store_true", + default=False, + dest="checkpoint_reset", + help=( + "Reset checkpoint experiment if it already exists " "(implies --checkpoint)." ), ) diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index b67bdbef4..cb316ae9b 100644 --- a/dvc/repo/experiments/__init__.py +++ b/dvc/repo/experiments/__init__.py @@ -66,6 +66,19 @@ class BaselineMismatchError(DvcException): self.expected_rev = expected +class CheckpointExistsError(DvcException): + def __init__(self, rev, continue_rev): + msg = ( + f"Checkpoint experiment containing '{rev[:7]}' already exists." + " To restart the experiment run:\n\n" + "\tdvc exp run --reset ...\n\n" + "To resume the experiment, run:\n\n" + f"\tdvc exp run --continue {continue_rev[:7]}\n" + ) + super().__init__(msg) + self.rev = rev + + class MultipleBranchError(DvcException): def __init__(self, rev): super().__init__( @@ -325,7 +338,12 @@ class Experiments: self.scm.add(list(params.keys())) def _commit( - self, exp_hash, check_exists=True, create_branch=True, checkpoint=False + self, + exp_hash, + check_exists=True, + create_branch=True, + checkpoint=False, + checkpoint_reset=False, ): """Commit stages as an experiment and return the commit SHA.""" if not self.scm.is_dirty(untracked_files=True): @@ -340,14 +358,14 @@ class Experiments: ) and exp_name in self.scm.list_branches(): branch_tip = self.scm.resolve_rev(exp_name) if checkpoint: - raise DvcException( - f"Checkpoint experiment containing '{rev[:7]}' " - "already exists. To resume the experiment, " - "run:\n\n" - f"\tdvc exp run --continue {branch_tip[:7]}" + self._reset_checkpoint_branch( + exp_name, rev, branch_tip, checkpoint_reset ) - logger.debug("Using existing experiment branch '%s'", exp_name) - return branch_tip + else: + logger.debug( + "Using existing experiment branch '%s'", exp_name + ) + return branch_tip self.scm.checkout(exp_name, create_new=True) logger.debug("Commit new experiment branch '%s'", exp_name) else: @@ -356,6 +374,13 @@ class Experiments: self.scm.commit(f"Add experiment {exp_name}") return self.scm.get_rev() + def _reset_checkpoint_branch(self, branch, rev, branch_tip, reset): + if not reset: + raise CheckpointExistsError(rev, branch_tip) + self._checkout_default_branch() + logger.debug("Removing existing checkpoint branch '%s'", branch) + self.scm.repo.git.branch(branch, D=True) + def reproduce_one(self, queue=False, **kwargs): """Reproduce and checkout a single experiment.""" checkpoint = kwargs.get("checkpoint", False) @@ -391,6 +416,7 @@ class Experiments: *args, checkpoint: Optional[bool] = False, checkpoint_continue: Optional[str] = None, + checkpoint_reset: Optional[bool] = False, branch: Optional[str] = None, **kwargs, ): @@ -423,7 +449,11 @@ class Experiments: try: stash_rev = self._stash_exp( - *args, branch=branch, allow_unchanged=checkpoint, **kwargs + *args, + branch=branch, + allow_unchanged=checkpoint, + checkpoint_reset=checkpoint_reset, + **kwargs, ) except UnchangedExperimentError as exc: logger.info("Reproducing existing experiment '%s'.", rev[:7]) @@ -477,14 +507,16 @@ class Experiments: for rev, item in to_run.items(): self._scm_checkout(item.baseline_rev) self.scm.repo.git.stash("apply", rev) - repro_args, repro_kwargs = self._unpack_args() + packed_args, packed_kwargs = self._unpack_args() + checkpoint_reset = packed_kwargs.pop("checkpoint_reset", False) executor = LocalExecutor( item.baseline_rev, branch=item.branch, - repro_args=repro_args, - repro_kwargs=repro_kwargs, + repro_args=packed_args, + repro_kwargs=packed_kwargs, dvc_dir=self.dvc_dir, cache_dir=self.repo.cache.local.cache_dir, + checkpoint_reset=checkpoint_reset, ) self._collect_input(executor) executors[rev] = executor @@ -615,7 +647,10 @@ class Experiments: try: create_branch = not executor.branch exp_rev = self._commit( - exp_hash, create_branch=create_branch, **kwargs + exp_hash, + create_branch=create_branch, + checkpoint_reset=executor.checkpoint_reset, + **kwargs, ) except UnchangedExperimentError: logger.debug( diff --git a/dvc/repo/experiments/executor.py b/dvc/repo/experiments/executor.py index 3c1122f51..54c8b583f 100644 --- a/dvc/repo/experiments/executor.py +++ b/dvc/repo/experiments/executor.py @@ -83,7 +83,12 @@ class ExperimentExecutor: class LocalExecutor(ExperimentExecutor): """Local machine experiment executor.""" - def __init__(self, baseline_rev: str, **kwargs): + def __init__( + self, + baseline_rev: str, + checkpoint_reset: Optional[bool] = False, + **kwargs, + ): from dvc.repo import Repo dvc_dir = kwargs.pop("dvc_dir") @@ -104,6 +109,7 @@ class LocalExecutor(ExperimentExecutor): # override default CACHE_MODE since files must be writable in order # to run repro self._tree.CACHE_MODE = 0o644 + self.checkpoint_reset = checkpoint_reset def _config(self, cache_dir): local_config = os.path.join(self.dvc_dir, "config.local") diff --git a/tests/func/experiments/test_experiments.py b/tests/func/experiments/test_experiments.py index 545dffb63..f114dee9e 100644 --- a/tests/func/experiments/test_experiments.py +++ b/tests/func/experiments/test_experiments.py @@ -336,3 +336,44 @@ def test_continue_checkpoint(tmp_dir, scm, dvc, mocker): assert ( tmp_dir / ".dvc" / "experiments" / "metrics.yaml" ).read_text().strip() == "foo: 2" + + +def test_reset_checkpoint(tmp_dir, scm, dvc, mocker): + from dvc.exceptions import ReproductionError + + tmp_dir.gen("checkpoint.py", CHECKPOINT_SCRIPT) + tmp_dir.gen("params.yaml", "foo: 1") + stage = dvc.run( + cmd="python checkpoint.py foo 5 params.yaml metrics.yaml", + metrics_no_cache=["metrics.yaml"], + params=["foo"], + outs_persist=["foo"], + always_changed=True, + name="checkpoint-file", + ) + scm.add( + [ + "dvc.yaml", + "dvc.lock", + "checkpoint.py", + "params.yaml", + "metrics.yaml", + ] + ) + scm.commit("init") + + dvc.experiments.run(stage.addressing, checkpoint=True) + scm.repo.git.reset(hard=True) + scm.repo.git.clean(force=True) + + with pytest.raises(ReproductionError): + dvc.experiments.run(stage.addressing, checkpoint=True) + + dvc.experiments.run( + stage.addressing, checkpoint=True, checkpoint_reset=True + ) + + assert (tmp_dir / "foo").read_text() == "5" + assert ( + tmp_dir / ".dvc" / "experiments" / "metrics.yaml" + ).read_text().strip() == "foo: 1" diff --git a/tests/unit/command/test_experiments.py b/tests/unit/command/test_experiments.py index 923342e9b..4e799bc61 100644 --- a/tests/unit/command/test_experiments.py +++ b/tests/unit/command/test_experiments.py @@ -71,6 +71,7 @@ def test_experiments_run(dvc, mocker): "jobs": None, "checkpoint": False, "checkpoint_continue": None, + "checkpoint_reset": False, } default_arguments.update(repro_arguments)