From 1f7b802f7986f853ed01fd7cc8972fc33ead8d0b Mon Sep 17 00:00:00 2001 From: Peter Rowlands Date: Wed, 29 Jun 2022 15:43:47 +0900 Subject: [PATCH] exp show: collect running standalone tempdir experiments --- dvc/repo/experiments/__init__.py | 19 +++++---- dvc/repo/experiments/queue/tempdir.py | 27 +++++++------ dvc/repo/experiments/queue/workspace.py | 4 +- dvc/repo/experiments/show.py | 4 +- tests/func/experiments/test_show.py | 51 ++++++++++++++++++++++++- 5 files changed, 84 insertions(+), 21 deletions(-) diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index abc81f985..7343c31cf 100644 --- a/dvc/repo/experiments/__init__.py +++ b/dvc/repo/experiments/__init__.py @@ -432,15 +432,20 @@ class Experiments: """Return info for running experiments.""" result = {} infofile = self.workspace_queue.get_infofile_path("workspace") - result.update(self._get_running_exp("workspace", infofile, fetch_refs)) - for entry in self.celery_queue.iter_active(): - infofile = self.celery_queue.get_infofile_path(entry.stash_rev) - result.update( - self._get_running_exp(entry.stash_rev, infofile, fetch_refs) - ) + result.update( + self._fetch_running_exp("workspace", infofile, fetch_refs) + ) + for queue in (self.tempdir_queue, self.celery_queue): + for entry in queue.iter_active(): + infofile = queue.get_infofile_path(entry.stash_rev) + result.update( + self._fetch_running_exp( + entry.stash_rev, infofile, fetch_refs + ) + ) return result - def _get_running_exp( + def _fetch_running_exp( self, rev: str, infofile: str, fetch_refs: bool ) -> Dict[str, Any]: from dvc.scm import InvalidRemoteSCMRepo diff --git a/dvc/repo/experiments/queue/tempdir.py b/dvc/repo/experiments/queue/tempdir.py index 240cb39c3..4133539cd 100644 --- a/dvc/repo/experiments/queue/tempdir.py +++ b/dvc/repo/experiments/queue/tempdir.py @@ -1,6 +1,6 @@ import logging import os -from typing import TYPE_CHECKING, Dict, Generator +from typing import TYPE_CHECKING, Dict, Generator, Optional from funcy import cached_property, first @@ -29,6 +29,8 @@ _STANDALONE_TMP_DIR = os.path.join(EXEC_TMP_DIR, "standalone") class TempDirQueue(WorkspaceQueue): """Standalone/tempdir exp queue implementation.""" + _EXEC_NAME: Optional[str] = None + @cached_property def _standalone_tmp_dir(self) -> str: return os.path.join(self.repo.tmp_dir, _STANDALONE_TMP_DIR) @@ -73,16 +75,19 @@ class TempDirQueue(WorkspaceQueue): 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, - ) + if not executor_info.collected and os.path.exists( + executor_info.root_dir + ): + 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( diff --git a/dvc/repo/experiments/queue/workspace.py b/dvc/repo/experiments/queue/workspace.py index 477e2676b..e8c662cd2 100644 --- a/dvc/repo/experiments/queue/workspace.py +++ b/dvc/repo/experiments/queue/workspace.py @@ -19,6 +19,8 @@ logger = logging.getLogger(__name__) class WorkspaceQueue(BaseStashQueue): + _EXEC_NAME: Optional[str] = "workspace" + def put(self, *args, **kwargs) -> QueueEntry: return self._stash_exp(*args, **kwargs) @@ -77,7 +79,7 @@ class WorkspaceQueue(BaseStashQueue): from dvc.stage.monitor import CheckpointKilledError results: Dict[str, Dict[str, str]] = defaultdict(dict) - exec_name = "workspace" + exec_name = self._EXEC_NAME or entry.stash_rev infofile = self.get_infofile_path(exec_name) try: rev = entry.stash_rev diff --git a/dvc/repo/experiments/show.py b/dvc/repo/experiments/show.py index 1ed7d76c8..2141a4dbd 100644 --- a/dvc/repo/experiments/show.py +++ b/dvc/repo/experiments/show.py @@ -193,8 +193,10 @@ def show( running=running, onerror=onerror, ) - # collect celery experiments + + # collect standalone & celery experiments for entry in chain( + repo.experiments.tempdir_queue.iter_active(), repo.experiments.celery_queue.iter_active(), repo.experiments.celery_queue.iter_queued(), ): diff --git a/tests/func/experiments/test_show.py b/tests/func/experiments/test_show.py index fc018f00c..c973a0016 100644 --- a/tests/func/experiments/test_show.py +++ b/tests/func/experiments/test_show.py @@ -13,6 +13,7 @@ from dvc.repo.experiments.executor.base import ( BaseExecutor, ExecutorInfo, ) +from dvc.repo.experiments.queue.base import QueueEntry from dvc.repo.experiments.refs import CELERY_STASH, ExpRefInfo from dvc.repo.experiments.utils import exp_refs_by_rev from dvc.utils.fs import makedirs @@ -364,7 +365,55 @@ def test_show_running_workspace(tmp_dir, scm, dvc, exp_stage, capsys): assert info.location in cap.out -def test_show_running_executor(tmp_dir, scm, dvc, exp_stage, mocker): +def test_show_running_tempdir(tmp_dir, scm, dvc, exp_stage, mocker): + baseline_rev = scm.get_rev() + run_results = dvc.experiments.run( + exp_stage.addressing, params=["foo=2"], tmp_dir=True + ) + exp_rev = first(run_results) + exp_ref = first(exp_refs_by_rev(scm, exp_rev)) + + queue = dvc.experiments.tempdir_queue + stash_rev = "abc123" + entries = [ + QueueEntry( + str(tmp_dir / ".dvc" / "tmp" / "foo"), + str(tmp_dir / ".dvc" / "tmp" / "foo"), + str(exp_ref), + stash_rev, + exp_ref.baseline_sha, + None, + exp_ref.name, + None, + ) + ] + mocker.patch.object( + dvc.experiments.tempdir_queue, + "iter_active", + return_value=entries, + ) + info = make_executor_info(location=BaseExecutor.DEFAULT_LOCATION) + pidfile = queue.get_infofile_path(stash_rev) + makedirs(os.path.dirname(pidfile), True) + (tmp_dir / pidfile).dump_json(info.asdict()) + mock_fetch = mocker.patch.object( + dvc.experiments, + "_fetch_running_exp", + return_value={exp_rev: info.asdict()}, + ) + + results = dvc.experiments.show() + mock_fetch.assert_has_calls( + [mocker.call(stash_rev, pidfile, True)], + ) + exp_data = get_in(results, [baseline_rev, exp_rev, "data"]) + assert exp_data["running"] + assert exp_data["executor"] == info.location + + assert not results["workspace"]["baseline"]["data"]["running"] + + +def test_show_running_celery(tmp_dir, scm, dvc, exp_stage, mocker): baseline_rev = scm.get_rev() dvc.experiments.run(exp_stage.addressing, params=["foo=2"], queue=True) exp_rev = dvc.experiments.scm.resolve_rev(f"{CELERY_STASH}@{{0}}")