From 98cd479876838bbb22f8140fb10af27f1eca93fe Mon Sep 17 00:00:00 2001
From: Peter Rowlands <peter@pmrowla.com>
Date: Wed, 10 May 2023 17:23:22 +0900
Subject: [PATCH] exp: add contextmanager for switching to active tempdir repo

(cherry picked from commit 8cb26460234d6ec3fc799cb48a58b81366fdc1ed)
---
 dvc/repo/experiments/brancher.py   | 43 ++++++++++++++++++++++++++++++
 dvc/repo/experiments/exceptions.py |  7 +++++
 dvc/repo/experiments/queue/base.py | 23 ++++++++++++++++
 3 files changed, 73 insertions(+)
 create mode 100644 dvc/repo/experiments/brancher.py

diff --git a/dvc/repo/experiments/brancher.py b/dvc/repo/experiments/brancher.py
new file mode 100644
index 000000000..41b3cbbfd
--- /dev/null
+++ b/dvc/repo/experiments/brancher.py
@@ -0,0 +1,43 @@
+from contextlib import ExitStack, contextmanager
+from typing import TYPE_CHECKING, Iterator, Tuple
+
+from dvc.repo.experiments.exceptions import InvalidExpRevError
+from dvc.scm import RevError
+
+if TYPE_CHECKING:
+    from dvc.repo import Repo
+
+
+@contextmanager
+def switch_repo(
+    repo: "Repo",
+    rev: str,
+) -> Iterator[Tuple["Repo", str]]:
+    """Return a repo instance (brancher) switched to rev.
+
+    If rev is the name of a running experiment, the returned instance will be
+    the live repo wherever the experiment is running.
+
+    NOTE: This will not resolve git SHA's that only exist in queued exp workspaces
+    (it will only match queued exp names).
+    """
+    try:
+        with repo.switch(rev):
+            yield repo, rev
+        return
+    except RevError as exc:
+        orig_exc = exc
+    exps = repo.experiments
+
+    for queue in (exps.tempdir_queue, exps.celery_queue):
+        try:
+            active_repo = queue.active_repo(rev)
+        except InvalidExpRevError:
+            continue
+        stack = ExitStack()
+        stack.enter_context(active_repo)
+        stack.enter_context(active_repo.switch("workspace"))
+        with stack:
+            yield active_repo, rev
+        return
+    raise orig_exc
diff --git a/dvc/repo/experiments/exceptions.py b/dvc/repo/experiments/exceptions.py
index d83d4a000..96ec34974 100644
--- a/dvc/repo/experiments/exceptions.py
+++ b/dvc/repo/experiments/exceptions.py
@@ -114,3 +114,10 @@ class UnresolvedRunningExpNamesError(UnresolvedExpNamesError):
 
 class ExpQueueEmptyError(DvcException):
     pass
+
+
+class ExpNotStartedError(DvcException):
+    def __init__(self, name: str):
+        super().__init__(
+            f"Queued experiment '{name}' exists but has not started running yet"
+        )
diff --git a/dvc/repo/experiments/queue/base.py b/dvc/repo/experiments/queue/base.py
index 0df531d69..da6ebcda3 100644
--- a/dvc/repo/experiments/queue/base.py
+++ b/dvc/repo/experiments/queue/base.py
@@ -790,3 +790,26 @@ class BaseStashQueue(ABC):
         Returns:
             Dict mapping baseline revision to list of queued experiments.
         """
+
+    def active_repo(self, name: str) -> "Repo":
+        """Return a Repo for the specified active experiment if it exists."""
+        from dvc.repo import Repo
+        from dvc.repo.experiments.exceptions import (
+            ExpNotStartedError,
+            InvalidExpRevError,
+        )
+        from dvc.repo.experiments.executor.base import ExecutorInfo, TaskStatus
+
+        for entry in self.iter_active():
+            if entry.name != name:
+                continue
+            infofile = self.get_infofile_path(entry.stash_rev)
+            executor_info = ExecutorInfo.load_json(infofile)
+            if executor_info.status < TaskStatus.RUNNING:
+                raise ExpNotStartedError(name)
+            dvc_root = os.path.join(executor_info.root_dir, executor_info.dvc_dir)
+            try:
+                return Repo(dvc_root)
+            except (FileNotFoundError, DvcException) as exc:
+                raise InvalidExpRevError(name) from exc
+        raise InvalidExpRevError(name)