From 58f74ae543439c503d634c58b14117622b192c54 Mon Sep 17 00:00:00 2001 From: Peter Rowlands Date: Thu, 20 Apr 2023 15:39:43 +0900 Subject: [PATCH] exp: standardize exp data collection --- dvc/repo/experiments/collect.py | 380 ++++++++++++++++++++++++ dvc/repo/experiments/queue/base.py | 55 ++++ dvc/repo/experiments/queue/celery.py | 135 ++++++++- dvc/repo/experiments/queue/tempdir.py | 54 +++- dvc/repo/experiments/queue/workspace.py | 148 +++++---- 5 files changed, 706 insertions(+), 66 deletions(-) create mode 100644 dvc/repo/experiments/collect.py diff --git a/dvc/repo/experiments/collect.py b/dvc/repo/experiments/collect.py new file mode 100644 index 000000000..dfcc1d2e1 --- /dev/null +++ b/dvc/repo/experiments/collect.py @@ -0,0 +1,380 @@ +import itertools +import logging +import os +from datetime import datetime +from typing import ( + TYPE_CHECKING, + Collection, + Dict, + Iterable, + Iterator, + List, + Optional, + Union, +) + +from funcy import first +from scmrepo.exceptions import SCMError as InnerSCMError + +from dvc.scm import Git, SCMError, iter_revs + +from .exceptions import InvalidExpRefError +from .refs import EXEC_BRANCH, ExpRefInfo +from .serialize import ExpRange, ExpState, SerializableError, SerializableExp + +if TYPE_CHECKING: + from dvc.repo import Repo + + from .cache import ExpCache + +logger = logging.getLogger(__name__) + + +def collect_rev( + repo: "Repo", + rev: str, + param_deps: bool = False, + force: bool = False, + cache: Optional["ExpCache"] = None, + **kwargs, +) -> ExpState: + """Collect experiment state for the given revision. + + Exp will be loaded from cache when available unless rev is 'workspace' or + force is set. + """ + from dvc.fs import LocalFileSystem + + cache = cache or repo.experiments.cache + assert cache + # TODO: support filtering serialized exp when param_deps is set + if rev != "workspace" and not (force or param_deps): + cached_exp = cache.get(rev) + if cached_exp: + if isinstance(cached_exp, SerializableError): + return ExpState(rev=rev, error=cached_exp) + return ExpState(rev=rev, data=cached_exp) + if rev == "workspace" and isinstance(repo.fs, LocalFileSystem): + orig_cwd: Optional[str] = os.getcwd() + os.chdir(repo.root_dir) + else: + orig_cwd = None + try: + data = _collect_rev( + repo, + rev, + param_deps=param_deps, + force=force, + **kwargs, + ) + if not (rev == "workspace" or param_deps or data.contains_error): + cache.put(data, force=True) + return ExpState(rev=rev, data=data) + except Exception as exc: # noqa: BLE001, pylint: disable=broad-except + logger.debug("", exc_info=True) + error = SerializableError(str(exc), type(exc).__name__) + return ExpState(rev=rev, error=error) + finally: + if orig_cwd: + os.chdir(orig_cwd) + + +def _collect_rev( + repo: "Repo", + rev: str, + param_deps: bool = False, + **kwargs, +) -> SerializableExp: + with repo.switch(rev) as rev: + if rev == "workspace": + timestamp: Optional[datetime] = None + else: + commit = repo.scm.resolve_commit(rev) + timestamp = datetime.fromtimestamp(commit.commit_time) + + return SerializableExp.from_repo( + repo, + rev=rev, + param_deps=param_deps, + timestamp=timestamp, + ) + + +def collect_branch( + repo: "Repo", + rev: str, + end_rev: Optional[str] = None, + **kwargs, +) -> Iterator["ExpState"]: + """Iterate over exp states in a Git branch. + + Git branch will be traversed in reverse, starting from rev. + + Args: + rev: Branch tip (head). + end_rev: If specified, traversal will stop when end_rev is reached + (exclusive, end_rev will not be collected). + """ + try: + for branch_rev in repo.scm.branch_revs(rev, end_rev): + yield collect_rev(repo, branch_rev, **kwargs) + except (SCMError, InnerSCMError): + pass + + +def collect_exec_branch( + repo: "Repo", + baseline_rev: str, + **kwargs, +) -> Iterator["ExpState"]: + """Iterate over active experiment branch for the current executor.""" + last_rev = repo.scm.get_ref(EXEC_BRANCH) or repo.scm.get_rev() + last_rev = repo.scm.get_rev() + yield collect_rev(repo, "workspace", **kwargs) + if last_rev != baseline_rev: + yield from collect_branch(repo, last_rev, baseline_rev, **kwargs) + + +def collect_queued( + repo: "Repo", + baseline_revs: Collection[str], + **kwargs, +) -> Dict[str, List["ExpRange"]]: + """Collect queued experiments derived from the specified revisions. + + Args: + repo: Repo. + baseline_revs: Resolved baseline Git SHAs. + + Returns: + Dict mapping baseline revision to list of queued experiments. + """ + if not baseline_revs: + return {} + return repo.experiments.celery_queue.collect_queued_data(baseline_revs, **kwargs) + + +def collect_active( + repo: "Repo", + baseline_revs: Collection[str], + **kwargs, +) -> Dict[str, List["ExpRange"]]: + """Collect active (running) experiments derived from the specified revisions. + + Args: + repo: Repo. + baseline_revs: Resolved baseline Git SHAs. + + Returns: + Dict mapping baseline revision to list of active experiments. + """ + if not baseline_revs: + return {} + result: Dict[str, List["ExpRange"]] = {} + exps = repo.experiments + for queue in (exps.workspace_queue, exps.tempdir_queue, exps.celery_queue): + for baseline, active_exps in queue.collect_active_data( + baseline_revs, **kwargs + ).items(): + if baseline in result: + result[baseline].extend(active_exps) + else: + result[baseline] = list(active_exps) + return result + + +def collect_failed( + repo: "Repo", + baseline_revs: Collection[str], + **kwargs, +) -> Dict[str, List["ExpRange"]]: + """Collect failed experiments derived from the specified revisions. + + Args: + repo: Repo. + baseline_revs: Resolved baseline Git SHAs. + + Returns: + Dict mapping baseline revision to list of active experiments. + """ + if not baseline_revs: + return {} + return repo.experiments.celery_queue.collect_failed_data(baseline_revs, **kwargs) + + +def collect_successful( + repo: "Repo", + baseline_revs: Collection[str], + **kwargs, +) -> Dict[str, List["ExpRange"]]: + """Collect successful experiments derived from the specified revisions. + + Args: + repo: Repo. + baseline_revs: Resolved baseline Git SHAs. + + Returns: + Dict mapping baseline revision to successful experiments. + """ + result: Dict[str, List["ExpRange"]] = {} + for baseline_rev in baseline_revs: + result[baseline_rev] = list(_collect_baseline(repo, baseline_rev, **kwargs)) + return result + + +def _collect_baseline( + repo: "Repo", + baseline_rev: str, + **kwargs, +) -> Iterator["ExpRange"]: + """Iterate over experiments derived from a baseline revision. + + Args: + repo: Repo. + baseline_revs: Resolved baseline Git SHAs. + + Yields: + Tuple of (timestamp, exp_range). + """ + ref_info = ExpRefInfo(baseline_sha=baseline_rev) + refs: Optional[Iterable[str]] = kwargs.get("refs") + if refs: + ref_it = (ref for ref in iter(refs) if ref.startswith(str(ref_info))) + else: + ref_it = repo.scm.iter_refs(base=str(ref_info)) + for ref in ref_it: + try: + ref_info = ExpRefInfo.from_ref(ref) + exp_rev = repo.scm.get_ref(ref) + if not exp_rev: + continue + except (InvalidExpRefError, SCMError, InnerSCMError): + continue + exps = list(collect_branch(repo, exp_rev, baseline_rev, **kwargs)) + if exps: + exps[0].name = ref_info.name + yield ExpRange(exps, name=ref_info.name) + + +def collect( + repo: "Repo", + revs: Union[List[str], str, None] = None, + all_branches: bool = False, + all_tags: bool = False, + all_commits: bool = False, + num: int = 1, + hide_queued: bool = False, + hide_failed: bool = False, + sha_only: bool = False, + **kwargs, +) -> List["ExpState"]: + """Collect baseline revisions and derived experiments.""" + assert isinstance(repo.scm, Git) + if repo.scm.no_commits: + return [] + if not any([revs, all_branches, all_tags, all_commits]): + revs = ["HEAD"] + if isinstance(revs, str): + revs = [revs] + cached_refs = list(repo.scm.iter_refs()) + baseline_revs = list( + iter_revs( + repo.scm, + revs=revs, + num=num, + all_branches=all_branches, + all_tags=all_tags, + all_commits=all_commits, + ) + ) + if sha_only: + baseline_names: Dict[str, Optional[str]] = {} + else: + baseline_names = _describe(repo.scm, baseline_revs, refs=cached_refs) + + workspace_data = collect_rev(repo, "workspace", **kwargs) + result: List["ExpState"] = [workspace_data] + queued = collect_queued(repo, baseline_revs, **kwargs) if not hide_queued else {} + active = collect_active(repo, baseline_revs, **kwargs) + failed = collect_failed(repo, baseline_revs, **kwargs) if not hide_failed else {} + successful = collect_successful(repo, baseline_revs, **kwargs) + + for baseline_rev in baseline_revs: + baseline_data = collect_rev(repo, baseline_rev) + experiments = list( + itertools.chain.from_iterable( + _sorted_ranges(collected.get(baseline_rev, [])) + for collected in (successful, active, queued, failed) + ) + ) + result.append( + ExpState( + rev=baseline_rev, + name=baseline_names.get(baseline_rev), + data=baseline_data.data, + error=baseline_data.error, + experiments=experiments if experiments else None, + ) + ) + return result + + +def _describe( + scm: "Git", + revs: Iterable[str], + refs: Optional[Iterable[str]] = None, +) -> Dict[str, Optional[str]]: + """Describe revisions using a tag, branch. + + The first matching name will be returned for each rev. Names are preferred in this + order: + - current branch (if rev matches HEAD and HEAD is a branch) + - tags + - branches + + Returns: + Dict mapping revisions from revs to a name. + """ + + head_rev = scm.get_rev() + head_ref = scm.get_ref("HEAD", follow=False) + if head_ref and head_ref.startswith("refs/heads/"): + head_branch = head_ref[len("refs/heads/") :] + else: + head_branch = None + + tags = {} + branches = {} + ref_it = iter(refs) if refs else scm.iter_refs() + for ref in ref_it: + is_tag = ref.startswith("refs/tags/") + is_branch = ref.startswith("refs/heads/") + if not (is_tag or is_branch): + continue + rev = scm.get_ref(ref) + if not rev: + logger.debug("unresolved ref %s", ref) + continue + if is_tag and rev not in tags: + tags[rev] = ref[len("refs/tags/") :] + if is_branch and rev not in branches: + branches[rev] = ref[len("refs/heads/") :] + names: Dict[str, Optional[str]] = {} + for rev in revs: + if rev == head_rev and head_branch: + names[rev] = head_branch + else: + names[rev] = tags.get(rev) or branches.get(rev) + return names + + +def _sorted_ranges(exp_ranges: Iterable["ExpRange"]) -> List["ExpRange"]: + """Return list of ExpRange sorted by timestamp.""" + + def _head_timestamp(exp_range: "ExpRange") -> datetime: + head_exp = first(exp_range.revs) + if head_exp and head_exp.data and head_exp.data.timestamp: + return head_exp.data.timestamp + return datetime.fromtimestamp(0) + + return sorted(exp_ranges, key=_head_timestamp, reverse=True) diff --git a/dvc/repo/experiments/queue/base.py b/dvc/repo/experiments/queue/base.py index e3921c33e..98ff3ace8 100644 --- a/dvc/repo/experiments/queue/base.py +++ b/dvc/repo/experiments/queue/base.py @@ -46,6 +46,7 @@ if TYPE_CHECKING: from dvc.repo import Repo from dvc.repo.experiments import Experiments from dvc.repo.experiments.executor.base import ExecutorResult + from dvc.repo.experiments.serialize import ExpRange from dvc.scm import Git logger = logging.getLogger(__name__) @@ -734,3 +735,57 @@ class BaseStashQueue(ABC): Args: fetch_ref (bool): fetch completed checkpoints or not. """ + + @abstractmethod + def collect_active_data( + self, + baseline_revs: Optional[Collection[str]], + fetch_refs: bool = False, + **kwargs, + ) -> Dict[str, List["ExpRange"]]: + """Collect data for active (running) experiments. + + Args: + baseline_revs: Optional resolved baseline Git SHAs. If set, only experiments + derived from the specified revisions will be collected. Defaults to + collecting all experiments. + fetch_refs: Whether or not to fetch completed checkpoint commits from Git + remote. + + Returns: + Dict mapping baseline revision to list of active experiments. + """ + + @abstractmethod + def collect_queued_data( + self, + baseline_revs: Optional[Collection[str]], + **kwargs, + ) -> Dict[str, List["ExpRange"]]: + """Collect data for queued experiments. + + Args: + baseline_revs: Optional resolved baseline Git SHAs. If set, only experiments + derived from the specified revisions will be collected. Defaults to + collecting all experiments. + + Returns: + Dict mapping baseline revision to list of queued experiments. + """ + + @abstractmethod + def collect_failed_data( + self, + baseline_revs: Optional[Collection[str]], + **kwargs, + ) -> Dict[str, List["ExpRange"]]: + """Collect data for failed experiments. + + Args: + baseline_revs: Optional resolved baseline Git SHAs. If set, only experiments + derived from the specified revisions will be collected. Defaults to + collecting all experiments. + + Returns: + Dict mapping baseline revision to list of queued experiments. + """ diff --git a/dvc/repo/experiments/queue/celery.py b/dvc/repo/experiments/queue/celery.py index cc9669e7b..1ca787719 100644 --- a/dvc/repo/experiments/queue/celery.py +++ b/dvc/repo/experiments/queue/celery.py @@ -2,6 +2,7 @@ import hashlib import locale import logging import os +from collections import defaultdict from typing import ( TYPE_CHECKING, Collection, @@ -40,6 +41,7 @@ if TYPE_CHECKING: from dvc.repo.experiments.executor.base import ExecutorResult from dvc.repo.experiments.refs import ExpRefInfo + from dvc.repo.experiments.serialize import ExpRange from dvc_task.app import FSApp from dvc_task.proc.manager import ProcessManager from dvc_task.worker import TemporaryWorker @@ -198,7 +200,9 @@ class LocalCeleryQueue(BaseStashQueue): continue args, kwargs, _embed = msg.decode() entry_dict = kwargs.get("entry_dict", args[0]) - logger.debug("Found queued task %s", entry_dict["stash_rev"]) + logger.trace( # type: ignore[attr-defined] + "Found queued task %s", entry_dict["stash_rev"] + ) yield _MessageEntry(msg, QueueEntry.from_dict(entry_dict)) def _iter_processed(self) -> Generator[_MessageEntry, None, None]: @@ -214,7 +218,9 @@ class LocalCeleryQueue(BaseStashQueue): task_id = msg.headers["id"] result: AsyncResult = AsyncResult(task_id) if not result.ready(): - logger.debug("Found active task %s", entry.stash_rev) + logger.trace( # type: ignore[attr-defined] + "Found active task %s", entry.stash_rev + ) yield _TaskEntry(result, entry) def _iter_done_tasks(self) -> Generator[_TaskEntry, None, None]: @@ -222,7 +228,9 @@ class LocalCeleryQueue(BaseStashQueue): task_id = msg.headers["id"] result: AsyncResult = AsyncResult(task_id) if result.ready(): - logger.debug("Found done task %s", entry.stash_rev) + logger.trace( # type: ignore[attr-defined] + "Found done task %s", entry.stash_rev + ) yield _TaskEntry(result, entry) def iter_active(self) -> Generator[QueueEntry, None, None]: @@ -512,3 +520,124 @@ class LocalCeleryQueue(BaseStashQueue): queue_entry = None if git_remote else queue_entry_match[exp_name] results[exp_name] = ExpRefAndQueueEntry(exp_ref, queue_entry) return results + + def collect_active_data( + self, + baseline_revs: Optional[Collection[str]], + fetch_refs: bool = False, + **kwargs, + ) -> Dict[str, List["ExpRange"]]: + from dvc.repo import Repo + from dvc.repo.experiments.collect import collect_exec_branch + from dvc.repo.experiments.serialize import ( + ExpExecutor, + ExpRange, + LocalExpExecutor, + ) + + result: Dict[str, List[ExpRange]] = defaultdict(list) + for entry in self.iter_active(): + if baseline_revs and entry.baseline_rev not in baseline_revs: + continue + if fetch_refs: + fetch_running_exp_from_temp_dir(self, entry.stash_rev, fetch_refs) + proc_info = self.proc.get(entry.stash_rev) + executor_info = self._load_info(entry.stash_rev) + if proc_info: + local_exec: Optional[LocalExpExecutor] = LocalExpExecutor( + root=executor_info.root_dir, + log=proc_info.stdout, + pid=proc_info.pid, + task_id=entry.stash_rev, + ) + else: + local_exec = None + dvc_root = os.path.join(executor_info.root_dir, executor_info.dvc_dir) + with Repo(dvc_root) as exec_repo: + kwargs["cache"] = self.repo.experiments.cache + exps = list( + collect_exec_branch(exec_repo, executor_info.baseline_rev, **kwargs) + ) + exps[0].rev = entry.stash_rev + exps[0].name = entry.name + result[entry.baseline_rev].append( + ExpRange( + exps, + executor=ExpExecutor( + "running", + name=executor_info.location, + local=local_exec, + ), + name=entry.name, + ) + ) + return result + + def collect_queued_data( + self, + baseline_revs: Optional[Collection[str]], + **kwargs, + ) -> Dict[str, List["ExpRange"]]: + from dvc.repo.experiments.collect import collect_rev + from dvc.repo.experiments.serialize import ( + ExpExecutor, + ExpRange, + LocalExpExecutor, + ) + + result: Dict[str, List[ExpRange]] = defaultdict(list) + for entry in self.iter_queued(): + if baseline_revs and entry.baseline_rev not in baseline_revs: + continue + exp = collect_rev(self.repo, entry.stash_rev, **kwargs) + exp.name = entry.name + local_exec: Optional[LocalExpExecutor] = LocalExpExecutor( + task_id=entry.stash_rev, + ) + result[entry.baseline_rev].append( + ExpRange( + [exp], + executor=ExpExecutor("queued", name="dvc-task", local=local_exec), + name=entry.name, + ) + ) + return result + + def collect_failed_data( + self, + baseline_revs: Optional[Collection[str]], + **kwargs, + ) -> Dict[str, List["ExpRange"]]: + from dvc.repo.experiments.collect import collect_rev + from dvc.repo.experiments.serialize import ( + ExpExecutor, + ExpRange, + LocalExpExecutor, + SerializableError, + ) + + result: Dict[str, List[ExpRange]] = defaultdict(list) + for entry, _ in self.iter_failed(): + if baseline_revs and entry.baseline_rev not in baseline_revs: + continue + proc_info = self.proc.get(entry.stash_rev) + if proc_info: + local_exec: Optional[LocalExpExecutor] = LocalExpExecutor( + log=proc_info.stdout, + pid=proc_info.pid, + returncode=proc_info.returncode, + task_id=entry.stash_rev, + ) + else: + local_exec = None + exp = collect_rev(self.repo, entry.stash_rev, **kwargs) + exp.name = entry.name + exp.error = SerializableError("Experiment run failed") + result[entry.baseline_rev].append( + ExpRange( + [exp], + executor=ExpExecutor("failed", local=local_exec), + name=entry.name, + ) + ) + return result diff --git a/dvc/repo/experiments/queue/tempdir.py b/dvc/repo/experiments/queue/tempdir.py index 944e83597..cbbdbed7c 100644 --- a/dvc/repo/experiments/queue/tempdir.py +++ b/dvc/repo/experiments/queue/tempdir.py @@ -1,7 +1,7 @@ import logging import os from collections import defaultdict -from typing import TYPE_CHECKING, Dict, Generator, List, Optional +from typing import TYPE_CHECKING, Collection, Dict, Generator, List, Optional from funcy import first @@ -19,6 +19,7 @@ from .workspace import WorkspaceQueue if TYPE_CHECKING: from dvc.repo.experiments import Experiments from dvc.repo.experiments.executor.base import BaseExecutor, ExecutorResult + from dvc.repo.experiments.serialize import ExpRange from dvc_task.proc.manager import ProcessManager logger = logging.getLogger(__name__) @@ -148,3 +149,54 @@ class TempDirQueue(WorkspaceQueue): fetch_running_exp_from_temp_dir(self, entry.stash_rev, fetch_refs) ) return result + + def collect_active_data( + self, + baseline_revs: Optional[Collection[str]], + fetch_refs: bool = False, + **kwargs, + ) -> Dict[str, List["ExpRange"]]: + from dvc.repo import Repo + from dvc.repo.experiments.collect import collect_exec_branch + from dvc.repo.experiments.serialize import ( + ExpExecutor, + ExpRange, + LocalExpExecutor, + ) + + result: Dict[str, List[ExpRange]] = defaultdict(list) + for entry in self.iter_active(): + if baseline_revs and entry.baseline_rev not in baseline_revs: + continue + if fetch_refs: + fetch_running_exp_from_temp_dir(self, entry.stash_rev, fetch_refs) + proc_info = self.proc.get(entry.stash_rev) + infofile = self.get_infofile_path(entry.stash_rev) + executor_info = ExecutorInfo.load_json(infofile) + if proc_info: + local_exec: Optional[LocalExpExecutor] = LocalExpExecutor( + root=executor_info.root_dir, + log=proc_info.stdout, + pid=proc_info.pid, + ) + else: + local_exec = None + dvc_root = os.path.join(executor_info.root_dir, executor_info.dvc_dir) + with Repo(dvc_root) as repo: + exps = list( + collect_exec_branch(repo, executor_info.baseline_rev, **kwargs) + ) + exps[0].rev = entry.stash_rev + exps[0].name = entry.name + result[entry.baseline_rev].append( + ExpRange( + exps, + executor=ExpExecutor( + "running", + name=executor_info.location, + local=local_exec, + ), + name=entry.name, + ) + ) + return result diff --git a/dvc/repo/experiments/queue/workspace.py b/dvc/repo/experiments/queue/workspace.py index 5ba6b29db..70c908478 100644 --- a/dvc/repo/experiments/queue/workspace.py +++ b/dvc/repo/experiments/queue/workspace.py @@ -1,28 +1,27 @@ import json import logging +import os from collections import defaultdict from typing import TYPE_CHECKING, Collection, Dict, Generator, List, Optional import psutil from funcy import first -from voluptuous import Invalid from dvc.exceptions import DvcException -from dvc.lock import make_lock from dvc.repo.experiments.exceptions import ExpQueueEmptyError from dvc.repo.experiments.executor.base import ExecutorInfo, TaskStatus from dvc.repo.experiments.executor.local import WorkspaceExecutor from dvc.repo.experiments.refs import EXEC_BRANCH, WORKSPACE_STASH from dvc.repo.experiments.utils import get_exp_rwlock -from dvc.rwlock import RWLOCK_FILE, RWLOCK_LOCK, SCHEMA -from dvc.utils import relpath from dvc.utils.fs import remove +from dvc.utils.serialize import load_json from .base import BaseStashQueue, QueueEntry, QueueGetResult if TYPE_CHECKING: from dvc.repo.experiments import Experiments from dvc.repo.experiments.executor.base import BaseExecutor, ExecutorResult + from dvc.repo.experiments.serialize import ExpRange from .base import QueueDoneResult @@ -101,9 +100,14 @@ class WorkspaceQueue(BaseStashQueue): ) -> Dict[str, Dict[str, str]]: kwargs.pop("copy_paths", None) from dvc.stage.monitor import CheckpointKilledError + from dvc_task.proc.process import ProcessInfo results: Dict[str, Dict[str, str]] = defaultdict(dict) exec_name = self._EXEC_NAME or entry.stash_rev + proc_info = ProcessInfo(os.getpid(), None, None, None, None) + proc_info_path = self._proc_info_path(exec_name) + os.makedirs(os.path.dirname(proc_info_path), exist_ok=True) + proc_info.dump(proc_info_path) infofile = self.get_infofile_path(exec_name) try: rev = entry.stash_rev @@ -129,8 +133,29 @@ class WorkspaceQueue(BaseStashQueue): raise DvcException(f"Failed to reproduce experiment '{rev[:7]}'") from exc finally: executor.cleanup(infofile) + remove(self._proc_info_path(exec_name)) return results + def _proc_info_path(self, name: str) -> str: + return os.path.join(self.pid_dir, name, f"{name}.json") + + @property + def _active_pid(self) -> Optional[int]: + from dvc_task.proc.process import ProcessInfo + + assert self._EXEC_NAME + name = self._EXEC_NAME + try: + proc_info = ProcessInfo.load(self._proc_info_path(name)) + pid = proc_info.pid + if psutil.pid_exists(pid): + return pid + logger.debug("Workspace exec PID '%d' no longer exists, removing.", pid) + remove(self._proc_info_path(name)) + except (FileNotFoundError, json.JSONDecodeError): + pass + return None + @staticmethod def collect_executor( # pylint: disable=unused-argument exp: "Experiments", @@ -163,70 +188,13 @@ class WorkspaceQueue(BaseStashQueue): ): raise NotImplementedError - def check_rwlock( - self, - hardlink: bool = False, - autocorrect: bool = False, - ) -> bool: - """Check and autocorrect the RWLock status for file paths. - - Args: - hardlink (bool): use hardlink lock to guard rwlock file when on - edit. - autocorrect (bool): autocorrect corrupted rwlock file. - - Return: - (bool): if the pid alive. - """ - path = self.repo.fs.path.join(self.repo.tmp_dir, RWLOCK_FILE) - - rwlock_guard = make_lock( - self.repo.fs.path.join(self.repo.tmp_dir, RWLOCK_LOCK), - tmp_dir=self.repo.tmp_dir, - hardlink_lock=hardlink, - ) - with rwlock_guard: - try: - with self.repo.fs.open(path, encoding="utf-8") as fobj: - lock: Dict[str, List[Dict]] = SCHEMA(json.load(fobj)) - file_path = first(lock["read"]) - if not file_path: - return False - lock_info = first(lock["read"][file_path]) - pid = int(lock_info["pid"]) - if psutil.pid_exists(pid): - return True - cmd = lock_info["cmd"] - logger.warning( - "Process '%s' with (Pid %s), in RWLock-file '%s' had been killed.", - cmd, - pid, - relpath(path), - ) - except FileNotFoundError: - return False - except json.JSONDecodeError: - logger.warning( - "Unable to read RWLock-file '%s'. JSON structure is corrupted", - relpath(path), - ) - except Invalid: - logger.warning("RWLock-file '%s' format error.", relpath(path)) - if autocorrect: - logger.warning("Delete corrupted RWLock-file '%s'", relpath(path)) - remove(path) - return False - def get_running_exps( self, fetch_refs: bool = True, # noqa: ARG002 ) -> Dict[str, Dict]: - from dvc.utils.serialize import load_json - assert self._EXEC_NAME result: Dict[str, Dict] = {} - - if not self.check_rwlock(autocorrect=True): + if self._active_pid is None: return result infofile = self.get_infofile_path(self._EXEC_NAME) @@ -247,3 +215,59 @@ class WorkspaceQueue(BaseStashQueue): else: result[self._EXEC_NAME] = info.asdict() return result + + def collect_active_data( + self, + baseline_revs: Optional[Collection[str]], + fetch_refs: bool = False, # noqa: ARG002 + **kwargs, + ) -> Dict[str, List["ExpRange"]]: + from dvc.repo.experiments.collect import collect_exec_branch + from dvc.repo.experiments.serialize import ( + ExpExecutor, + ExpRange, + LocalExpExecutor, + ) + + result: Dict[str, List[ExpRange]] = defaultdict(list) + pid = self._active_pid + if pid is None: + return result + + assert self._EXEC_NAME + infofile = self.get_infofile_path(self._EXEC_NAME) + try: + info = ExecutorInfo.from_dict(load_json(infofile)) + except OSError: + return result + + if ( + (not baseline_revs or info.baseline_rev in baseline_revs) + and info.status < TaskStatus.FAILED + and info.status != TaskStatus.SUCCESS + ): + local_exec = LocalExpExecutor(root=info.root_dir, pid=pid) + exps = list(collect_exec_branch(self.repo, info.baseline_rev, **kwargs)) + exps[0].name = info.name + result[info.baseline_rev] = [ + ExpRange( + exps, + executor=ExpExecutor("running", name="workspace", local=local_exec), + name=info.name, + ) + ] + return result + + def collect_queued_data( + self, + baseline_revs: Optional[Collection[str]], + **kwargs, + ) -> Dict[str, List["ExpRange"]]: + raise NotImplementedError + + def collect_failed_data( + self, + baseline_revs: Optional[Collection[str]], + **kwargs, + ) -> Dict[str, List["ExpRange"]]: + raise NotImplementedError