diff --git a/dvc/repo/experiments/__init__.py b/dvc/repo/experiments/__init__.py index 8d3fab61e..bebf87626 100644 --- a/dvc/repo/experiments/__init__.py +++ b/dvc/repo/experiments/__init__.py @@ -2,7 +2,7 @@ import logging import os import re import time -from typing import Any, Dict, Iterable, Optional +from typing import Dict, Iterable, Optional from funcy import cached_property, first @@ -16,7 +16,7 @@ from .exceptions import ( InvalidExpRefError, MultipleBranchError, ) -from .executor.base import BaseExecutor, ExecutorInfo, TaskStatus +from .executor.base import BaseExecutor from .queue.base import BaseStashQueue, QueueEntry from .queue.celery import LocalCeleryQueue from .queue.tempdir import TempDirQueue @@ -25,7 +25,6 @@ from .refs import ( CELERY_FAILED_STASH, CELERY_STASH, EXEC_APPLY, - EXEC_BRANCH, EXEC_CHECKPOINT, EXEC_NAMESPACE, EXPS_NAMESPACE, @@ -438,74 +437,15 @@ class Experiments: result[rev] = name return result - def get_running_exps(self, fetch_refs: bool = True) -> Dict[str, Any]: + def get_running_exps(self, fetch_refs: bool = True) -> Dict[str, Dict]: """Return info for running experiments.""" result = {} - infofile = self.workspace_queue.get_infofile_path("workspace") - 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 _fetch_running_exp( - self, rev: str, infofile: str, fetch_refs: bool - ) -> Dict[str, Any]: - from dvc.scm import InvalidRemoteSCMRepo - from dvc.utils.serialize import load_json - - from .executor.local import TempDirExecutor - - result: Dict[str, Any] = {} - try: - info = ExecutorInfo.from_dict(load_json(infofile)) - except OSError: - return result - if info.status < TaskStatus.FAILED: - if rev == "workspace": - # If we are appending to a checkpoint branch in a workspace - # run, show the latest checkpoint as running. - if info.status == TaskStatus.SUCCESS: - return result - last_rev = self.scm.get_ref(EXEC_BRANCH) - if last_rev: - result[last_rev] = info.asdict() - else: - result[rev] = info.asdict() - else: - 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] + for queue in ( + self.workspace_queue, + self.tempdir_queue, + self.celery_queue, + ): + result.update(queue.get_running_exps(fetch_refs)) return result def apply(self, *args, **kwargs): diff --git a/dvc/repo/experiments/queue/base.py b/dvc/repo/experiments/queue/base.py index e7739d284..df67967cd 100644 --- a/dvc/repo/experiments/queue/base.py +++ b/dvc/repo/experiments/queue/base.py @@ -31,7 +31,9 @@ from ..executor.base import ( EXEC_PID_DIR, EXEC_TMP_DIR, BaseExecutor, + ExecutorInfo, ExecutorResult, + TaskStatus, ) from ..executor.local import WorkspaceExecutor from ..refs import ExpRefInfo @@ -706,3 +708,57 @@ class BaseStashQueue(ABC): entry.stash_rev, 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 + + 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/tempdir.py b/dvc/repo/experiments/queue/tempdir.py index 1b88a8e12..4d3e9d33c 100644 --- a/dvc/repo/experiments/queue/tempdir.py +++ b/dvc/repo/experiments/queue/tempdir.py @@ -98,3 +98,6 @@ class TempDirQueue(WorkspaceQueue): exec_result: ExecutorResult, ) -> Dict[str, str]: 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) diff --git a/dvc/repo/experiments/queue/workspace.py b/dvc/repo/experiments/queue/workspace.py index 463d41d5e..43aa61304 100644 --- a/dvc/repo/experiments/queue/workspace.py +++ b/dvc/repo/experiments/queue/workspace.py @@ -7,7 +7,12 @@ from funcy import first from dvc.exceptions import DvcException from ..exceptions import ExpQueueEmptyError -from ..executor.base import BaseExecutor, ExecutorResult +from ..executor.base import ( + BaseExecutor, + ExecutorInfo, + ExecutorResult, + TaskStatus, +) from ..executor.local import WorkspaceExecutor from ..refs import EXEC_BRANCH from .base import BaseStashQueue, QueueDoneResult, QueueEntry, QueueGetResult @@ -150,3 +155,28 @@ class WorkspaceQueue(BaseStashQueue): follow: bool = False, ): raise NotImplementedError + + def get_running_exps(self, fetch_refs: bool = True) -> Dict[str, Dict]: + from dvc.utils.serialize import load_json + + assert self._EXEC_NAME + + result: Dict[str, Dict] = {} + infofile = self.get_infofile_path(self._EXEC_NAME) + + try: + info = ExecutorInfo.from_dict(load_json(infofile)) + except OSError: + return result + + if info.status < TaskStatus.FAILED: + # If we are appending to a checkpoint branch in a workspace + # run, show the latest checkpoint as running. + if info.status == TaskStatus.SUCCESS: + return result + last_rev = self.scm.get_ref(EXEC_BRANCH) + if last_rev: + result[last_rev] = info.asdict() + else: + result[self._EXEC_NAME] = info.asdict() + return result diff --git a/dvc/repo/experiments/show.py b/dvc/repo/experiments/show.py index 8a1da3b99..ee96a83c3 100644 --- a/dvc/repo/experiments/show.py +++ b/dvc/repo/experiments/show.py @@ -287,7 +287,9 @@ def show( iter_revs(repo.scm, revs, num, all_branches, all_tags, all_commits) ) - running = repo.experiments.get_running_exps(fetch_refs=fetch_running) + running: Dict[str, Dict] = repo.experiments.get_running_exps( + fetch_refs=fetch_running + ) queued_experiment = ( _collect_queued_experiment( diff --git a/tests/func/experiments/test_show.py b/tests/func/experiments/test_show.py index d39868da0..6851ff9ff 100644 --- a/tests/func/experiments/test_show.py +++ b/tests/func/experiments/test_show.py @@ -487,7 +487,7 @@ def test_show_running_tempdir(tmp_dir, scm, dvc, exp_stage, mocker): makedirs(os.path.dirname(pidfile), True) (tmp_dir / pidfile).dump_json(info.asdict()) mock_fetch = mocker.patch.object( - dvc.experiments, + dvc.experiments.tempdir_queue, "_fetch_running_exp", return_value={exp_rev: info.asdict()}, )