From 9ef1d1d6e6c95bd9cd72c6a33d0778705e4214c2 Mon Sep 17 00:00:00 2001 From: karajan1001 Date: Thu, 27 Oct 2022 18:44:05 +0800 Subject: [PATCH] Move away tempdir related change from basequeue. --- dvc/repo/experiments/queue/base.py | 50 +---------------------- dvc/repo/experiments/queue/celery.py | 16 ++++++++ dvc/repo/experiments/queue/tempdir.py | 10 ++++- dvc/repo/experiments/queue/utils.py | 58 +++++++++++++++++++++++++++ tests/func/experiments/test_show.py | 4 +- 5 files changed, 86 insertions(+), 52 deletions(-) create mode 100644 dvc/repo/experiments/queue/utils.py diff --git a/dvc/repo/experiments/queue/base.py b/dvc/repo/experiments/queue/base.py index df67967cd..6ddf0fc71 100644 --- a/dvc/repo/experiments/queue/base.py +++ b/dvc/repo/experiments/queue/base.py @@ -31,9 +31,7 @@ from ..executor.base import ( EXEC_PID_DIR, EXEC_TMP_DIR, BaseExecutor, - ExecutorInfo, ExecutorResult, - TaskStatus, ) from ..executor.local import WorkspaceExecutor from ..refs import ExpRefInfo @@ -709,56 +707,10 @@ class BaseStashQueue(ABC): message=f"commit: {msg}", ) - def _fetch_running_exp( - self, rev: str, infofile: str, fetch_refs: bool - ) -> Dict[str, Dict]: - from dvc.scm import InvalidRemoteSCMRepo - from dvc.utils.serialize import load_json - - from ..executor.local import TempDirExecutor - - result: Dict[str, Dict] = {} - try: - info = ExecutorInfo.from_dict(load_json(infofile)) - except OSError: - return result - if info.status < TaskStatus.FAILED: - result[rev] = info.asdict() - if ( - info.git_url - and fetch_refs - and info.status > TaskStatus.PREPARING - ): - - def on_diverged(_ref: str, _checkpoint: bool): - return False - - executor = TempDirExecutor.from_info(info) - try: - for ref in executor.fetch_exps( - self.scm, - on_diverged=on_diverged, - ): - logger.debug("Updated running experiment '%s'.", ref) - last_rev = self.scm.get_ref(ref) - result[rev]["last"] = last_rev - if last_rev: - result[last_rev] = info.asdict() - except InvalidRemoteSCMRepo: - # ignore stale info files - del result[rev] - return result - + @abstractmethod def get_running_exps(self, fetch_refs: bool = True) -> Dict[str, Dict]: """Get the execution info of the currently running experiments Args: fetch_ref (bool): fetch completed checkpoints or not. """ - result: Dict[str, Dict] = {} - for entry in self.iter_active(): - infofile = self.get_infofile_path(entry.stash_rev) - result.update( - self._fetch_running_exp(entry.stash_rev, infofile, fetch_refs) - ) - return result diff --git a/dvc/repo/experiments/queue/celery.py b/dvc/repo/experiments/queue/celery.py index 0e85e9393..c50978088 100644 --- a/dvc/repo/experiments/queue/celery.py +++ b/dvc/repo/experiments/queue/celery.py @@ -31,6 +31,7 @@ from ..executor.base import ( ) from .base import BaseStashQueue, QueueDoneResult, QueueEntry, QueueGetResult from .tasks import run_exp +from .utils import fetch_running_exp_from_temp_dir if TYPE_CHECKING: from dvc_task.app import FSApp @@ -373,3 +374,18 @@ class LocalCeleryQueue(BaseStashQueue): from .remove import celery_remove return celery_remove(self, *args, **kwargs) + + def get_running_exps(self, fetch_refs: bool = True) -> Dict[str, Dict]: + """Get the execution info of the currently running experiments + + Args: + fetch_ref (bool): fetch completed checkpoints or not. + """ + result: Dict[str, Dict] = {} + for entry in self.iter_active(): + result.update( + fetch_running_exp_from_temp_dir( + self, entry.stash_rev, fetch_refs + ) + ) + return result diff --git a/dvc/repo/experiments/queue/tempdir.py b/dvc/repo/experiments/queue/tempdir.py index 4d3e9d33c..d5ff16812 100644 --- a/dvc/repo/experiments/queue/tempdir.py +++ b/dvc/repo/experiments/queue/tempdir.py @@ -15,6 +15,7 @@ from ..executor.base import ( ) from ..executor.local import TempDirExecutor from .base import BaseStashQueue, QueueEntry, QueueGetResult +from .utils import fetch_running_exp_from_temp_dir from .workspace import WorkspaceQueue if TYPE_CHECKING: @@ -100,4 +101,11 @@ class TempDirQueue(WorkspaceQueue): return BaseStashQueue.collect_executor(exp, executor, exec_result) def get_running_exps(self, fetch_refs: bool = True) -> Dict[str, Dict]: - return super(WorkspaceQueue, self).get_running_exps(fetch_refs) + result: Dict[str, Dict] = {} + for entry in self.iter_active(): + result.update( + fetch_running_exp_from_temp_dir( + self, entry.stash_rev, fetch_refs + ) + ) + return result diff --git a/dvc/repo/experiments/queue/utils.py b/dvc/repo/experiments/queue/utils.py new file mode 100644 index 000000000..d77d33648 --- /dev/null +++ b/dvc/repo/experiments/queue/utils.py @@ -0,0 +1,58 @@ +import logging +from typing import TYPE_CHECKING, Dict + +from ..executor.base import ExecutorInfo, TaskStatus + +logger = logging.getLogger(__name__) + + +if TYPE_CHECKING: + from .base import BaseStashQueue + + +def fetch_running_exp_from_temp_dir( + queue: "BaseStashQueue", rev: str, fetch_refs: bool +) -> Dict[str, Dict]: + """Fetch status of running exps out of current working directory + + Args: + queue (BaseStashQueue): + rev (str): stash revision of the experiment + fetch_refs (bool): fetch running checkpoint results to local or not. + + Returns: + Dict[str, Dict]: _description_ + """ + from dvc.scm import InvalidRemoteSCMRepo + from dvc.utils.serialize import load_json + + from ..executor.local import TempDirExecutor + + result: Dict[str, Dict] = {} + infofile = queue.get_infofile_path(rev) + try: + info = ExecutorInfo.from_dict(load_json(infofile)) + except OSError: + return result + if info.status < TaskStatus.FAILED: + result[rev] = info.asdict() + if info.git_url and fetch_refs and info.status > TaskStatus.PREPARING: + + def on_diverged(_ref: str, _checkpoint: bool): + return False + + executor = TempDirExecutor.from_info(info) + try: + for ref in executor.fetch_exps( + queue.scm, + on_diverged=on_diverged, + ): + logger.debug("Updated running experiment '%s'.", ref) + last_rev = queue.scm.get_ref(ref) + result[rev]["last"] = last_rev + if last_rev: + result[last_rev] = info.asdict() + except InvalidRemoteSCMRepo: + # ignore stale info files + del result[rev] + return result diff --git a/tests/func/experiments/test_show.py b/tests/func/experiments/test_show.py index 08ee4c8b7..2869053fe 100644 --- a/tests/func/experiments/test_show.py +++ b/tests/func/experiments/test_show.py @@ -524,13 +524,13 @@ def test_show_running_tempdir(tmp_dir, scm, dvc, exp_stage, mocker): (tmp_dir / pidfile).dump_json(info.asdict()) mock_fetch = mocker.patch.object( dvc.experiments.tempdir_queue, - "_fetch_running_exp", + "get_running_exps", return_value={exp_rev: info.asdict()}, ) results = dvc.experiments.show() mock_fetch.assert_has_calls( - [mocker.call(stash_rev, pidfile, True)], + [mocker.call(True)], ) exp_data = get_in(results, [baseline_rev, exp_rev, "data"]) assert exp_data["status"] == "Running"