exp queue: add separate standalone/tempdir implementation

dvc/repo/experiments/__init__.py

@@ -19,6 +19,7 @@ from .exceptions import (
 from .executor.base import BaseExecutor, ExecutorInfo
 from .queue.base import BaseStashQueue, QueueEntry
 from .queue.celery import LocalCeleryQueue
+from .queue.tempdir import TempDirQueue
 from .queue.workspace import WorkspaceQueue
 from .refs import (
     CELERY_FAILED_STASH,
@@ -79,6 +80,12 @@ class Experiments:
     def workspace_queue(self) -> WorkspaceQueue:
         return WorkspaceQueue(self.repo, WORKSPACE_STASH)
 
+    @cached_property
+    def tempdir_queue(self) -> TempDirQueue:
+        # NOTE: tempdir and workspace stash is shared since both
+        # implementations immediately push -> pop (queue length is only 0 or 1)
+        return TempDirQueue(self.repo, WORKSPACE_STASH)
+
     @cached_property
     def celery_queue(self) -> LocalCeleryQueue:
         return LocalCeleryQueue(self.repo, CELERY_STASH, CELERY_FAILED_STASH)

dvc/repo/experiments/executor/local.py

@@ -120,9 +120,10 @@ class TempDirExecutor(BaseLocalExecutor):
         repo: "Repo",
         stash_rev: str,
         entry: "ExpStashEntry",
+        wdir: Optional[str] = None,
         **kwargs,
     ):
-        tmp_dir = mkdtemp(dir=os.path.join(repo.tmp_dir, EXEC_TMP_DIR))
+        tmp_dir = mkdtemp(dir=wdir or os.path.join(repo.tmp_dir, EXEC_TMP_DIR))
         try:
             executor = cls._from_stash_entry(repo, stash_rev, entry, tmp_dir)
             logger.debug("Init temp dir executor in '%s'", tmp_dir)

dvc/repo/experiments/queue/base.py

@@ -542,6 +542,7 @@ class BaseStashQueue(ABC):
         exp: "Experiments",
         queue_entry: QueueEntry,
         executor_cls: Type[BaseExecutor] = WorkspaceExecutor,
+        **kwargs,
     ) -> BaseExecutor:
         scm = exp.scm
         stash = ExpStash(scm, queue_entry.stash_ref)
@@ -563,7 +564,9 @@ class BaseStashQueue(ABC):
         #   EXEC_MERGE - the unmerged changes (from our stash)
         #       to be reproduced
         #   EXEC_BASELINE - the baseline commit for this experiment
