checkpoints: add --reset option (#4731)

* experiments: add --reset option to reset/restart checkpoint experiments

* add CheckpointExistsError internal exception

* update tests

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)."
         ),
     )

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(

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")

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"

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)