From f386f7c082717f986e2bbd8b53e05eb87616ccb2 Mon Sep 17 00:00:00 2001
From: karajan1001 <mishanyo1001@gmail.com>
Date: Mon, 19 Dec 2022 18:59:00 +0800
Subject: [PATCH] `exp apply` can accept exp names and failed exps as target

fix: #8182

1. `exp apply` can apply to failed experiments.
2. `exp apply` can accept exp names as target.
3. Add a new functional test for it.
---
 dvc/repo/experiments/apply.py              | 34 +++++++++++++++-------
 tests/func/experiments/test_experiments.py | 31 +++++++++++++++++---
 2 files changed, 51 insertions(+), 14 deletions(-)

diff --git a/dvc/repo/experiments/apply.py b/dvc/repo/experiments/apply.py
index bdf923174..6eb2920c2 100644
--- a/dvc/repo/experiments/apply.py
+++ b/dvc/repo/experiments/apply.py
@@ -16,9 +16,10 @@ from .executor.base import BaseExecutor
 from .refs import EXEC_APPLY
 
 if TYPE_CHECKING:
-    from scmrepo import Git
+    from scmrepo.git import Git
 
     from dvc.repo import Repo
+    from dvc.repo.experiments import Experiments
 
 logger = logging.getLogger(__name__)
 
@@ -31,20 +32,33 @@ def apply(repo: "Repo", rev: str, force: bool = True, **kwargs):
     from dvc.repo.checkout import checkout as dvc_checkout
     from dvc.scm import GitMergeError, RevError, resolve_rev
 
-    exps = repo.experiments
+    exps: "Experiments" = repo.experiments
+
+    is_stash: bool = False
 
     try:
         exp_rev = resolve_rev(repo.scm, rev)
         exps.check_baseline(exp_rev)
-    except (RevError, BaselineMismatchError) as exc:
+    except RevError as exc:
+        (
+            exp_ref_info,
+            queue_entry,
+        ) = exps.celery_queue.get_ref_and_entry_by_names(rev)[rev]
+        baseline_sha = repo.scm.get_rev()
+        if exp_ref_info:
+            if baseline_sha != exp_ref_info.baseline_sha:
+                raise InvalidExpRevError(rev)
+            exp_rev = repo.scm.get_ref(str(exp_ref_info))
+        elif queue_entry:
+            if queue_entry.baseline_rev != baseline_sha:
+                raise InvalidExpRevError(rev)
+            exp_rev = queue_entry.stash_rev
+            is_stash = True
+        else:
+            raise InvalidExpRevError(rev) from exc
+    except BaselineMismatchError as exc:
         raise InvalidExpRevError(rev) from exc
 
-    stash_rev = exp_rev in exps.stash_revs
-    if not stash_rev and not exps.get_branch_by_rev(
-        exp_rev, allow_multiple=True
-    ):
-        raise InvalidExpRevError(exp_rev)
-
     # NOTE: we don't use scmrepo's stash_workspace() here since we need
     # finer control over the merge behavior when we unstash everything
     with _apply_workspace(repo, rev, force):
@@ -55,7 +69,7 @@ def apply(repo: "Repo", rev: str, force: bool = True, **kwargs):
 
     repo.scm.reset()
 
-    if stash_rev:
+    if is_stash:
         args_path = os.path.join(repo.tmp_dir, BaseExecutor.PACKED_ARGS_FILE)
         if os.path.exists(args_path):
             remove(args_path)
diff --git a/tests/func/experiments/test_experiments.py b/tests/func/experiments/test_experiments.py
index 4de2e232f..2f0568e15 100644
--- a/tests/func/experiments/test_experiments.py
+++ b/tests/func/experiments/test_experiments.py
@@ -11,6 +11,7 @@ from dvc.dvcfile import PIPELINE_FILE
 from dvc.exceptions import ReproductionError
 from dvc.repo.experiments.exceptions import ExperimentExistsError
 from dvc.repo.experiments.queue.base import BaseStashQueue
+from dvc.repo.experiments.refs import CELERY_STASH
 from dvc.repo.experiments.utils import exp_refs_by_rev
 from dvc.scm import resolve_rev
 from dvc.stage.exceptions import StageFileDoesNotExistError
@@ -114,12 +115,12 @@ def test_apply(tmp_dir, scm, dvc, exp_stage):
     exp_a = first(results)
 
     results = dvc.experiments.run(
-        exp_stage.addressing, params=["foo=3"], tmp_dir=True
+        exp_stage.addressing, params=["foo=3"], tmp_dir=True, name="foo"
     )
     exp_b = first(results)
 
     with pytest.raises(InvalidArgumentError):
-        dvc.experiments.apply("foo")
+        dvc.experiments.apply("bar")
 
     dvc.experiments.apply(exp_a)
     assert (tmp_dir / "params.yaml").read_text().strip() == "foo: 2"
@@ -131,11 +132,34 @@ def test_apply(tmp_dir, scm, dvc, exp_stage):
         assert (tmp_dir / "params.yaml").read_text().strip() == "foo: 2"
         assert (tmp_dir / "metrics.yaml").read_text().strip() == "foo: 2"
 
-    dvc.experiments.apply(exp_b)
+    dvc.experiments.apply("foo")
     assert (tmp_dir / "params.yaml").read_text().strip() == "foo: 3"
     assert (tmp_dir / "metrics.yaml").read_text().strip() == "foo: 3"
 
 
+def test_apply_failed(tmp_dir, scm, dvc, failed_exp_stage, test_queue):
+    from dvc.exceptions import InvalidArgumentError
+
+    dvc.experiments.run(
+        failed_exp_stage.addressing, params=["foo=2"], queue=True
+    )
+    exp_rev = dvc.experiments.scm.resolve_rev(f"{CELERY_STASH}@{{0}}")
+
+    dvc.experiments.run(
+        failed_exp_stage.addressing, params=["foo=3"], queue=True, name="foo"
+    )
+    dvc.experiments.run(run_all=True)
+
+    with pytest.raises(InvalidArgumentError):
+        dvc.experiments.apply("bar")
+
+    dvc.experiments.apply(exp_rev)
+    assert (tmp_dir / "params.yaml").read_text().strip() == "foo: 2"
+
+    dvc.experiments.apply("foo")
+    assert (tmp_dir / "params.yaml").read_text().strip() == "foo: 3"
+
+
 def test_apply_queued(tmp_dir, scm, dvc, exp_stage):
     from dvc.repo.experiments.exceptions import ApplyConflictError
 
@@ -183,7 +207,6 @@ def test_apply_untracked(tmp_dir, scm, dvc, exp_stage):
 
 
 def test_get_baseline(tmp_dir, scm, dvc, exp_stage):
-    from dvc.repo.experiments.refs import CELERY_STASH
 
     init_rev = scm.get_rev()
     assert dvc.experiments.get_baseline(init_rev) is None