From 79232dfedf731038e1d7f62c6f247d00e6ec53a8 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?=
 <peter@pmrowla.com>
Date: Wed, 21 Oct 2020 17:25:16 +0900
Subject: [PATCH] checkpoints: `--continue` from last run experiment by default
 (#4758)

* checkpoints: `--continue` with no args resumes most recent experiment

* add test for --continue :last
---
 dvc/command/experiments.py                 |  6 +++-
 dvc/repo/experiments/__init__.py           | 32 +++++++++++++++++-----
 tests/func/experiments/test_experiments.py |  8 ++++--
 3 files changed, 36 insertions(+), 10 deletions(-)

diff --git a/dvc/command/experiments.py b/dvc/command/experiments.py
index ebe84c15d..b07c34051 100644
--- a/dvc/command/experiments.py
+++ b/dvc/command/experiments.py
@@ -710,12 +710,16 @@ def add_parser(subparsers, parent_parser):
     experiments_run_parser.add_argument(
         "--continue",
         type=str,
+        nargs="?",
         default=None,
+        const=":last",
         dest="checkpoint_continue",
         help=(
             "Continue from the specified checkpoint experiment "
-            "(implies --checkpoint)."
+            "(implies --checkpoint). If no experiment revision is provided, "
+            "the most recently run checkpoint experiment will be used."
         ),
+        metavar="<experiment_rev>",
     )
     experiments_run_parser.add_argument(
         "--reset",
diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py
index 0a455461e..e2468704c 100644
--- a/dvc/repo/experiments/__init__.py
+++ b/dvc/repo/experiments/__init__.py
@@ -426,13 +426,31 @@ class Experiments:
         workspace.
         """
         if checkpoint_continue:
-            rev = self.scm.resolve_rev(checkpoint_continue)
-            branch = self._get_branch_containing(rev)
-            if not branch:
-                raise DvcException(
-                    "Could not find checkpoint experiment "
-                    f"'{checkpoint_continue}'"
-                )
+            branch = None
+            if checkpoint_continue == ":last":
+                # Continue from most recently committed checkpoint
+                for head in sorted(
+                    self.scm.repo.heads,
+                    key=lambda h: h.commit.committed_date,
+                    reverse=True,
+                ):
+                    exp_branch = head.name
+                    m = self.BRANCH_RE.match(exp_branch)
+                    if m and m.group("checkpoint"):
+                        branch = exp_branch
+                        break
+                if not branch:
+                    raise DvcException(
+                        "No existing checkpoint experiment to continue"
+                    )
+            else:
+                rev = self.scm.resolve_rev(checkpoint_continue)
+                branch = self._get_branch_containing(rev)
+                if not branch:
+                    raise DvcException(
+                        "Could not find checkpoint experiment "
+                        f"'{checkpoint_continue}'"
+                    )
             logger.debug(
                 "Continuing checkpoint experiment '%s'", checkpoint_continue
             )
diff --git a/tests/func/experiments/test_experiments.py b/tests/func/experiments/test_experiments.py
index f114dee9e..afe970481 100644
--- a/tests/func/experiments/test_experiments.py
+++ b/tests/func/experiments/test_experiments.py
@@ -301,7 +301,8 @@ def test_new_checkpoint(tmp_dir, scm, dvc, mocker):
     ).read_text().strip() == "foo: 2"
 
 
-def test_continue_checkpoint(tmp_dir, scm, dvc, mocker):
+@pytest.mark.parametrize("last", [True, False])
+def test_continue_checkpoint(tmp_dir, scm, dvc, mocker, last):
     tmp_dir.gen("checkpoint.py", CHECKPOINT_SCRIPT)
     tmp_dir.gen("params.yaml", "foo: 1")
     stage = dvc.run(
@@ -326,7 +327,10 @@ def test_continue_checkpoint(tmp_dir, scm, dvc, mocker):
     results = dvc.experiments.run(
         stage.addressing, checkpoint=True, params=["foo=2"]
     )
-    exp_rev = first(results)
+    if last:
+        exp_rev = ":last"
+    else:
+        exp_rev = first(results)
 
     dvc.experiments.run(
         stage.addressing, checkpoint=True, checkpoint_continue=exp_rev,