-        return executor_cls.from_stash_entry(exp.repo, stash_rev, stash_entry)
+        return executor_cls.from_stash_entry(
+            exp.repo, stash_rev, stash_entry, **kwargs
+        )
 
     def get_infofile_path(self, name: str) -> str:
         return os.path.join(

dvc/repo/experiments/queue/celery.py

@@ -22,12 +22,7 @@ from dvc.exceptions import DvcException
 from dvc.ui import ui
 
 from ..exceptions import UnresolvedQueueExpNamesError
-from ..executor.base import (
-    EXEC_PID_DIR,
-    EXEC_TMP_DIR,
-    ExecutorInfo,
-    ExecutorResult,
-)
+from ..executor.base import EXEC_TMP_DIR, ExecutorInfo, ExecutorResult
 from ..stash import ExpStashEntry
 from .base import BaseStashQueue, QueueDoneResult, QueueEntry, QueueGetResult
 from .tasks import run_exp
@@ -89,8 +84,7 @@ class LocalCeleryQueue(BaseStashQueue):
     def proc(self) -> "ProcessManager":
         from dvc_task.proc.manager import ProcessManager
 
-        pid_dir = os.path.join(self.repo.tmp_dir, EXEC_TMP_DIR, EXEC_PID_DIR)
-        return ProcessManager(pid_dir)
+        return ProcessManager(self.pid_dir)
 
     @cached_property
     def worker(self) -> "TemporaryWorker":

dvc/repo/experiments/queue/tasks.py

@@ -62,11 +62,9 @@ def collect_exp(
     exec_result = executor_info.result
     try:
         if exec_result is not None:
-            results = BaseStashQueue.collect_executor(
+            BaseStashQueue.collect_executor(
                 repo.experiments, executor, exec_result
             )
-            for rev in results:
-                logger.debug("Collected experiment '%s'", rev[:7])
         else:
             logger.debug("Experiment failed (Exec result was None)")
             celery_queue.stash_failed(entry)

dvc/repo/experiments/queue/tempdir.py

@@ -0,0 +1,93 @@
+import logging
+import os
+from typing import TYPE_CHECKING, Dict, Generator
+
+from funcy import cached_property, first
+
+from ..exceptions import ExpQueueEmptyError
+from ..executor.base import (
+    EXEC_PID_DIR,
+    EXEC_TMP_DIR,
+    BaseExecutor,
+    ExecutorInfo,
+    ExecutorResult,
+)
+from ..executor.local import TempDirExecutor
+from .base import BaseStashQueue, QueueEntry, QueueGetResult
+from .workspace import WorkspaceQueue
+
+if TYPE_CHECKING:
+    from dvc.repo.experiments import Experiments
+    from dvc_task.proc.manager import ProcessManager
+
+logger = logging.getLogger(__name__)
+
+
+_STANDALONE_TMP_DIR = os.path.join(EXEC_TMP_DIR, "standalone")
+
+
+class TempDirQueue(WorkspaceQueue):
+    """Standalone/tempdir exp queue implementation."""
+
+    @cached_property
+    def _standalone_tmp_dir(self) -> str:
+        return os.path.join(self.repo.tmp_dir, _STANDALONE_TMP_DIR)
+
+    @cached_property
+    def pid_dir(self) -> str:
+        return os.path.join(self._standalone_tmp_dir, EXEC_PID_DIR)
+
+    @cached_property
+    def proc(self) -> "ProcessManager":
+        from dvc_task.proc.manager import ProcessManager
+
+        return ProcessManager(self.pid_dir)
+
+    def get(self) -> QueueGetResult:
+        revs = self.stash.stash_revs
+        if not revs:
+            raise ExpQueueEmptyError("No stashed standalone experiments.")
+        stash_rev, stash_entry = first(revs.items())
+        entry = QueueEntry(
+            self.repo.root_dir,
+            self.scm.root_dir,
+            self.ref,
+            stash_rev,
+            stash_entry.baseline_rev,
+            stash_entry.branch,
+            stash_entry.name,
+            stash_entry.head_rev,
+        )
+        executor = self.setup_executor(
+            self.repo.experiments,
+            entry,
+            TempDirExecutor,
+            wdir=self._standalone_tmp_dir,
+        )
+        return QueueGetResult(entry, executor)
+
+    def iter_active(self) -> Generator[QueueEntry, None, None]:
+        # NOTE: Yielded queue entries are not complete for performance reasons.
+        # Retrieving exec ref information is unavailable without doing a
+        # git-fetch, and is unneeded in the common use cases for iter_active.
+        for stash_rev in self.proc:
+            infofile = self.get_infofile_path(stash_rev)
+            executor_info = ExecutorInfo.load_json(infofile)
+            yield QueueEntry(
+                self.repo.root_dir,
+                self.scm.root_dir,
+                self.ref,
+                stash_rev,
+                executor_info.baseline_rev,
+                None,  # branch unavailable without doing a git-fetch
+                executor_info.name,
+                None,
+            )
+
+    @staticmethod
+    def collect_executor(
+        exp: "Experiments",
+        executor: BaseExecutor,
+        exec_result: ExecutorResult,
+    ) -> Dict[str, str]:
+        return BaseStashQueue.collect_executor(exp, executor, exec_result)

tests/func/experiments/test_remove.py

@@ -34,12 +34,10 @@ def test_remove_experiments_by_ref(tmp_dir, scm, dvc, exp_stage, caplog):
 
 def test_remove_all_queued_experiments(tmp_dir, scm, dvc, exp_stage):
     queue_length = 3
-    ref_list = []
     for i in range(queue_length):
-        result = dvc.experiments.run(
+        dvc.experiments.run(
             exp_stage.addressing, params=[f"foo={i}"], queue=True
         )
-        ref_list.append(first(result)[:7])
 
     results = dvc.experiments.run(
         exp_stage.addressing, params=[f"foo={queue_length}"]
@@ -47,41 +45,40 @@ def test_remove_all_queued_experiments(tmp_dir, scm, dvc, exp_stage):
     ref_info = first(exp_refs_by_rev(scm, first(results)))
 
     assert len(dvc.experiments.stash_revs) == queue_length
-    assert set(dvc.experiments.remove(queue=True)) == set(ref_list)
+    assert len(dvc.experiments.remove(queue=True)) == queue_length
     assert len(dvc.experiments.stash_revs) == 0
     assert scm.get_ref(str(ref_info)) is not None
 
 
 def test_remove_special_queued_experiments(tmp_dir, scm, dvc, exp_stage):
-    results = dvc.experiments.run(
+    dvc.experiments.run(
         exp_stage.addressing, params=["foo=1"], queue=True, name="queue1"
     )
-    rev1 = first(results)
-    results = dvc.experiments.run(
+    dvc.experiments.run(
         exp_stage.addressing, params=["foo=2"], queue=True, name="queue2"
     )
-    rev2 = first(results)
-    results = dvc.experiments.run(
+    dvc.experiments.run(
         exp_stage.addressing, params=["foo=3"], queue=True, name="queue3"
     )
-    rev3 = first(results)
+    queue_revs = {
+        entry.name: entry.stash_rev
+        for entry in dvc.experiments.celery_queue.iter_queued()
+    }
+    assert len(queue_revs) == 3
+
     results = dvc.experiments.run(exp_stage.addressing, params=["foo=4"])
     ref_info1 = first(exp_refs_by_rev(scm, first(results)))
     results = dvc.experiments.run(exp_stage.addressing, params=["foo=5"])
     ref_info2 = first(exp_refs_by_rev(scm, first(results)))
 
-    assert rev1 in dvc.experiments.stash_revs
-    assert rev2 in dvc.experiments.stash_revs
-    assert rev3 in dvc.experiments.stash_revs
     assert scm.get_ref(str(ref_info1)) is not None
     assert scm.get_ref(str(ref_info2)) is not None
 
+    rev2 = queue_revs["queue2"]
     assert set(
         dvc.experiments.remove(["queue1", rev2[:5], str(ref_info1)])
     ) == {"queue1", rev2[:5], str(ref_info1)}
-    assert rev1 not in dvc.experiments.stash_revs
-    assert rev2 not in dvc.experiments.stash_revs
-    assert rev3 in dvc.experiments.stash_revs
+    assert len(list(dvc.experiments.celery_queue.iter_queued())) == 1
     assert scm.get_ref(str(ref_info1)) is None
     assert scm.get_ref(str(ref_info2)) is not None
 
@@ -141,23 +138,25 @@ def test_remove_experiments_by_rev(tmp_dir, scm, dvc, exp_stage):
     results = dvc.experiments.run(exp_stage.addressing, params=["foo=1"])
     baseline_exp_ref = first(exp_refs_by_rev(scm, first(results)))
 
-    results = dvc.experiments.run(
+    dvc.experiments.run(
         exp_stage.addressing, params=["foo=2"], queue=True, name="queue2"
     )
-    baseline_queue_ref = first(results)
-
     scm.commit("new_baseline")
 
     results = dvc.experiments.run(exp_stage.addressing, params=["foo=3"])
     ref_info = first(exp_refs_by_rev(scm, first(results)))
     new_exp_ref = str(ref_info)
 
-    results = dvc.experiments.run(
+    dvc.experiments.run(
         exp_stage.addressing, params=["foo=4"], queue=True, name="queue4"
     )
-    new_queue_ref = first(results)
+
     assert dvc.experiments.remove(rev=baseline) == [baseline_exp_ref.name]
+    queue_revs = {
+        entry.name: entry.stash_rev
+        for entry in dvc.experiments.celery_queue.iter_queued()
+    }
     assert scm.get_ref(str(baseline_exp_ref)) is None
-    assert baseline_queue_ref in dvc.experiments.stash_revs
+    assert "queue2" in queue_revs
     assert scm.get_ref(new_exp_ref) is not None
-    assert new_queue_ref in dvc.experiments.stash_revs
+    assert "queue4" in queue_revs