exp: standardize exp data collection

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)

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.
+        """

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

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

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