diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index 8647d6f3c..43a1e4c76 100644 --- a/dvc/repo/experiments/__init__.py +++ b/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) diff --git a/dvc/repo/experiments/executor/local.py b/dvc/repo/experiments/executor/local.py index 778285f52..3490c1394 100644 --- a/dvc/repo/experiments/executor/local.py +++ b/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) diff --git a/dvc/repo/experiments/queue/base.py b/dvc/repo/experiments/queue/base.py index ea1f27530..a5caa9363 100644 --- a/dvc/repo/experiments/queue/base.py +++ b/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( diff --git a/dvc/repo/experiments/queue/celery.py b/dvc/repo/experiments/queue/celery.py index b8ef0ded8..b5e0356a5 100644 --- a/dvc/repo/experiments/queue/celery.py +++ b/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": diff --git a/dvc/repo/experiments/queue/tasks.py b/dvc/repo/experiments/queue/tasks.py index 273efcc17..9d0f758fb 100644 --- a/dvc/repo/experiments/queue/tasks.py +++ b/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) diff --git a/dvc/repo/experiments/queue/tempdir.py b/dvc/repo/experiments/queue/tempdir.py new file mode 100644 index 000000000..240cb39c3 --- /dev/null +++ b/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) diff --git a/tests/func/experiments/test_remove.py b/tests/func/experiments/test_remove.py index cee9f41ac..43fed7d72 100644 --- a/tests/func/experiments/test_remove.py +++ b/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