plots: pass all types of plots in --json option

Related: iterative/vscode-dvc#1757

.github/workflows/tests.yaml

@@ -66,8 +66,6 @@ jobs:
       run: |
         pip install --upgrade pip setuptools wheel
         pip install -e ".[dev]"
-      env:
-        AIOHTTP_NO_EXTENSIONS: ${{ matrix.pyv == '3.11-dev' && '1' }}
     - name: run tests
       timeout-minutes: 40
       run: pytest -nauto -ra --durations 100 --cov=dvc --cov-report=xml --cov-report=term -k "${{ matrix.pytest-filter }}"

.pre-commit-config.yaml

@@ -16,9 +16,9 @@ repos:
       - id: codespell
         args: 
           - --ignore-words-list
-          - ba,datas,fo,uptodate,cachable
+          - ba,datas,fo,uptodate,cachable,falsy
     repo: https://github.com/codespell-project/codespell
-    rev: v2.1.0
+    rev: v2.2.1
   - hooks:
       - id: isort
         language_version: python3

README.rst

@@ -82,7 +82,7 @@ To share and back up the *data cache*, DVC supports multiple remote storage plat
 
 |Flowchart|
 
-`DVC pipelines <https://dvc.org/doc/start/data-management/pipelines>`_ (computational graphs) connect code and data together.
+`DVC pipelines <https://dvc.org/doc/start/data-management/data-pipelines>`_ (computational graphs) connect code and data together.
 They specify all steps required to produce a model: input dependencies including code, data, commands to run; and output information to be saved.
 
 Last but not least, `DVC Experiment Versioning <https://dvc.org/doc/start/experiments>`_ lets you prepare and run a large number of experiments.

dvc/commands/checkout.py

@@ -110,7 +110,7 @@ def add_parser(subparsers, parent_parser):
         nargs="*",
         help=(
             "Limit command scope to these tracked files/directories, "
-            ".dvc files, or stage names."
+            ".dvc files and stage or foreach-group names."
         ),
     ).complete = completion.DVC_FILE
     checkout_parser.set_defaults(func=CmdCheckout)

dvc/commands/commit.py

@@ -70,7 +70,9 @@ def add_parser(subparsers, parent_parser):
     commit_parser.add_argument(
         "targets",
         nargs="*",
-        help="stages or .dvc files to commit. Optional. "
-        "(Finds all DVC files in the workspace by default.)",
+        help=(
+            "Limit command scope to these tracked files/directories, "
+            ".dvc files and stage or foreach-group names."
+        ),
     ).complete = completion.DVCFILES_AND_STAGE
     commit_parser.set_defaults(func=CmdCommit)

dvc/commands/dag.py

@@ -204,7 +204,7 @@ def add_parser(subparsers, parent_parser):
     dag_parser.add_argument(
         "target",
         nargs="?",
-        help="Stage or output to show pipeline for (optional). "
+        help="Stage name, foreach-group name or output to show pipeline for. "
         "Finds all stages in the workspace by default.",
     )
     dag_parser.set_defaults(func=CmdDAG)

dvc/commands/data_sync.py

@@ -120,7 +120,7 @@ def shared_parent_parser():
         nargs="*",
         help=(
             "Limit command scope to these tracked files/directories, "
-            ".dvc files, or stage names."
+            ".dvc files and stage or foreach-group names."
         ),
     ).complete = completion.DVC_FILE
 
@@ -195,7 +195,7 @@ def add_parser(subparsers, _parent_parser):
         "--glob",
         action="store_true",
         default=False,
-        help="Pull cache for targets matching shell-style wildcards.",
+        help=argparse.SUPPRESS,
     )
     pull_parser.set_defaults(func=CmdDataPull)
 

dvc/commands/git_hook.py

@@ -74,7 +74,7 @@ class CmdMergeDriver(CmdHookBase):
             our = Dvcfile(dvc, self.args.our, verify=False)
             their = Dvcfile(dvc, self.args.their, verify=False)
 
-            our.merge(ancestor, their)
+            our.merge(ancestor, their, allowed=["add", "remove", "change"])
 
             return 0
         finally:

dvc/commands/plots.py

@@ -25,7 +25,7 @@ def _show_json(renderers, split=False):
 
 
 def _adjust_vega_renderers(renderers):
-    from dvc.render import VERSION_FIELD
+    from dvc.render import REVISION_FIELD, VERSION_FIELD
     from dvc_render import VegaRenderer
 
     for r in renderers:
@@ -42,6 +42,7 @@ def _adjust_vega_renderers(renderers):
                         dp["rev"] = "::".join(vi.values())
             else:
                 for dp in r.datapoints:
+                    dp[REVISION_FIELD] = dp[VERSION_FIELD]["revision"]
                     dp.pop(VERSION_FIELD, {})
 
 
@@ -66,18 +67,6 @@ def _data_versions_count(renderer):
     return len(set(x))
 
 
-def _filter_unhandled_renderers(renderers):
-    # filtering out renderers currently unhandled by vscode extension
-    from dvc_render import VegaRenderer
-
-    def _is_json_viable(r):
-        return not (
-            isinstance(r, VegaRenderer) and _data_versions_count(r) > 1
-        )
-
-    return list(filter(_is_json_viable, renderers))
-
-
 class CmdPlots(CmdBase):
     def _func(self, *args, **kwargs):
         raise NotImplementedError
@@ -156,12 +145,10 @@ class CmdPlots(CmdBase):
                 templates_dir=self.repo.plots.templates_dir,
             )
             if self.args.json:
-                renderers = _filter_unhandled_renderers(renderers)
                 _show_json(renderers, self.args.split)
                 return 0
 
             _adjust_vega_renderers(renderers)
-
             if self.args.show_vega:
                 renderer = first(filter(lambda r: r.TYPE == "vega", renderers))
                 if renderer:

dvc/commands/repro.py

@@ -55,7 +55,16 @@ def add_arguments(repro_parser):
     repro_parser.add_argument(
         "targets",
         nargs="*",
-        help="Stages to reproduce. 'dvc.yaml' by default.",
+        help=(
+            """\
+Stages to reproduce. 'dvc.yaml' by default.
+The targets can be path to a dvc.yaml file or `.dvc` file,
+or a stage name or a foreach-group name from dvc.yaml file from
+current working directory. To run a stage/foreach-group from dvc.yaml
+from other directories, the target must be a path followed by colon `:`
+and then the stage name or foreach-group name.
+"""
+        ),
     ).complete = completion.DVCFILES_AND_STAGE
     repro_parser.add_argument(
         "-f",

dvc/dvcfile.py

@@ -172,7 +172,7 @@ class FileMixin:
     def dump(self, stage, **kwargs):
         raise NotImplementedError
 
-    def merge(self, ancestor, other):
+    def merge(self, ancestor, other, allowed=None):
         raise NotImplementedError
 
 
@@ -204,12 +204,12 @@ class SingleStageFile(FileMixin):
     def remove_stage(self, stage):  # pylint: disable=unused-argument
         self.remove()
 
-    def merge(self, ancestor, other):
+    def merge(self, ancestor, other, allowed=None):
         assert isinstance(ancestor, SingleStageFile)
         assert isinstance(other, SingleStageFile)
 
         stage = self.stage
-        stage.merge(ancestor.stage, other.stage)
+        stage.merge(ancestor.stage, other.stage, allowed=allowed)
         self.dump(stage)
 
 
@@ -237,10 +237,10 @@ class PipelineFile(FileMixin):
             self._dump_pipeline_file(stage)
 
         if update_lock:
-            self._dump_lockfile(stage)
+            self._dump_lockfile(stage, **kwargs)
 
-    def _dump_lockfile(self, stage):
-        self._lockfile.dump(stage)
+    def _dump_lockfile(self, stage, **kwargs):
+        self._lockfile.dump(stage, **kwargs)
 
     @staticmethod
     def _check_if_parametrized(stage, action: str = "dump") -> None:
@@ -310,7 +310,7 @@ class PipelineFile(FileMixin):
         else:
             super().remove()
 
-    def merge(self, ancestor, other):
+    def merge(self, ancestor, other, allowed=None):
         raise NotImplementedError
 
 
@@ -366,7 +366,7 @@ class Lockfile(FileMixin):
         return {SCHEMA_KWD: version}
 
     def dump(self, stage, **kwargs):
-        stage_data = serialize.to_lockfile(stage)
+        stage_data = serialize.to_lockfile(stage, **kwargs)
 
         with modify_yaml(self.path, fs=self.repo.fs) as data:
             version = LOCKFILE_VERSION.from_dict(data)
@@ -411,7 +411,7 @@ class Lockfile(FileMixin):
         else:
             self.remove()
 
-    def merge(self, ancestor, other):
+    def merge(self, ancestor, other, allowed=None):
         raise NotImplementedError
 
 

dvc/fs/callbacks.py

@@ -24,6 +24,7 @@ class RichCallback(Callback):
         bytes: bool = False,  # pylint: disable=redefined-builtin
         unit: str = None,
         disable: bool = False,
+        transient: bool = True,
     ) -> None:
         self._progress = progress
         self.disable = disable
@@ -35,6 +36,7 @@ class RichCallback(Callback):
             "visible": False,
             "progress_type": None if bytes else "summary",
         }
+        self._transient = transient
         self._stack = ExitStack()
         super().__init__(size=size, value=value)
 
@@ -47,7 +49,7 @@ class RichCallback(Callback):
             return self._progress
 
         progress = RichTransferProgress(
-            transient=True,
+            transient=self._transient,
             disable=self.disable,
             console=ui.error_console,
         )
@@ -61,7 +63,8 @@ class RichCallback(Callback):
         return self
 
     def close(self):
-        self.progress.clear_task(self.task)
+        if self._transient:
+            self.progress.clear_task(self.task)
         self._stack.close()
 
     def call(self, hook_name=None, **kwargs):
@@ -74,6 +77,9 @@ class RichCallback(Callback):
 
     def branch(self, path_1, path_2, kwargs, child: Optional[Callback] = None):
         child = child or RichCallback(
-            progress=self.progress, desc=path_1, bytes=True
+            progress=self.progress,
+            desc=path_1,
+            bytes=True,
+            transient=self._transient,
         )
         return super().branch(path_1, path_2, kwargs, child=child)

dvc/fs/dvc.py

@@ -10,9 +10,7 @@ from typing import TYPE_CHECKING, Any, Callable, Optional, Tuple, Type, Union
 from fsspec.spec import AbstractFileSystem
 from funcy import cached_property, wrap_prop, wrap_with
 
-from dvc.utils.fs import makedirs
 from dvc_objects.fs.base import FileSystem
-from dvc_objects.fs.callbacks import DEFAULT_CALLBACK
 from dvc_objects.fs.path import Path
 
 from .data import DataFileSystem
@@ -66,15 +64,7 @@ def _get_dvc_path(dvc_fs, subkey):
 
 
 class _DVCFileSystem(AbstractFileSystem):  # pylint:disable=abstract-method
-    """DVC + git-tracked files fs.
-
-    Args:
-        repo: DVC or git repo.
-        subrepos: traverse to subrepos (by default, it ignores subrepos)
-        repo_factory: A function to initialize subrepo with, default is Repo.
-        kwargs: Additional keyword arguments passed to the `DataFileSystem()`.
-    """
-
+    cachable = False
     root_marker = "/"
 
     def __init__(
@@ -86,6 +76,41 @@ class _DVCFileSystem(AbstractFileSystem):  # pylint:disable=abstract-method
         repo_factory: Optional[RepoFactory] = None,
         **repo_kwargs: Any,
     ) -> None:
+        """DVC + git-tracked files fs.
+
+        Args:
+            path (str, optional): URL or path to a DVC/Git repository.
+                Defaults to a DVC repository in the current working directory.
+                Both HTTP and SSH protocols are supported for remote Git repos
+                (e.g. [user@]server:project.git).
+            rev (str, optional): Any Git revision such as a branch or tag name,
+                a commit hash or a dvc experiment name.
+                Defaults to the default branch in case of remote repositories.
+                In case of a local repository, if rev is unspecified, it will
+                default to the working directory.
+                If the repo is not a Git repo, this option is ignored.
+            repo (:obj:`Repo`, optional): `Repo` instance.
+            subrepos (bool): traverse to subrepos.
+                By default, it ignores subrepos.
+            repo_factory (callable): A function to initialize subrepo with.
+                The default is `Repo`.
+
+        Examples:
+            - Opening a filesystem from repo in current working directory
+
+            >>> fs = DVCFileSystem()
+
+            - Opening a filesystem from local repository
+
+            >>> fs = DVCFileSystem("path/to/local/repository")
+
+            - Opening a remote repository
+
+            >>> fs = DVCFileSystem(
+            ...    "https://github.com/iterative/example-get-started",
+            ...    rev="main",
+            ... )
+        """
         from pygtrie import Trie
 
         super().__init__()
@@ -238,7 +263,8 @@ class _DVCFileSystem(AbstractFileSystem):  # pylint:disable=abstract-method
         dvc_path = _get_dvc_path(dvc_fs, subkey)
         return dvc_fs.open(dvc_path, mode=mode)
 
-    def isdvc(self, path, **kwargs):
+    def isdvc(self, path, **kwargs) -> bool:
+        """Is this entry dvc-tracked?"""
         key = self._get_key_from_relative(path)
         _, dvc_fs, subkey = self._get_subrepo_info(key)
         dvc_path = _get_dvc_path(dvc_fs, subkey)
@@ -295,27 +321,6 @@ class _DVCFileSystem(AbstractFileSystem):  # pylint:disable=abstract-method
 
         return infos
 
-    def get_file(  # pylint: disable=arguments-differ
-        self, rpath, lpath, callback=DEFAULT_CALLBACK, **kwargs
-    ):
-        key = self._get_key_from_relative(rpath)
-        fs_path = self._from_key(key)
-        fs = self.repo.fs
-
-        if self.isdir(rpath):
-            makedirs(lpath, exist_ok=True)
-            return None
-
-        try:
-            fs.get_file(fs_path, lpath, callback=callback, **kwargs)
-            return
-        except FileNotFoundError:
-            _, dvc_fs, subkey = self._get_subrepo_info(key)
-            if not dvc_fs:
-                raise
-        dvc_path = _get_dvc_path(dvc_fs, subkey)
-        dvc_fs.get_file(dvc_path, lpath, callback=callback, **kwargs)
-
     def info(self, path, **kwargs):
         key = self._get_key_from_relative(path)
         ignore_subrepos = kwargs.get("ignore_subrepos", True)

dvc/output.py

@@ -765,11 +765,8 @@ class Output:
         if self.remote:
             ret[self.PARAM_REMOTE] = self.remote
 
-        if (
-            self.use_cache
-            and self.is_in_repo
-            and self.hash_info.isdir
-            and (kwargs.get("with_files") or self.files is not None)
+        if self.hash_info.isdir and (
+            kwargs.get("with_files") or self.files is not None
         ):
             if self.obj:
                 obj = self.obj
@@ -1127,7 +1124,7 @@ class Output:
                 "unable to auto-merge outputs that are not directories"
             )
 
-    def merge(self, ancestor, other):
+    def merge(self, ancestor, other, allowed=None):
         from dvc_data.hashfile.tree import MergeError as TreeMergeError
         from dvc_data.hashfile.tree import du, merge
 
@@ -1144,7 +1141,11 @@ class Output:
 
         try:
             merged = merge(
-                self.odb, ancestor_info, self.hash_info, other.hash_info
+                self.odb,
+                ancestor_info,
+                self.hash_info,
+                other.hash_info,
+                allowed=allowed,
             )
         except TreeMergeError as exc:
             raise MergeError(str(exc)) from exc

dvc/render/convert.py

@@ -2,7 +2,13 @@ import json
 from collections import defaultdict
 from typing import Dict, List, Union
 
-from dvc.render import REVISION_FIELD, REVISIONS_KEY, SRC_FIELD, TYPE_KEY
+from dvc.render import (
+    REVISION_FIELD,
+    REVISIONS_KEY,
+    SRC_FIELD,
+    TYPE_KEY,
+    VERSION_FIELD,
+)
 from dvc.render.converter.image import ImageConverter
 from dvc.render.converter.vega import VegaConverter
 
@@ -23,7 +29,7 @@ def _get_converter(
 def _group_by_rev(datapoints):
     grouped = defaultdict(list)
     for datapoint in datapoints:
-        rev = datapoint.pop(REVISION_FIELD)
+        rev = datapoint.get(VERSION_FIELD, {}).get("revision")
         grouped[rev].append(datapoint)
     return dict(grouped)
 

dvc/render/converter/vega.py

@@ -127,13 +127,20 @@ class VegaConverter(Converter):
                 ),
             )
         )
+        self.steps.append(
+            (
+                "infer_y",
+                partial(
+                    self._infer_y,
+                ),
+            )
+        )
 
         self.steps.append(
             (
                 "generate_y",
                 partial(
                     self._generate_y_values,
-                    y_values=self.plot_properties.get("y", None),
                 ),
             )
         )
@@ -156,7 +163,7 @@ class VegaConverter(Converter):
             } - {None}
             self.inferred_properties["fields"] = fields
 
-    def _infer_y(self, datapoints: List[Dict]):
+    def _infer_y(self, datapoints: List[Dict], **kwargs):
         if "y" not in self.plot_properties:
             data_fields = list(first(datapoints))
             skip = (
@@ -178,6 +185,7 @@ class VegaConverter(Converter):
                     )
             else:
                 self.inferred_properties["y"] = inferred_y
+        return datapoints
 
     def convert(
         self,
@@ -197,7 +205,6 @@ class VegaConverter(Converter):
         processed = deepcopy(data)
 
         for step_name, step in self.steps:
-
             if step_name not in skip:
                 processed = step(  # type: ignore
                     processed,
@@ -205,25 +212,28 @@ class VegaConverter(Converter):
                     filename=filename,
                 )
 
-        self._infer_y(processed)  # type: ignore
         return processed, {**self.plot_properties, **self.inferred_properties}
 
     def _generate_y_values(  # noqa: C901
         self,
         datapoints: List[Dict],
-        y_values: Optional[Union[str, List, Dict]],
         revision: str,
         filename: str,
         **kwargs,
     ) -> List[Dict]:
 
+        y_values = self.plot_properties.get(
+            "y", None
+        ) or self.inferred_properties.get("y", None)
+
+        assert y_values is not None
+
         result = []
         properties_update = {}
 
         def _add_version_info(datapoint, version_info):
             tmp = datapoint.copy()
             tmp[VERSION_FIELD] = version_info
-            tmp[REVISION_FIELD] = version_info["revision"]
             return tmp
 
         def _version_info(revision, filename=None, field=None):
@@ -250,12 +260,6 @@ class VegaConverter(Converter):
 
             return tmp
 
-        if not y_values:
-            for dp in datapoints:
-                result.append(
-                    _add_version_info(dp, _version_info(revision, filename))
-                )
-
         if isinstance(y_values, str):
             for datapoint in datapoints:
                 result.append(
@@ -290,13 +294,15 @@ class VegaConverter(Converter):
             ):
                 # if we use the same field from all files,
                 # we dont have to generate it
+                y_field = all_fields.pop()
                 for datapoint in datapoints:
                     result.append(
                         _add_version_info(
-                            datapoint, _version_info(revision, filename)
+                            datapoint,
+                            _version_info(revision, filename, y_field),
                         )
                     )
-                properties_update.update({"y": all_fields.pop()})
+                properties_update.update({"y": y_field})
             else:
                 for def_filename, val in y_values.items():
                     if isinstance(val, str):

dvc/repo/experiments/__init__.py

@@ -16,7 +16,7 @@ from .exceptions import (
     InvalidExpRefError,
     MultipleBranchError,
 )
-from .executor.base import BaseExecutor, ExecutorInfo
+from .executor.base import BaseExecutor, ExecutorInfo, TaskStatus
 from .queue.base import BaseStashQueue, QueueEntry
 from .queue.celery import LocalCeleryQueue
 from .queue.tempdir import TempDirQueue
@@ -463,6 +463,8 @@ class Experiments:
             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.RUNNING:
+                    return result
                 last_rev = self.scm.get_ref(EXEC_BRANCH)
                 if last_rev:
                     result[last_rev] = info.asdict()

dvc/repo/experiments/executor/base.py

@@ -4,6 +4,7 @@ import pickle
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
 from dataclasses import asdict, dataclass
+from enum import IntEnum
 from functools import partial
 from typing import (
     TYPE_CHECKING,
@@ -35,6 +36,8 @@ from ..refs import (
     EXEC_BASELINE,
     EXEC_BRANCH,
     EXEC_CHECKPOINT,
+    EXEC_HEAD,
+    EXEC_MERGE,
     EXEC_NAMESPACE,
     EXPS_NAMESPACE,
     EXPS_STASH,
@@ -47,10 +50,9 @@ if TYPE_CHECKING:
     from scmrepo.git import Git
 
     from dvc.repo import Repo
+    from dvc.repo.experiments.stash import ExpStashEntry
     from dvc.stage import PipelineStage
 
-    from ..base import ExpStashEntry
-
 logger = logging.getLogger(__name__)
 
 
@@ -64,6 +66,16 @@ class ExecutorResult(NamedTuple):
     force: bool
 
 
+class TaskStatus(IntEnum):
+    PENDING = 0
+    PREPARING = 1
+    RUNNING = 2
+    SUCCESS = 3
+    FAILED = 4
+    CANCELED = 5
+    FINISHED = 6
+
+
 @dataclass
 class ExecutorInfo:
     git_url: str
@@ -73,13 +85,16 @@ class ExecutorInfo:
     dvc_dir: str
     name: Optional[str] = None
     wdir: Optional[str] = None
-    collected: bool = False
     result_hash: Optional[str] = None
     result_ref: Optional[str] = None
     result_force: bool = False
+    status: TaskStatus = TaskStatus.PENDING
 
     @classmethod
     def from_dict(cls, d):
+        if "collected" in d:
+            if d.pop("collected"):
+                d["status"] = TaskStatus.FINISHED
         return cls(**d)
 
     def asdict(self):
@@ -136,6 +151,7 @@ class BaseExecutor(ABC):
         root_dir: str,
         dvc_dir: str,
         baseline_rev: str,
+        status: TaskStatus,
         wdir: Optional[str] = None,
         name: Optional[str] = None,
         location: Optional[str] = None,
@@ -149,9 +165,17 @@ class BaseExecutor(ABC):
         self.baseline_rev = baseline_rev
         self.location: str = location or self.DEFAULT_LOCATION
         self.result = result
+        self.status = status
 
     @abstractmethod
-    def init_git(self, scm: "Git", branch: Optional[str] = None):
+    def init_git(
+        self,
+        scm: "Git",
+        stash_rev: str,
+        entry: "ExpStashEntry",
+        infofile: Optional[str],
+        branch: Optional[str] = None,
+    ):
         """Init git repo and populate it using exp refs from the specified
         SCM instance.
         """
@@ -191,6 +215,7 @@ class BaseExecutor(ABC):
             dvc_dir=self.dvc_dir,
             name=self.name,
             wdir=self.wdir,
+            status=self.status,
             **result_dict,
         )
 
@@ -212,6 +237,7 @@ class BaseExecutor(ABC):
             root_dir=info.root_dir,
             dvc_dir=info.dvc_dir,
             baseline_rev=info.baseline_rev,
+            status=info.status,
             name=info.name,
             wdir=info.wdir,
             result=result,
@@ -222,7 +248,6 @@ class BaseExecutor(ABC):
     def from_stash_entry(
         cls: Type[_T],
         repo: "Repo",
-        stash_rev: str,
         entry: "ExpStashEntry",
         **kwargs,
     ) -> _T:
@@ -232,7 +257,6 @@ class BaseExecutor(ABC):
     def _from_stash_entry(
         cls: Type[_T],
         repo: "Repo",
-        stash_rev: str,
         entry: "ExpStashEntry",
         root_dir: str,
         **kwargs,
@@ -241,12 +265,11 @@ class BaseExecutor(ABC):
             root_dir=root_dir,
             dvc_dir=relpath(repo.dvc_dir, repo.scm.root_dir),
             baseline_rev=entry.baseline_rev,
+            status=TaskStatus.PREPARING,
             name=entry.name,
             wdir=relpath(os.getcwd(), repo.scm.root_dir),
             **kwargs,
         )
-        executor.init_git(repo.scm, branch=entry.branch)
-        executor.init_cache(repo, stash_rev)
         return executor
 
     @staticmethod
@@ -259,8 +282,12 @@ class BaseExecutor(ABC):
                 exp_data.update(to_lockfile(stage))
         return dict_sha256(exp_data)
 
-    def cleanup(self):
-        pass
+    def cleanup(self, infofile: str):
+        if infofile is not None:
+            info = ExecutorInfo.load_json(infofile)
+            if info.status < TaskStatus.FAILED:
+                info.status = TaskStatus.FINISHED
+            info.dump_json(infofile)
 
     # TODO: come up with better way to stash repro arguments
     @staticmethod
@@ -324,11 +351,12 @@ class BaseExecutor(ABC):
                 logger.debug("Replacing existing experiment '%s'", orig_ref)
                 return True
 
+            if on_diverged:
+                return on_diverged(orig_ref, has_checkpoint)
+
             self._raise_ref_conflict(
                 dest_scm, orig_ref, new_rev, has_checkpoint
             )
-            if on_diverged:
-                on_diverged(orig_ref, has_checkpoint)
             logger.debug("Reproduced existing experiment '%s'", orig_ref)
             return False
 
@@ -411,11 +439,9 @@ class BaseExecutor(ABC):
         exp_ref: Optional["ExpRefInfo"] = None
         repro_force: bool = False
 
-        if infofile is not None:
-            info.dump_json(infofile)
-
         with cls._repro_dvc(
             info,
+            infofile,
             log_errors=log_errors,
             **kwargs,
         ) as dvc:
@@ -491,9 +517,6 @@ class BaseExecutor(ABC):
             info.result_ref = ref
             info.result_force = repro_force
 
-        if infofile is not None:
-            info.dump_json(infofile)
-
         # ideally we would return stages here like a normal repro() call, but
         # stages is not currently picklable and cannot be returned across
         # multiprocessing calls
@@ -550,6 +573,7 @@ class BaseExecutor(ABC):
     def _repro_dvc(
         cls,
         info: "ExecutorInfo",
+        infofile: str = None,
         log_errors: bool = True,
         **kwargs,
     ):
@@ -557,6 +581,9 @@ class BaseExecutor(ABC):
         from dvc.stage.monitor import CheckpointKilledError
 
         dvc = Repo(os.path.join(info.root_dir, info.dvc_dir))
+        info.status = TaskStatus.RUNNING
+        if infofile is not None:
+            info.dump_json(infofile)
         if cls.QUIET:
             dvc.scm_context.quiet = cls.QUIET
         old_cwd = os.getcwd()
@@ -568,15 +595,28 @@ class BaseExecutor(ABC):
         try:
             logger.debug("Running repro in '%s'", os.getcwd())
             yield dvc
+            info.status = TaskStatus.SUCCESS
+            if infofile is not None:
+                info.dump_json(infofile)
+
         except CheckpointKilledError:
+            info.status = TaskStatus.FAILED
+            if infofile is not None:
+                info.dump_json(infofile)
             raise
         except DvcException:
             if log_errors:
                 logger.exception("")
+            info.status = TaskStatus.FAILED
+            if infofile is not None:
+                info.dump_json(infofile)
             raise
         except Exception:
             if log_errors:
                 logger.exception("unexpected error")
+            info.status = TaskStatus.FAILED
+            if infofile is not None:
+                info.dump_json(infofile)
             raise
         finally:
             dvc.close()
@@ -714,3 +754,23 @@ class BaseExecutor(ABC):
         disable_other_loggers()
         if level is not None:
             dvc_logger.setLevel(level)
+
+    @contextmanager
+    def set_exec_refs(
+        self, scm: "Git", stash_rev: str, entry: "ExpStashEntry"
+    ):
+        try:
+            # Executor will be initialized with an empty git repo that
+            # we populate by pushing:
+            #   EXEC_HEAD - the base commit for this experiment
+            #   EXEC_MERGE - the unmerged changes (from our stash)
+            #       to be reproduced
+            #   EXEC_BASELINE - the baseline commit for this experiment
+            scm.set_ref(EXEC_HEAD, entry.head_rev)
+            scm.set_ref(EXEC_MERGE, stash_rev)
+            scm.set_ref(EXEC_BASELINE, entry.baseline_rev)
+            yield
+        finally:
+            for ref in (EXEC_HEAD, EXEC_MERGE, EXEC_BASELINE):
+                if scm.get_ref(ref):
+                    scm.remove_ref(ref)

dvc/repo/experiments/executor/local.py

@@ -19,7 +19,7 @@ from ..refs import (
     EXEC_MERGE,
     EXEC_NAMESPACE,
 )
-from .base import EXEC_TMP_DIR, BaseExecutor
+from .base import EXEC_TMP_DIR, BaseExecutor, TaskStatus
 
 if TYPE_CHECKING:
     from scmrepo.git import Git
@@ -46,8 +46,8 @@ class BaseLocalExecutor(BaseExecutor):
     def scm(self):
         return SCM(self.root_dir)
 
-    def cleanup(self):
-        super().cleanup()
+    def cleanup(self, infofile: str):
+        super().cleanup(infofile)
         self.scm.close()
         del self.scm
 
@@ -67,15 +67,28 @@ class TempDirExecutor(BaseLocalExecutor):
     QUIET = True
     DEFAULT_LOCATION = "tempdir"
 
-    def init_git(self, scm: "Git", branch: Optional[str] = None):
+    def init_git(
+        self,
+        scm: "Git",
+        stash_rev: str,
+        entry: "ExpStashEntry",
+        infofile: Optional[str],
+        branch: Optional[str] = None,
+    ):
         from dulwich.repo import Repo as DulwichRepo
 
         from ..utils import push_refspec
 
         DulwichRepo.init(os.fspath(self.root_dir))
 
-        refspec = f"{EXEC_NAMESPACE}/"
-        push_refspec(scm, self.git_url, refspec, refspec)
+        self.status = TaskStatus.PREPARING
+        if infofile:
+            self.info.dump_json(infofile)
+
+        with self.set_exec_refs(scm, stash_rev, entry):
+            refspec = f"{EXEC_NAMESPACE}/"
+            push_refspec(scm, self.git_url, refspec, refspec)
+
         if branch:
             push_refspec(scm, self.git_url, branch, branch)
             self.scm.set_ref(EXEC_BRANCH, branch, symbolic=True)
@@ -84,17 +97,12 @@ class TempDirExecutor(BaseLocalExecutor):
 
         if self.scm.get_ref(EXEC_CHECKPOINT):
             self.scm.remove_ref(EXEC_CHECKPOINT)
-
         # checkout EXEC_HEAD and apply EXEC_MERGE on top of it without
         # committing
         head = EXEC_BRANCH if branch else EXEC_HEAD
         self.scm.checkout(head, detach=True)
         merge_rev = self.scm.get_ref(EXEC_MERGE)
 
-        for ref in (EXEC_HEAD, EXEC_MERGE, EXEC_BASELINE):
-            if scm.get_ref(ref):
-                scm.remove_ref(ref)
-
         try:
             self.scm.merge(merge_rev, squash=True, commit=False)
         except _SCMError as exc:
@@ -114,8 +122,8 @@ class TempDirExecutor(BaseLocalExecutor):
         """Initialize DVC cache."""
         self._config(repo.odb.repo.path)
 
-    def cleanup(self):
-        super().cleanup()
+    def cleanup(self, infofile: str):
+        super().cleanup(infofile)
         logger.debug("Removing tmpdir '%s'", self.root_dir)
         remove(self.root_dir)
 
@@ -123,7 +131,6 @@ class TempDirExecutor(BaseLocalExecutor):
     def from_stash_entry(
         cls,
         repo: "Repo",
-        stash_rev: str,
         entry: "ExpStashEntry",
         wdir: Optional[str] = None,
         **kwargs,
@@ -132,9 +139,7 @@ class TempDirExecutor(BaseLocalExecutor):
         makedirs(parent_dir, exist_ok=True)
         tmp_dir = mkdtemp(dir=parent_dir)
         try:
-            executor = cls._from_stash_entry(
-                repo, stash_rev, entry, tmp_dir, **kwargs
-            )
+            executor = cls._from_stash_entry(repo, entry, tmp_dir, **kwargs)
             logger.debug("Init temp dir executor in '%s'", tmp_dir)
             return executor
         except Exception:
@@ -152,18 +157,31 @@ class WorkspaceExecutor(BaseLocalExecutor):
     def from_stash_entry(
         cls,
         repo: "Repo",
-        stash_rev: str,
         entry: "ExpStashEntry",
         **kwargs,
     ):
         root_dir = repo.scm.root_dir
-        executor = cls._from_stash_entry(
-            repo, stash_rev, entry, root_dir, **kwargs
+        executor: "WorkspaceExecutor" = cls._from_stash_entry(
+            repo, entry, root_dir, **kwargs
         )
         logger.debug("Init workspace executor in '%s'", root_dir)
         return executor
 
-    def init_git(self, scm: "Git", branch: Optional[str] = None):
+    def init_git(
+        self,
+        scm: "Git",
+        stash_rev: str,
+        entry: "ExpStashEntry",
+        infofile: Optional[str],
+        branch: Optional[str] = None,
+    ):
+        self.status = TaskStatus.PREPARING
+        if infofile:
+            self.info.dump_json(infofile)
+
+        scm.set_ref(EXEC_HEAD, entry.head_rev)
+        scm.set_ref(EXEC_MERGE, stash_rev)
+        scm.set_ref(EXEC_BASELINE, entry.baseline_rev)
         self._detach_stack.enter_context(
             self.scm.detach_head(
                 self.scm.get_ref(EXEC_HEAD),
@@ -184,7 +202,9 @@ class WorkspaceExecutor(BaseLocalExecutor):
     def init_cache(self, repo: "Repo", rev: str, run_cache: bool = True):
         pass
 
-    def cleanup(self):
+    def cleanup(self, infofile: str):
+        super().cleanup(infofile)
+        remove(os.path.dirname(infofile))
         with self._detach_stack:
             self.scm.remove_ref(EXEC_BASELINE)
             self.scm.remove_ref(EXEC_MERGE)
@@ -193,4 +213,3 @@ class WorkspaceExecutor(BaseLocalExecutor):
             checkpoint = self.scm.get_ref(EXEC_CHECKPOINT)
             if checkpoint and checkpoint != self._orig_checkpoint:
                 self.scm.set_ref(EXEC_APPLY, checkpoint)
-        super().cleanup()

dvc/repo/experiments/executor/manager/base.py

@@ -1,269 +0,0 @@
-import logging
-import os
-from abc import ABC
-from collections import defaultdict, deque
-from collections.abc import Mapping
-from typing import TYPE_CHECKING, Deque, Dict, Generator, Optional, Tuple, Type
-
-from dvc.proc.manager import ProcessManager
-
-from ...exceptions import CheckpointExistsError, ExperimentExistsError
-from ...refs import EXEC_BASELINE, EXEC_HEAD, EXEC_MERGE, ExpRefInfo
-from ...stash import ExpStashEntry
-from ..base import EXEC_PID_DIR, BaseExecutor
-from ..local import TempDirExecutor, WorkspaceExecutor
-
-if TYPE_CHECKING:
-    from scmrepo.git import Git
-
-    from dvc.repo import Repo
-
-logger = logging.getLogger(__name__)
-
-
-class BaseExecutorManager(ABC, Mapping):  # noqa: B024
-    """Manages executors for a collection of experiments to be run."""
-
-    EXECUTOR_CLS: Type = BaseExecutor
-
-    def __init__(
-        self,
-        scm: "Git",
-        wdir: str,
-        **kwargs,
-    ):
-        from dvc.utils.fs import makedirs
-
-        self.scm = scm
-        makedirs(wdir, exist_ok=True)
-        self.wdir = wdir
-        self.proc = ProcessManager(self.pid_dir)
-        self._attached: Dict[str, "BaseExecutor"] = {}
-        self._detached: Dict[str, "BaseExecutor"] = dict(self._load_infos())
-        self._queue: Deque[Tuple[str, "BaseExecutor"]] = deque()
-
-    def __getitem__(self, key: str) -> "BaseExecutor":
-        try:
-            return self._attached[key]
-        except KeyError:
-            pass
-        return self._detached[key]
-
-    def __iter__(self):
-        yield from self._attached
-        yield from self._detached
-
-    def __len__(self):
-        return len(self._attached) + len(self._detached)
-
-    @property
-    def pid_dir(self) -> str:
-        return os.path.join(self.wdir, EXEC_PID_DIR)
-
-    def enqueue(self, rev: str, executor: "BaseExecutor"):
-        assert rev not in self
-        self._queue.append((rev, executor))
-
-    def _load_infos(self) -> Generator[Tuple[str, "BaseExecutor"], None, None]:
-        import json
-        from urllib.parse import urlparse
-
-        from ..base import ExecutorInfo
-        from ..ssh import SSHExecutor
-
-        def make_executor(info: "ExecutorInfo"):
-            if info.git_url:
-                scheme = urlparse(info.git_url).scheme
-                if scheme == "file":
-                    cls: Type = TempDirExecutor
-                elif scheme == "ssh":
-                    cls = SSHExecutor
-                else:
-                    raise NotImplementedError
-            else:
-                cls = WorkspaceExecutor
-            return cls.from_info(info)
-
-        for name in self.proc:
-            infofile = self.get_infofile_path(name)
-            try:
-                with open(infofile, encoding="utf-8") as fobj:
-                    info = ExecutorInfo.from_dict(json.load(fobj))
-                yield name, make_executor(info)
-            except OSError:
-                continue
-
-    def get_infofile_path(self, name: str) -> str:
-        return os.path.join(
-            self.pid_dir,
-            name,
-            f"{name}{BaseExecutor.INFOFILE_EXT}",
-        )
-
-    @classmethod
-    def from_stash_entries(
-        cls,
-        scm: "Git",
-        wdir: str,
-        repo: "Repo",
-        to_run: Dict[str, ExpStashEntry],
-        **kwargs,
-    ):
-        manager = cls(scm, wdir)
-        manager._enqueue_stash_entries(scm, repo, to_run, **kwargs)
-        return manager
-
-    def _enqueue_stash_entries(
-        self,
-        scm: "Git",
-        repo: "Repo",
-        to_run: Dict[str, ExpStashEntry],
-        **kwargs,
-    ):
-        try:
-            for stash_rev, entry in to_run.items():
-                scm.set_ref(EXEC_HEAD, entry.head_rev)
-                scm.set_ref(EXEC_MERGE, stash_rev)
-                scm.set_ref(EXEC_BASELINE, entry.baseline_rev)
-
-                # Executor will be initialized with an empty git repo that
-                # we populate by pushing:
-                #   EXEC_HEAD - the base commit for this experiment
-                #   EXEC_MERGE - the unmerged changes (from our stash)
-                #       to be reproduced
-                #   EXEC_BASELINE - the baseline commit for this experiment
-                executor = self.EXECUTOR_CLS.from_stash_entry(
-                    repo,
-                    stash_rev,
-                    entry,
-                    **kwargs,
-                )
-                self.enqueue(stash_rev, executor)
-        finally:
-            for ref in (EXEC_HEAD, EXEC_MERGE, EXEC_BASELINE):
-                scm.remove_ref(ref)
-
-    def exec_queue(
-        self, repo: "Repo", jobs: Optional[int] = 1, detach: bool = False
-    ):
-        """Run dvc repro for queued executors in parallel."""
-        if detach:
-            raise NotImplementedError
-            # TODO use ProcessManager.spawn() to support detached runs
-        return self._exec_attached(repo, jobs=jobs)
-
-    def _exec_attached(self, repo: "Repo", jobs: Optional[int] = 1):
-        import signal
-        from concurrent.futures import (
-            CancelledError,
-            ProcessPoolExecutor,
-            wait,
-        )
-        from multiprocessing import Manager
-
-        from dvc.stage.monitor import CheckpointKilledError
-
-        result: Dict[str, Dict[str, str]] = defaultdict(dict)
-
-        manager = Manager()
-        pid_q = manager.Queue()
-
-        with ProcessPoolExecutor(max_workers=jobs) as workers:
-            futures = {}
-            while self._queue:
-                rev, executor = self._queue.popleft()
-                infofile = self.get_infofile_path(rev)
-                future = workers.submit(
-                    executor.reproduce,
-                    info=executor.info,
-                    rev=rev,
-                    queue=pid_q,
-                    infofile=infofile,
-                    log_level=logger.getEffectiveLevel(),
-                )
-                futures[future] = (rev, executor)
-                self._attached[rev] = executor
-
-            try:
-                wait(futures)
-            except KeyboardInterrupt:
-                # forward SIGINT to any running executor processes and
-                # cancel any remaining futures
-                workers.shutdown(wait=False)
-                pids = {}
-                for future, (rev, _) in futures.items():
-                    if future.running():
-                        # if future has already been started by the scheduler
-                        # we still have to wait until it tells us its PID
-                        while rev not in pids:
-                            rev, pid = pid_q.get()
-                            pids[rev] = pid
-                        os.kill(pids[rev], signal.SIGINT)
-                    elif not future.done():
-                        future.cancel()
-
-            for future, (rev, executor) in futures.items():
-                rev, executor = futures[future]
-
-                try:
-                    exc = future.exception()
-                    if exc is None:
-                        exec_result = future.result()
-                        result[rev].update(
-                            self._collect_executor(repo, executor, exec_result)
-                        )
-                    elif not isinstance(exc, CheckpointKilledError):
-                        logger.error(
-                            "Failed to reproduce experiment '%s'", rev[:7]
-                        )
-                except CancelledError:
-                    logger.error(
-                        "Cancelled before attempting to reproduce experiment "
-                        "'%s'",
-                        rev[:7],
-                    )
-                finally:
-                    self.cleanup_executor(rev, executor)
-
-        return result
-
-    def _collect_executor(self, repo, executor, exec_result) -> Dict[str, str]:
-        # NOTE: GitPython Repo instances cannot be re-used
-        # after process has received SIGINT or SIGTERM, so we
-        # need this hack to re-instantiate git instances after
-        # checkpoint runs. See:
-        # https://github.com/gitpython-developers/GitPython/issues/427
-        # del self.repo.scm
-
-        results = {}
-
-        def on_diverged(ref: str, checkpoint: bool):
-            ref_info = ExpRefInfo.from_ref(ref)
-            if checkpoint:
-                raise CheckpointExistsError(ref_info.name)
-            raise ExperimentExistsError(ref_info.name)
-
-        for ref in executor.fetch_exps(
-            self.scm,
-            force=exec_result.force,
-            on_diverged=on_diverged,
-        ):
-            exp_rev = self.scm.get_ref(ref)
-            if exp_rev:
-                logger.debug("Collected experiment '%s'.", exp_rev[:7])
-                results[exp_rev] = exec_result.exp_hash
-
-        if exec_result.ref_info is not None:
-            executor.collect_cache(repo, exec_result.ref_info)
-
-        return results
-
-    def cleanup_executor(self, rev: str, executor: "BaseExecutor"):
-        from dvc.utils.fs import remove
-
-        executor.cleanup()
-        try:
-            self.proc.remove(rev)
-        except KeyError:
-            pass
-        remove(os.path.join(self.pid_dir, rev))

dvc/repo/experiments/executor/manager/local.py

@@ -1,107 +0,0 @@
-import logging
-from collections import defaultdict
-from typing import TYPE_CHECKING, Dict, Optional
-
-from ...refs import EXEC_BASELINE, EXEC_BRANCH, EXEC_HEAD, EXEC_MERGE
-from ...stash import ExpStashEntry
-from ..local import TempDirExecutor, WorkspaceExecutor
-from .base import BaseExecutorManager
-
-if TYPE_CHECKING:
-    from scmrepo.git import Git
-
-    from dvc.repo import Repo
-
-logger = logging.getLogger(__name__)
-
-
-class TempDirExecutorManager(BaseExecutorManager):
-    EXECUTOR_CLS = TempDirExecutor
-
-
-class WorkspaceExecutorManager(BaseExecutorManager):
-    EXECUTOR_CLS = WorkspaceExecutor
-
-    @classmethod
-    def from_stash_entries(
-        cls,
-        scm: "Git",
-        wdir: str,
-        repo: "Repo",
-        to_run: Dict[str, ExpStashEntry],
-        **kwargs,
-    ):
-        manager = cls(scm, wdir)
-        try:
-            assert len(to_run) == 1
-            for stash_rev, entry in to_run.items():
-                scm.set_ref(EXEC_HEAD, entry.head_rev)
-                scm.set_ref(EXEC_MERGE, stash_rev)
-                scm.set_ref(EXEC_BASELINE, entry.baseline_rev)
-
-                executor = cls.EXECUTOR_CLS.from_stash_entry(
-                    repo,
-                    stash_rev,
-                    entry,
-                    **kwargs,
-                )
-                manager.enqueue(stash_rev, executor)
-        finally:
-            for ref in (EXEC_MERGE,):
-                scm.remove_ref(ref)
-        return manager
-
-    def _collect_executor(self, repo, executor, exec_result) -> Dict[str, str]:
-        results = {}
-        exp_rev = self.scm.get_ref(EXEC_BRANCH)
-        if exp_rev:
-            logger.debug("Collected experiment '%s'.", exp_rev[:7])
-            results[exp_rev] = exec_result.exp_hash
-        return results
-
-    def exec_queue(
-        self, repo: "Repo", jobs: Optional[int] = 1, detach: bool = False
-    ):
-        """Run a single WorkspaceExecutor.
-
-        Workspace execution is done within the main DVC process
-        (rather than in multiprocessing context)
-        """
-        from dvc.exceptions import DvcException
-        from dvc.stage.monitor import CheckpointKilledError
-
-        assert len(self._queue) == 1
-        assert not detach
-        result: Dict[str, Dict[str, str]] = defaultdict(dict)
-        rev, executor = self._queue.popleft()
-
-        exec_name = "workspace"
-        infofile = self.get_infofile_path(exec_name)
-        try:
-            exec_result = executor.reproduce(
-                info=executor.info,
-                rev=rev,
-                infofile=infofile,
-                log_level=logger.getEffectiveLevel(),
-                log_errors=not isinstance(executor, WorkspaceExecutor),
-            )
-            if not exec_result.exp_hash:
-                raise DvcException(
-                    f"Failed to reproduce experiment '{rev[:7]}'"
-                )
-            if exec_result.ref_info:
-                result[rev].update(
-                    self._collect_executor(repo, executor, exec_result)
-                )
-        except CheckpointKilledError:
-            # Checkpoint errors have already been logged
-            return {}
-        except DvcException:
-            raise
-        except Exception as exc:
-            raise DvcException(
-                f"Failed to reproduce experiment '{rev[:7]}'"
-            ) from exc
-        finally:
-            self.cleanup_executor(exec_name, executor)
-        return result

dvc/repo/experiments/executor/manager/ssh.py

@@ -1,109 +0,0 @@
-import logging
-import posixpath
-from collections import defaultdict
-from typing import TYPE_CHECKING, Callable, Dict, Generator, Optional, Tuple
-
-from ...stash import ExpStashEntry
-from ..base import BaseExecutor
-from ..ssh import SSHExecutor, _sshfs
-from .base import BaseExecutorManager
-
-if TYPE_CHECKING:
-    from scmrepo.git import Git
-
-    from dvc.repo import Repo
-
-logger = logging.getLogger(__name__)
-
-
-class SSHExecutorManager(BaseExecutorManager):
-    EXECUTOR_CLS = SSHExecutor
-
-    def __init__(
-        self,
-        scm: "Git",
-        wdir: str,
-        host: Optional[str] = None,
-        port: Optional[int] = None,
-        username: Optional[str] = None,
-        fs_factory: Optional[Callable] = None,
-        **kwargs,
-    ):
-        assert host
-        super().__init__(scm, wdir, **kwargs)
-        self.host = host
-        self.port = port
-        self.username = username
-        self._fs_factory = fs_factory
-
-    def _load_infos(self) -> Generator[Tuple[str, "BaseExecutor"], None, None]:
-        # TODO: load existing infos using sshfs
-        yield from []
-
-    @classmethod
-    def from_stash_entries(
-        cls,
-        scm: "Git",
-        wdir: str,
-        repo: "Repo",
-        to_run: Dict[str, ExpStashEntry],
-        **kwargs,
-    ):
-        machine_name: Optional[str] = kwargs.get("machine_name", None)
-        manager = cls(
-            scm, wdir, **repo.machine.get_executor_kwargs(machine_name)
-        )
-        manager._enqueue_stash_entries(scm, repo, to_run, **kwargs)
-        return manager
-
-    def sshfs(self):
-        return _sshfs(self._fs_factory, host=self.host, port=self.port)
-
-    def get_infofile_path(self, name: str) -> str:
-        return f"{name}{BaseExecutor.INFOFILE_EXT}"
-
-    def _exec_attached(self, repo: "Repo", jobs: Optional[int] = 1):
-        from dvc.exceptions import DvcException
-        from dvc.stage.monitor import CheckpointKilledError
-
-        assert len(self._queue) == 1
-        result: Dict[str, Dict[str, str]] = defaultdict(dict)
-        rev, executor = self._queue.popleft()
-        info = executor.info
-        infofile = posixpath.join(
-            info.root_dir,
-            info.dvc_dir,
-            "tmp",
-            self.get_infofile_path(rev),
-        )
-        try:
-            exec_result = executor.reproduce(
-                info=executor.info,
-                rev=rev,
-                infofile=infofile,
-                log_level=logger.getEffectiveLevel(),
-                fs_factory=self._fs_factory,
-            )
-            if not exec_result.exp_hash:
-                raise DvcException(
-                    f"Failed to reproduce experiment '{rev[:7]}'"
-                )
-            if exec_result.ref_info:
-                result[rev].update(
-                    self._collect_executor(repo, executor, exec_result)
-                )
-        except CheckpointKilledError:
-            # Checkpoint errors have already been logged
-            return {}
-        except DvcException:
-            raise
-        except Exception as exc:
-            raise DvcException(
-                f"Failed to reproduce experiment '{rev[:7]}'"
-            ) from exc
-        finally:
-            self.cleanup_executor(rev, executor)
-        return result
-
-    def cleanup_executor(self, rev: str, executor: "BaseExecutor"):
-        executor.cleanup()

dvc/repo/experiments/executor/ssh.py

@@ -15,7 +15,7 @@ from ..refs import (
     EXEC_MERGE,
     EXEC_NAMESPACE,
 )
-from .base import BaseExecutor, ExecutorInfo, ExecutorResult
+from .base import BaseExecutor, ExecutorInfo, ExecutorResult, TaskStatus
 
 if TYPE_CHECKING:
     from queue import Queue
@@ -23,8 +23,8 @@ if TYPE_CHECKING:
     from scmrepo.git import Git
 
     from dvc.repo import Repo
-
-    from ..base import ExpRefInfo, ExpStashEntry
+    from dvc.repo.experiments.refs import ExpRefInfo
+    from dvc.repo.experiments.stash import ExpStashEntry
 
 logger = logging.getLogger(__name__)
 
@@ -75,14 +75,12 @@ class SSHExecutor(BaseExecutor):
     def from_stash_entry(
         cls,
         repo: "Repo",
-        stash_rev: str,
         entry: "ExpStashEntry",
         **kwargs,
     ):
         machine_name: Optional[str] = kwargs.pop("machine_name", None)
         executor = cls._from_stash_entry(
             repo,
-            stash_rev,
             entry,
             cls.gen_dirname(entry.name),
             location=machine_name,
@@ -119,9 +117,20 @@ class SSHExecutor(BaseExecutor):
         }
         return kwargs
 
-    def init_git(self, scm: "Git", branch: Optional[str] = None):
+    def init_git(
+        self,
+        scm: "Git",
+        stash_rev: str,
+        entry: "ExpStashEntry",
+        infofile: Optional[str],
+        branch: Optional[str] = None,
+    ):
         from ..utils import push_refspec
 
+        self.status = TaskStatus.PREPARING
+        if infofile:
+            self.info.dump_json(infofile)
+
         with self.sshfs() as fs:
             fs.makedirs(self.root_dir)
             self._ssh_cmd(fs, "git init .")
@@ -137,8 +146,11 @@ class SSHExecutor(BaseExecutor):
             # TODO: support multiple client key retries in git backends
             # (see https://github.com/iterative/dvc/issues/6508)
             kwargs = self._git_client_args(fs)
-            refspec = f"{EXEC_NAMESPACE}/"
-            push_refspec(scm, self.git_url, refspec, refspec, **kwargs)
+
+            with self.set_exec_refs(scm, stash_rev, entry):
+                refspec = f"{EXEC_NAMESPACE}/"
+                push_refspec(scm, self.git_url, refspec, refspec, **kwargs)
+
             if branch:
                 push_refspec(scm, self.git_url, branch, branch, **kwargs)
                 self._ssh_cmd(fs, f"git symbolic-ref {EXEC_BRANCH} {branch}")

dvc/repo/experiments/queue/base.py

@@ -33,7 +33,7 @@ from ..executor.base import (
     ExecutorResult,
 )
 from ..executor.local import WorkspaceExecutor
-from ..refs import EXEC_BASELINE, EXEC_HEAD, EXEC_MERGE, ExpRefInfo
+from ..refs import ExpRefInfo
 from ..stash import ExpStash, ExpStashEntry
 from ..utils import exp_refs_by_rev, scm_locked
 
@@ -561,7 +561,7 @@ class BaseStashQueue(ABC):
 
     @staticmethod
     @scm_locked
-    def setup_executor(
+    def init_executor(
         exp: "Experiments",
         queue_entry: QueueEntry,
         executor_cls: Type[BaseExecutor] = WorkspaceExecutor,
@@ -576,21 +576,23 @@ class BaseStashQueue(ABC):
         )
         if stash_entry.stash_index is not None:
             stash.drop(stash_entry.stash_index)
-
-        scm.set_ref(EXEC_HEAD, stash_entry.head_rev)
-        scm.set_ref(EXEC_MERGE, stash_rev)
-        scm.set_ref(EXEC_BASELINE, stash_entry.baseline_rev)
-
-        # Executor will be initialized with an empty git repo that
-        # we populate by pushing:
-        #   EXEC_HEAD - the base commit for this experiment
-        #   EXEC_MERGE - the unmerged changes (from our stash)
-        #       to be reproduced
-        #   EXEC_BASELINE - the baseline commit for this experiment
-        return executor_cls.from_stash_entry(
-            exp.repo, stash_rev, stash_entry, **kwargs
+        executor = executor_cls.from_stash_entry(
+            exp.repo, stash_entry, **kwargs
         )
 
+        infofile = exp.celery_queue.get_infofile_path(stash_rev)
+        executor.init_git(
+            exp.repo.scm,
+            stash_rev,
+            stash_entry,
+            infofile,
+            branch=stash_entry.branch,
+        )
+
+        executor.init_cache(exp.repo, stash_rev)
+
+        return executor
+
     def get_infofile_path(self, name: str) -> str:
         return os.path.join(
             self.pid_dir,

dvc/repo/experiments/queue/celery.py

@@ -23,7 +23,12 @@ from dvc.exceptions import DvcException
 from dvc.ui import ui
 
 from ..exceptions import UnresolvedQueueExpNamesError
-from ..executor.base import EXEC_TMP_DIR, ExecutorInfo, ExecutorResult
+from ..executor.base import (
+    EXEC_TMP_DIR,
+    ExecutorInfo,
+    ExecutorResult,
+    TaskStatus,
+)
 from .base import BaseStashQueue, QueueDoneResult, QueueEntry, QueueGetResult
 from .tasks import run_exp
 
@@ -248,7 +253,7 @@ class LocalCeleryQueue(BaseStashQueue):
 
         def _load_collected(rev: str) -> Optional[ExecutorResult]:
             executor_info = _load_info(rev)
-            if executor_info.collected:
+            if executor_info.status > TaskStatus.SUCCESS:
                 return executor_info.result
             raise FileNotFoundError
 

dvc/repo/experiments/queue/tasks.py

@@ -3,8 +3,6 @@ from typing import Any, Dict
 from celery import shared_task
 from celery.utils.log import get_task_logger
 
-from dvc.utils.fs import remove
-
 from ..executor.base import ExecutorInfo
 from ..executor.local import TempDirExecutor
 from .base import BaseStashQueue, QueueEntry
@@ -13,7 +11,7 @@ logger = get_task_logger(__name__)
 
 
 @shared_task
-def setup_exp(entry_dict: Dict[str, Any]) -> str:
+def setup_exp(entry_dict: Dict[str, Any]) -> TempDirExecutor:
     """Setup an experiment.
 
     Arguments:
@@ -28,7 +26,7 @@ def setup_exp(entry_dict: Dict[str, Any]) -> str:
     repo = Repo(entry.dvc_root)
     # TODO: split executor.init_cache into separate subtask - we can release
     # exp.scm_lock before DVC push
-    executor = BaseStashQueue.setup_executor(
+    executor = BaseStashQueue.init_executor(
         repo.experiments,
         entry,
         TempDirExecutor,
@@ -36,7 +34,7 @@ def setup_exp(entry_dict: Dict[str, Any]) -> str:
     )
     infofile = repo.experiments.celery_queue.get_infofile_path(entry.stash_rev)
     executor.info.dump_json(infofile)
-    return executor.root_dir
+    return executor
 
 
 @shared_task
@@ -79,25 +77,14 @@ def collect_exp(
 
 
 @shared_task
-def cleanup_exp(  # pylint: disable=unused-argument
-    tmp_dir: str, entry_dict: Dict[str, Any]
-) -> None:
+def cleanup_exp(executor: TempDirExecutor, infofile: str) -> None:
     """Cleanup after an experiment.
 
     Arguments:
         tmp_dir: Temp directory to be removed.
         entry_dict: Serialized QueueEntry for this experiment.
     """
-    remove(tmp_dir)
-
-
-def _set_collected(infofile: str):
-    try:
-        executor_info = ExecutorInfo.load_json(infofile)
-        executor_info.collected = True
-        executor_info.dump_json(infofile)
-    except FileNotFoundError:
-        pass
+    executor.cleanup(infofile)
 
 
 @shared_task
@@ -115,11 +102,10 @@ def run_exp(entry_dict: Dict[str, Any]) -> None:
     repo = Repo(entry.dvc_root)
     queue = repo.experiments.celery_queue
     infofile = queue.get_infofile_path(entry.stash_rev)
-    root_dir = setup_exp.s(entry_dict)()
+    executor = setup_exp.s(entry_dict)()
     try:
         cmd = ["dvc", "exp", "exec-run", "--infofile", infofile]
         proc_dict = queue.proc.run_signature(cmd, name=entry.stash_rev)()
         collect_exp.s(proc_dict, entry_dict)()
     finally:
-        cleanup_exp.s(root_dir, entry_dict)()
-        _set_collected(infofile)
+        cleanup_exp.s(executor, infofile)()

dvc/repo/experiments/queue/tempdir.py

@@ -11,6 +11,7 @@ from ..executor.base import (
     BaseExecutor,
     ExecutorInfo,
     ExecutorResult,
+    TaskStatus,
 )
 from ..executor.local import TempDirExecutor
 from .base import BaseStashQueue, QueueEntry, QueueGetResult
@@ -60,7 +61,7 @@ class TempDirQueue(WorkspaceQueue):
             stash_entry.name,
             stash_entry.head_rev,
         )
-        executor = self.setup_executor(
+        executor = self.init_executor(
             self.repo.experiments,
             entry,
             TempDirExecutor,
@@ -75,8 +76,9 @@ class TempDirQueue(WorkspaceQueue):
         for stash_rev in self.proc:
             infofile = self.get_infofile_path(stash_rev)
             executor_info = ExecutorInfo.load_json(infofile)
-            if not executor_info.collected and os.path.exists(
-                executor_info.root_dir
+            if (
+                not executor_info.status <= TaskStatus.SUCCESS
+                and os.path.exists(executor_info.root_dir)
             ):
                 yield QueueEntry(
                     self.repo.root_dir,

dvc/repo/experiments/queue/workspace.py

@@ -1,12 +1,10 @@
 import logging
-import os
 from collections import defaultdict
 from typing import TYPE_CHECKING, Collection, Dict, Generator, Optional
 
 from funcy import first
 
 from dvc.exceptions import DvcException
-from dvc.utils.fs import remove
 
 from ..exceptions import ExpQueueEmptyError
 from ..executor.base import BaseExecutor, ExecutorResult
@@ -41,7 +39,7 @@ class WorkspaceQueue(BaseStashQueue):
             stash_entry.name,
             stash_entry.head_rev,
         )
-        executor = self.setup_executor(self.repo.experiments, entry)
+        executor = self.init_executor(self.repo.experiments, entry)
         return QueueGetResult(entry, executor)
 
     def iter_queued(self) -> Generator[QueueEntry, None, None]:
@@ -118,9 +116,7 @@ class WorkspaceQueue(BaseStashQueue):
                 f"Failed to reproduce experiment '{rev[:7]}'"
             ) from exc
         finally:
-            if self._EXEC_NAME == exec_name:
-                remove(os.path.join(self.pid_dir, exec_name))
-            executor.cleanup()
+            executor.cleanup(infofile)
         return results
 
     @staticmethod
@@ -135,6 +131,7 @@ class WorkspaceQueue(BaseStashQueue):
             assert exec_result.exp_hash
             logger.debug("Collected experiment '%s'.", exp_rev[:7])
             results[exp_rev] = exec_result.exp_hash
+
         return results
 
     def get_result(self, entry: QueueEntry) -> Optional[ExecutorResult]:

dvc/repo/fetch.py

@@ -25,22 +25,6 @@ def _fetch_worktree(repo, remote):
             *key,
         )
     save(index)
-
-    for stage in repo.index.stages:
-        for out in stage.outs:
-            if not out.use_cache:
-                continue
-
-            if not out.is_in_repo:
-                continue
-
-            workspace, key = out.index_key
-            entry = repo.index.data[workspace][key]
-            out.hash_info = entry.hash_info
-            out.meta = entry.meta
-
-        stage.dvcfile.dump(stage)
-
     return len(index)
 
 

dvc/repo/index.py

@@ -192,9 +192,12 @@ class Index:
 
             data_index = by_workspace[workspace]
 
+            if out.files:
+                out.obj = out.get_obj()
+
             data_index[key] = DataIndexEntry(
                 meta=out.meta,
-                obj=out.get_obj() if out.files else out.obj,
+                obj=out.obj,
                 hash_info=out.hash_info,
                 odb=out.odb,
                 cache=out.odb,

dvc/repo/plots/__init__.py

@@ -100,7 +100,7 @@ class Plots:
                             "filename":
                             {
                                 "data_source": callable loading the data,
-                                "props": propreties for the file if it is
+                                "props": properties for the file if it is
                                          plots type output
                             }
                         }

dvc/repo/push.py

@@ -11,13 +11,11 @@ if TYPE_CHECKING:
 
 def _push_worktree(repo, remote):
     from dvc_data.hashfile.tree import tree_from_index
-    from dvc_data.index import build, checkout, collect
+    from dvc_data.index import checkout, collect
 
     index = repo.index.data["repo"]
     checkout(index, remote.path, remote.fs, force=True)
-    build(index, remote.path, remote.fs)
-    if any(out.isdir() for out in repo.index.outs):
-        collect(index, remote.path, remote.fs)
+    collect(index, remote.path, remote.fs, update=True)
 
     for stage in repo.index.stages:
         for out in stage.outs:

dvc/stage/__init__.py

@@ -568,6 +568,8 @@ class Stage(params.StageParams):
             if kwargs.get("checkpoint_func", None) or no_download:
                 allow_missing = True
             self.save(allow_missing=allow_missing)
+            if no_download:
+                self.ignore_outs()
             if not no_commit:
                 self.commit(allow_missing=allow_missing)
 
@@ -706,7 +708,7 @@ class Stage(params.StageParams):
                 "unable to auto-merge DVC files with deleted outputs"
             )
 
-    def merge(self, ancestor, other):
+    def merge(self, ancestor, other, allowed=None):
         assert other
 
         if not other.outs:
@@ -726,7 +728,7 @@ class Stage(params.StageParams):
         self._check_can_merge(self, ancestor_out)
         self._check_can_merge(other, ancestor_out)
 
-        self.outs[0].merge(ancestor_out, other.outs[0])
+        self.outs[0].merge(ancestor_out, other.outs[0], allowed=allowed)
 
     def dump(self, **kwargs):
         self.dvcfile.dump(self, **kwargs)
@@ -768,5 +770,5 @@ class PipelineStage(Stage):
     def _changed_stage_entry(self):
         return f"'cmd' of {self} has changed."
 
-    def merge(self, ancestor, other):
+    def merge(self, ancestor, other, allowed=None):
         raise NotImplementedError

dvc/stage/loader.py

@@ -73,6 +73,7 @@ class StageLoader(Mapping):
             item.hash_info = HashInfo(item.hash_name, hash_value)
             if item.hash_info and item.hash_info.isdir:
                 item.meta.isdir = True
+            item.files = get_in(checksums, [key, path, item.PARAM_FILES])
 
     @classmethod
     def load_stage(cls, dvcfile, name, stage_data, lock_data=None):

dvc/stage/serialize.py

@@ -152,7 +152,7 @@ def to_pipeline_file(stage: "PipelineStage"):
     }
 
 
-def to_single_stage_lockfile(stage: "Stage") -> dict:
+def to_single_stage_lockfile(stage: "Stage", **kwargs) -> dict:
     assert stage.cmd
 
     def _dumpd(item):
@@ -164,6 +164,13 @@ def to_single_stage_lockfile(stage: "Stage") -> dict:
             *meta_d.items(),
         ]
 
+        if item.hash_info.isdir and kwargs.get("with_files"):
+            if item.obj:
+                obj = item.obj
+            else:
+                obj = item.get_obj()
+            ret.append((item.PARAM_FILES, obj.as_list(with_meta=True)))
+
         return OrderedDict(ret)
 
     res = OrderedDict([("cmd", stage.cmd)])
@@ -183,9 +190,9 @@ def to_single_stage_lockfile(stage: "Stage") -> dict:
     return res
 
 
-def to_lockfile(stage: "PipelineStage") -> dict:
+def to_lockfile(stage: "PipelineStage", **kwargs) -> dict:
     assert stage.name
-    return {stage.name: to_single_stage_lockfile(stage)}
+    return {stage.name: to_single_stage_lockfile(stage, **kwargs)}
 
 
 def to_single_stage_file(stage: "Stage", **kwargs):

dvc/utils/hydra.py

@@ -86,8 +86,13 @@ def to_hydra_overrides(path_overrides):
 def get_hydra_sweeps(path_overrides):
     merged_overrides = []
     for path, overrides in path_overrides.items():
+        # `.` is reserved character in hydra syntax
+        # _merge_ is required to support sweeps across multiple files.
         merged_overrides.extend(
-            [f"{path}___{override}" for override in overrides]
+            [
+                f"{path.replace('.', '_')}_merge_{override}"
+                for override in overrides
+            ]
         )
 
     hydra_overrides = to_hydra_overrides(merged_overrides)
@@ -103,7 +108,7 @@ def get_hydra_sweeps(path_overrides):
     for split in splits:
         sweep_overrides = defaultdict(list)
         for merged_override in split:
-            path, override = merged_override.split("___")
-            sweep_overrides[path].append(override)
+            path, override = merged_override.split("_merge_")
+            sweep_overrides[path.replace("_", ".")].append(override)
         sweeps.append(dict(sweep_overrides))
     return sweeps

setup.cfg

@@ -22,6 +22,7 @@ classifiers =
     Programming Language :: Python :: 3.8
     Programming Language :: Python :: 3.9
     Programming Language :: Python :: 3.10
+    Programming Language :: Python :: 3.11
 
 [options]
 setup_requires =
@@ -66,8 +67,8 @@ install_requires =
     dvc-render==0.0.11
     dvc-task==0.1.2
     dvclive>=0.10.0
-    dvc-data==0.10.0
-    dvc-http==2.19.1
+    dvc-data==0.13.0
+    dvc-http==2.27.2
     hydra-core>=1.1.0
     iterative-telemetry==0.0.5
 

tests/func/experiments/conftest.py

@@ -1,122 +1,12 @@
-from textwrap import dedent
-
 import pytest
 
-from tests.func.test_repro_multistage import COPY_SCRIPT
-from tests.unit.repo.experiments.conftest import (  # noqa pylint disable=unused-argument
+from tests.unit.repo.experiments.conftest import (  # noqa, pylint disable=unused-argument
+    checkpoint_stage,
+    exp_stage,
+    failed_exp_stage,
     test_queue,
 )
 
-DEFAULT_ITERATIONS = 2
-CHECKPOINT_SCRIPT_FORMAT = dedent(
-    """\
-    import os
-    import sys
-    import shutil
-
-    from dvc.api import make_checkpoint
-
-    checkpoint_file = {}
-    checkpoint_iterations = int({})
-    if os.path.exists(checkpoint_file):
-        with open(checkpoint_file) as fobj:
-            try:
-                value = int(fobj.read())
-            except ValueError:
-                value = 0
-    else:
-        with open(checkpoint_file, "w"):
-            pass
-        value = 0
-
-    shutil.copyfile({}, {})
-
-    if os.getenv("DVC_CHECKPOINT"):
-        for _ in range(checkpoint_iterations):
-            value += 1
-            with open(checkpoint_file, "w") as fobj:
-                fobj.write(str(value))
-            make_checkpoint()
-"""
-)
-CHECKPOINT_SCRIPT = CHECKPOINT_SCRIPT_FORMAT.format(
-    "sys.argv[1]", "sys.argv[2]", "sys.argv[3]", "sys.argv[4]"
-)
-
-
-@pytest.fixture
-def exp_stage(tmp_dir, scm, dvc):
-    tmp_dir.gen("copy.py", COPY_SCRIPT)
-    tmp_dir.gen("params.yaml", "foo: 1")
-    stage = dvc.run(
-        cmd="python copy.py params.yaml metrics.yaml",
-        metrics_no_cache=["metrics.yaml"],
-        params=["foo"],
-        name="copy-file",
-        deps=["copy.py"],
-    )
-    scm.add(
-        [
-            "dvc.yaml",
-            "dvc.lock",
-            "copy.py",
-            "params.yaml",
-            "metrics.yaml",
-            ".gitignore",
-        ]
-    )
-    scm.commit("init")
-    return stage
-
-
-@pytest.fixture
-def checkpoint_stage(tmp_dir, scm, dvc, mocker):
-    mocker.patch("dvc.stage.run.Monitor.AWAIT", 0.01)
-
-    tmp_dir.gen("checkpoint.py", CHECKPOINT_SCRIPT)
-    tmp_dir.gen("params.yaml", "foo: 1")
-    stage = dvc.run(
-        cmd=(
-            f"python checkpoint.py foo {DEFAULT_ITERATIONS} "
-            "params.yaml metrics.yaml"
-        ),
-        metrics_no_cache=["metrics.yaml"],
-        params=["foo"],
-        checkpoints=["foo"],
-        deps=["checkpoint.py"],
-        no_exec=True,
-        name="checkpoint-file",
-    )
-    scm.add(["dvc.yaml", "checkpoint.py", "params.yaml", ".gitignore"])
-    scm.commit("init")
-    stage.iterations = DEFAULT_ITERATIONS
-    return stage
-
-
-@pytest.fixture
-def failed_exp_stage(tmp_dir, scm, dvc):
-    tmp_dir.gen("copy.py", COPY_SCRIPT)
-    tmp_dir.gen("params.yaml", "foo: 1")
-    stage = dvc.stage.add(
-        cmd="python -c 'import sys; sys.exit(1)'",
-        metrics_no_cache=["failed-metrics.yaml"],
-        params=["foo"],
-        name="failed-copy-file",
-        deps=["copy.py"],
-    )
-    scm.add(
-        [
-            "dvc.yaml",
-            "dvc.lock",
-            "copy.py",
-            "params.yaml",
-            "failed-metrics.yaml",
-            ".gitignore",
-        ]
-    )
-    scm.commit("init")
-    return stage
-
 
 @pytest.fixture
 def http_auth_patch(mocker):

tests/func/experiments/executor/test_ssh.py

@@ -29,7 +29,6 @@ def test_init_from_stash(tmp_dir, scm, dvc, machine_instance, mocker):
     mock_entry.name = ""
     SSHExecutor.from_stash_entry(
         dvc,
-        "abc123",
         mock_entry,
         machine_name="foo",
     )
@@ -39,14 +38,14 @@ def test_init_from_stash(tmp_dir, scm, dvc, machine_instance, mocker):
 
 @pytest.mark.needs_internet
 @pytest.mark.parametrize("cloud", [pytest.lazy_fixture("git_ssh")])
-def test_init_git(tmp_dir, scm, cloud):
+def test_init_git(tmp_dir, scm, cloud, mocker):
     tmp_dir.scm_gen({"foo": "foo", "dir": {"bar": "bar"}}, commit="init")
     baseline_rev = scm.get_rev()
-    scm.set_ref(EXEC_HEAD, baseline_rev)
     tmp_dir.gen("foo", "stashed")
     scm.gitpython.git.stash()
     rev = scm.resolve_rev("stash@{0}")
-    scm.set_ref(EXEC_MERGE, rev)
+
+    mock = mocker.Mock(baseline_rev=baseline_rev, head_rev=baseline_rev)
 
     root_url = cloud / SSHExecutor.gen_dirname()
 
@@ -59,7 +58,8 @@ def test_init_git(tmp_dir, scm, cloud):
         username=TEST_SSH_USER,
         fs_factory=partial(_ssh_factory, cloud),
     )
-    executor.init_git(scm)
+    infofile = str((root_url / "foo.run").path)
+    executor.init_git(scm, rev, mock, infofile=infofile)
     assert root_url.path == executor._repo_abspath
 
     fs = cloud._ssh
@@ -117,7 +117,6 @@ def test_reproduce(tmp_dir, scm, dvc, cloud, exp_stage, mocker):
     SSHExecutor.reproduce(
         info,
         rev,
-        infofile=infofile,
         fs_factory=partial(_ssh_factory, cloud),
     )
     assert mock_execute.called_once()

tests/func/experiments/test_show.py

@@ -12,6 +12,7 @@ from dvc.repo.experiments.executor.base import (
     EXEC_TMP_DIR,
     BaseExecutor,
     ExecutorInfo,
+    TaskStatus,
 )
 from dvc.repo.experiments.queue.base import QueueEntry
 from dvc.repo.experiments.refs import CELERY_STASH, ExpRefInfo
@@ -338,9 +339,16 @@ def test_show_sort(tmp_dir, scm, dvc, exp_stage, caplog):
 
 
 @pytest.mark.vscode
-def test_show_running_workspace(tmp_dir, scm, dvc, exp_stage, capsys):
+@pytest.mark.parametrize(
+    "status, running", [(TaskStatus.RUNNING, True), (TaskStatus.FAILED, False)]
+)
+def test_show_running_workspace(
+    tmp_dir, scm, dvc, exp_stage, capsys, status, running
+):
     pid_dir = os.path.join(dvc.tmp_dir, EXEC_TMP_DIR, EXEC_PID_DIR)
-    info = make_executor_info(location=BaseExecutor.DEFAULT_LOCATION)
+    info = make_executor_info(
+        location=BaseExecutor.DEFAULT_LOCATION, status=status
+    )
     pidfile = os.path.join(
         pid_dir,
         "workspace",
@@ -349,7 +357,8 @@ def test_show_running_workspace(tmp_dir, scm, dvc, exp_stage, capsys):
     makedirs(os.path.dirname(pidfile), True)
     (tmp_dir / pidfile).dump_json(info.asdict())
 
-    assert dvc.experiments.show()["workspace"] == {
+    print(dvc.experiments.show())
+    assert dvc.experiments.show().get("workspace") == {
         "baseline": {
             "data": {
                 "deps": {
@@ -363,18 +372,18 @@ def test_show_running_workspace(tmp_dir, scm, dvc, exp_stage, capsys):
                 "params": {"params.yaml": {"data": {"foo": 1}}},
                 "outs": {},
                 "queued": False,
-                "running": True,
-                "executor": info.location,
+                "running": True if running else False,
+                "executor": info.location if running else None,
                 "timestamp": None,
             }
         }
     }
-
     capsys.readouterr()
     assert main(["exp", "show", "--csv"]) == 0
     cap = capsys.readouterr()
-    assert "Running" in cap.out
-    assert info.location in cap.out
+    if running:
+        assert "Running" in cap.out
+        assert info.location in cap.out
 
 
 def test_show_running_tempdir(tmp_dir, scm, dvc, exp_stage, mocker):

tests/func/test_fs.py

@@ -1,7 +1,6 @@
 import os
 from operator import itemgetter
 
-from dvc.fs.callbacks import Callback
 from dvc.repo import Repo
 
 
@@ -42,57 +41,3 @@ def test_walk_dont_ignore_subrepos(tmp_dir, scm, dvc):
 
     assert set(get_dirs(next(dvc_fs.walk(path)))) == {".dvc", "subdir", ".git"}
     assert set(get_dirs(next(scm_fs.walk("/")))) == {".dvc", "subdir"}
-
-
-def test_callback_on_dvcfs(tmp_dir, dvc, scm, mocker):
-    tmp_dir.dvc_gen({"dir": {"bar": "bar"}}, commit="dvc")
-    tmp_dir.scm_gen({"dir": {"foo": "foo"}}, commit="git")
-
-    fs = dvc.dvcfs
-
-    callback = Callback()
-    fs.get(
-        "dir",
-        (tmp_dir / "dir2").fs_path,
-        callback=callback,
-    )
-
-    assert (tmp_dir / "dir2").read_text() == {"foo": "foo", "bar": "bar"}
-    assert callback.size == 2
-    assert callback.value == 2
-
-    callback = Callback()
-    branch = mocker.spy(callback, "branch")
-    fs.get(
-        os.path.join("dir", "foo"),
-        (tmp_dir / "foo").fs_path,
-        callback=callback,
-    )
-
-    size = os.path.getsize(tmp_dir / "dir" / "foo")
-    assert (tmp_dir / "foo").read_text() == "foo"
-    assert callback.size == 1
-    assert callback.value == 1
-
-    assert branch.call_count == 1
-    assert branch.spy_return.size == size
-    assert branch.spy_return.value == size
-
-    branch.reset_mock()
-
-    callback = Callback()
-    branch = mocker.spy(callback, "branch")
-    fs.get(
-        os.path.join("dir", "bar"),
-        (tmp_dir / "bar").fs_path,
-        callback=callback,
-    )
-
-    size = os.path.getsize(tmp_dir / "dir" / "bar")
-    assert (tmp_dir / "bar").read_text() == "bar"
-    assert callback.size == 1
-    assert callback.value == 1
-
-    assert branch.call_count == 1
-    assert branch.spy_return.size == size
-    assert branch.spy_return.value == size

tests/func/test_import_url.py

@@ -282,7 +282,7 @@ def test_import_url_to_remote_status(tmp_dir, dvc, local_cloud, local_remote):
     assert len(status) == 0
 
 
-def test_import_url_no_download(tmp_dir, dvc, local_workspace):
+def test_import_url_no_download(tmp_dir, scm, dvc, local_workspace):
     local_workspace.gen("file", "file content")
     dst = tmp_dir / "file"
     stage = dvc.imp_url(
@@ -292,6 +292,7 @@ def test_import_url_no_download(tmp_dir, dvc, local_workspace):
     assert stage.deps[0].hash_info.value == "d10b4c3ff123b26dc068d43a8bef2d23"
 
     assert not dst.exists()
+    assert scm.is_ignored(dst)
 
     out = stage.outs[0]
     assert not out.hash_info

tests/func/test_merge_driver.py

@@ -39,6 +39,22 @@ def _gen(tmp_dir, struct, name):
         (None, {"foo": "foo"}, {"bar": "bar"}, {"foo": "foo", "bar": "bar"}),
         (None, None, {"bar": "bar"}, {"bar": "bar"}),
         (None, {"foo": "foo"}, None, {"foo": "foo"}),
+        (
+            {"foo": "foo"},
+            {"foo": "bar"},
+            {"foo": "foo", "baz": "baz"},
+            {"foo": "bar", "baz": "baz"},
+        ),
+        ({"foo": "foo"}, {}, {"foo": "foo", "bar": "bar"}, {"bar": "bar"}),
+        (
+            {"common": "common", "subdir": {"foo": "foo", "bar": "bar"}},
+            {"common": "common", "subdir": {"foo": "foo", "bar": "baz"}},
+            {"common": "common", "subdir": {"bar": "bar", "bizz": "bizz"}},
+            {
+                "common": "common",
+                "subdir": {"bar": "baz", "bizz": "bizz"},
+            },
+        ),
     ],
 )
 def test_merge(tmp_dir, dvc, ancestor, our, their, merged):
@@ -74,18 +90,15 @@ def test_merge(tmp_dir, dvc, ancestor, our, their, merged):
             {"foo": "foo"},
             {"foo": "bar"},
             {"foo": "baz"},
-            (
-                "unable to auto-merge directories with "
-                "diff that contains 'change'ed files"
-            ),
+            ("unable to auto-merge the following paths:\nfoo"),
         ),
         (
             {"common": "common", "foo": "foo"},
             {"common": "common", "bar": "bar"},
             {"baz": "baz"},
             (
-                "unable to auto-merge directories with "
-                "diff that contains 'remove'ed files"
+                "unable to auto-merge the following paths:\n"
+                "both deleted: ('foo',)"
             ),
         ),
     ],

tests/func/utils/test_hydra.py

@@ -183,6 +183,13 @@ def test_compose_and_dump(tmp_dir, suffix, overrides, expected):
 @pytest.mark.parametrize(
     "overrides, expected",
     [
+        (
+            {"params.yaml": ["defaults/foo=1,2"]},
+            [
+                {"params.yaml": ["defaults/foo=1"]},
+                {"params.yaml": ["defaults/foo=2"]},
+            ],
+        ),
         (
             {"params.yaml": ["foo=1,2", "bar=3,4"]},
             [

tests/integration/plots/test_plots.py

@@ -11,6 +11,7 @@ from bs4 import BeautifulSoup
 from funcy import first
 
 from dvc.cli import main
+from dvc.render import REVISION_FIELD, VERSION_FIELD
 
 JSON_OUT = "vis_data"
 
@@ -60,17 +61,6 @@ def extract_vega_specs(html_path, plots_ids):
     return result
 
 
-def filter_fields(datapoints: List[Dict], fields: List[str]):
-
-    tmp = deepcopy(datapoints)
-    for datapoint in tmp:
-        keys = set(datapoint.keys())
-        for key in keys:
-            if key not in fields:
-                datapoint.pop(key)
-    return tmp
-
-
 def drop_fields(datapoints: List[Dict], fields: List[str]):
     tmp = deepcopy(datapoints)
     for datapoint in tmp:
@@ -219,10 +209,35 @@ def test_repo_with_plots(
     html_result = extract_vega_specs(
         html_path, ["linear.json", "confusion.json"]
     )
+
+    assert json_result["linear.json"][0]["content"]["data"][
+        "values"
+    ] == _update_datapoints(
+        linear_v1,
+        {
+            VERSION_FIELD: {
+                "revision": "workspace",
+                "filename": "linear.json",
+                "field": "y",
+            },
+        },
+    )
     assert html_result["linear.json"]["data"]["values"] == _update_datapoints(
         linear_v1,
         {
-            "rev": "workspace",
+            REVISION_FIELD: "workspace",
+        },
+    )
+    assert json_result["confusion.json"][0]["content"]["data"][
+        "values"
+    ] == _update_datapoints(
+        confusion_v1,
+        {
+            VERSION_FIELD: {
+                "revision": "workspace",
+                "filename": "confusion.json",
+                "field": "actual",
+            },
         },
     )
     assert html_result["confusion.json"]["data"][
@@ -230,10 +245,9 @@ def test_repo_with_plots(
     ] == _update_datapoints(
         confusion_v1,
         {
-            "rev": "workspace",
+            REVISION_FIELD: "workspace",
         },
     )
-
     verify_image(
         tmp_dir, "workspace", "image.png", image_v1, html_path, json_result
     )
@@ -280,17 +294,59 @@ def test_repo_with_plots(
             html_path,
             ["../linear.json", "../confusion.json"],
         )
+        assert json_result["../linear.json"][0]["content"]["data"][
+            "values"
+        ] == _update_datapoints(
+            linear_v2,
+            {
+                VERSION_FIELD: {
+                    "revision": "workspace",
+                    "filename": "../linear.json",
+                    "field": "y",
+                },
+            },
+        ) + _update_datapoints(
+            linear_v1,
+            {
+                VERSION_FIELD: {
+                    "revision": "HEAD",
+                    "filename": "../linear.json",
+                    "field": "y",
+                },
+            },
+        )
         assert html_result["../linear.json"]["data"][
             "values"
         ] == _update_datapoints(
             linear_v2,
             {
-                "rev": "workspace",
+                REVISION_FIELD: "workspace",
             },
         ) + _update_datapoints(
             linear_v1,
             {
-                "rev": "HEAD",
+                REVISION_FIELD: "HEAD",
+            },
+        )
+        assert json_result["../confusion.json"][0]["content"]["data"][
+            "values"
+        ] == _update_datapoints(
+            confusion_v2,
+            {
+                VERSION_FIELD: {
+                    "revision": "workspace",
+                    "filename": "../confusion.json",
+                    "field": "actual",
+                },
+            },
+        ) + _update_datapoints(
+            confusion_v1,
+            {
+                VERSION_FIELD: {
+                    "revision": "HEAD",
+                    "filename": "../confusion.json",
+                    "field": "actual",
+                },
             },
         )
         assert html_result["../confusion.json"]["data"][
@@ -298,12 +354,12 @@ def test_repo_with_plots(
         ] == _update_datapoints(
             confusion_v2,
             {
-                "rev": "workspace",
+                REVISION_FIELD: "workspace",
             },
         ) + _update_datapoints(
             confusion_v1,
             {
-                "rev": "HEAD",
+                REVISION_FIELD: "HEAD",
             },
         )
 
@@ -390,11 +446,9 @@ def test_repo_with_config_plots(tmp_dir, capsys, repo_with_config_plots):
     repo_state = repo_with_config_plots()
     plots = next(repo_state)
 
-    html_path, json_result, split_json_result = call(capsys)
+    html_path, _, __ = call(capsys)
 
     assert os.path.exists(html_path)
-    assert json_result == {}
-    assert split_json_result == {}
     html_result = extract_vega_specs(
         html_path,
         [
@@ -405,9 +459,15 @@ def test_repo_with_config_plots(tmp_dir, capsys, repo_with_config_plots):
     assert html_result["dvc.yaml::linear_train_vs_test"]["data"][
         "values"
     ] == _update_datapoints(
-        plots["data"]["linear_train.json"], {"rev": "linear_train.json"}
+        plots["data"]["linear_train.json"],
+        {
+            REVISION_FIELD: "linear_train.json",
+        },
     ) + _update_datapoints(
-        plots["data"]["linear_test.json"], {"rev": "linear_test.json"}
+        plots["data"]["linear_test.json"],
+        {
+            REVISION_FIELD: "linear_test.json",
+        },
     )
 
     # TODO check json results once vscode is able to handle flexible plots

tests/unit/command/test_plots.py

@@ -334,10 +334,6 @@ def test_plots_diff_json(dvc, mocker, capsys):
     )
     render_mock = mocker.patch("dvc_render.render_html")
 
-    show_json_mock = mocker.patch(
-        "dvc.commands.plots._filter_unhandled_renderers",
-        return_value=renderers,
-    )
     show_json_mock = mocker.patch("dvc.commands.plots._show_json")
 
     assert cmd.run() == 0

tests/unit/command/test_run.py

@@ -8,7 +8,7 @@ def test_run(mocker, dvc):
         [
             "run",
             "--name",
-            "nam",
+            "name",
             "--deps",
             "deps",
             "--outs",
@@ -86,7 +86,7 @@ def test_run(mocker, dvc):
         no_commit=True,
         always_changed=True,
         cmd="command",
-        name="nam",
+        name="name",
         single_stage=False,
         external=True,
         desc="description",

tests/unit/render/test_convert.py

@@ -1,4 +1,10 @@
-from dvc.render import REVISION_FIELD, REVISIONS_KEY, SRC_FIELD, TYPE_KEY
+from dvc.render import (
+    REVISION_FIELD,
+    REVISIONS_KEY,
+    SRC_FIELD,
+    TYPE_KEY,
+    VERSION_FIELD,
+)
 from dvc.render.convert import to_json
 
 
@@ -7,8 +13,18 @@ def test_to_json_vega(mocker):
     vega_renderer.TYPE = "vega"
     vega_renderer.get_filled_template.return_value = '{"this": "is vega"}'
     vega_renderer.datapoints = [
-        {"x": 1, "y": 2, REVISION_FIELD: "foo", "filename": "foo.json"},
-        {"x": 2, "y": 1, REVISION_FIELD: "bar", "filename": "foo.json"},
+        {
+            "x": 1,
+            "y": 2,
+            VERSION_FIELD: {"revision": "foo"},
+            "filename": "foo.json",
+        },
+        {
+            "x": 2,
+            "y": 1,
+            VERSION_FIELD: {"revision": "bar"},
+            "filename": "foo.json",
+        },
     ]
     result = to_json(vega_renderer)
     assert result[0] == {
@@ -17,10 +33,20 @@ def test_to_json_vega(mocker):
         "content": {"this": "is vega"},
         "datapoints": {
             "foo": [
-                {"x": 1, "y": 2, "filename": "foo.json"},
+                {
+                    "x": 1,
+                    "y": 2,
+                    "filename": "foo.json",
+                    VERSION_FIELD: {"revision": "foo"},
+                },
             ],
             "bar": [
-                {"x": 2, "y": 1, "filename": "foo.json"},
+                {
+                    "x": 2,
+                    "y": 1,
+                    "filename": "foo.json",
+                    VERSION_FIELD: {"revision": "bar"},
+                },
             ],
         },
     }
@@ -34,8 +60,18 @@ def test_to_json_vega_split(mocker):
         '{"this": "is split vega"}'
     )
     vega_renderer.datapoints = [
-        {"x": 1, "y": 2, REVISION_FIELD: "foo", "filename": "foo.json"},
-        {"x": 2, "y": 1, REVISION_FIELD: "bar", "filename": "foo.json"},
+        {
+            "x": 1,
+            "y": 2,
+            VERSION_FIELD: {"revision": "foo"},
+            "filename": "foo.json",
+        },
+        {
+            "x": 2,
+            "y": 1,
+            VERSION_FIELD: {"revision": "bar"},
+            "filename": "foo.json",
+        },
     ]
     result = to_json(vega_renderer, split=True)
     assert result[0] == {
@@ -43,8 +79,22 @@ def test_to_json_vega_split(mocker):
         REVISIONS_KEY: ["bar", "foo"],
         "content": {"this": "is split vega"},
         "datapoints": {
-            "foo": [{"x": 1, "y": 2, "filename": "foo.json"}],
-            "bar": [{"x": 2, "y": 1, "filename": "foo.json"}],
+            "foo": [
+                {
+                    "x": 1,
+                    "y": 2,
+                    "filename": "foo.json",
+                    VERSION_FIELD: {"revision": "foo"},
+                }
+            ],
+            "bar": [
+                {
+                    "x": 2,
+                    "y": 1,
+                    "filename": "foo.json",
+                    VERSION_FIELD: {"revision": "bar"},
+                }
+            ],
         },
     }
     vega_renderer.get_filled_template.assert_called_with(skip_anchors=["data"])

tests/unit/render/test_match.py

@@ -89,14 +89,20 @@ def test_match_renderers(mocker):
     assert len(renderers) == 1
     assert renderers[0].datapoints == [
         {
-            VERSION_FIELD: {"revision": "v1", "filename": "file.json"},
-            "rev": "v1",
+            VERSION_FIELD: {
+                "revision": "v1",
+                "filename": "file.json",
+                "field": "y",
+            },
             "x": 1,
             "y": 1,
         },
         {
-            VERSION_FIELD: {"revision": "v1", "filename": "file.json"},
-            "rev": "v1",
+            VERSION_FIELD: {
+                "revision": "v1",
+                "filename": "file.json",
+                "field": "y",
+            },
             "x": 2,
             "y": 2,
         },

tests/unit/render/test_vega_converter.py

@@ -57,42 +57,53 @@ def test_finding_lists(dictionary, expected_result):
 @pytest.mark.parametrize(
     "input_data,properties,expected_datapoints,expected_properties",
     [
-        (
-            # default x and y
+        pytest.param(
             {"metric": [{"v": 1}, {"v": 2}]},
             {},
             [
                 {
                     "v": 1,
                     "step": 0,
-                    VERSION_FIELD: {"revision": "r", "filename": "f"},
-                    "rev": "r",
+                    VERSION_FIELD: {
+                        "revision": "r",
+                        "filename": "f",
+                        "field": "v",
+                    },
                 },
                 {
                     "v": 2,
                     "step": 1,
-                    VERSION_FIELD: {"revision": "r", "filename": "f"},
-                    "rev": "r",
+                    VERSION_FIELD: {
+                        "revision": "r",
+                        "filename": "f",
+                        "field": "v",
+                    },
                 },
             ],
             {"x": "step", "y": "v"},
+            id="default_x_y",
         ),
-        (
-            # filter fields
+        pytest.param(
             {"metric": [{"v": 1, "v2": 0.1}, {"v": 2, "v2": 0.2}]},
             {"fields": {"v"}},
             [
                 {
                     "v": 1,
                     "step": 0,
-                    VERSION_FIELD: {"revision": "r", "filename": "f"},
-                    "rev": "r",
+                    VERSION_FIELD: {
+                        "revision": "r",
+                        "filename": "f",
+                        "field": "v",
+                    },
                 },
                 {
                     "v": 2,
                     "step": 1,
-                    VERSION_FIELD: {"revision": "r", "filename": "f"},
-                    "rev": "r",
+                    VERSION_FIELD: {
+                        "revision": "r",
+                        "filename": "f",
+                        "field": "v",
+                    },
                 },
             ],
             {
@@ -100,9 +111,9 @@ def test_finding_lists(dictionary, expected_result):
                 "y": "v",
                 "fields": {"v", "step"},
             },
+            id="filter_fields",
         ),
-        (
-            # choose x and y
+        pytest.param(
             {"metric": [{"v": 1, "v2": 0.1}, {"v": 2, "v2": 0.2}]},
             {"x": "v", "y": "v2"},
             [
@@ -114,7 +125,6 @@ def test_finding_lists(dictionary, expected_result):
                         "filename": "f",
                         "field": "v2",
                     },
-                    "rev": "r",
                 },
                 {
                     "v": 2,
@@ -124,13 +134,12 @@ def test_finding_lists(dictionary, expected_result):
                         "filename": "f",
                         "field": "v2",
                     },
-                    "rev": "r",
                 },
             ],
             {"x": "v", "y": "v2"},
+            id="choose_x_y",
         ),
-        (
-            # append x and y to filtered fields
+        pytest.param(
             {
                 "metric": [
                     {"v": 1, "v2": 0.1, "v3": 0.01, "v4": 0.001},
@@ -148,7 +157,6 @@ def test_finding_lists(dictionary, expected_result):
                         "filename": "f",
                         "field": "v4",
                     },
-                    "rev": "r",
                 },
                 {
                     "v": 2,
@@ -159,13 +167,12 @@ def test_finding_lists(dictionary, expected_result):
                         "filename": "f",
                         "field": "v4",
                     },
-                    "rev": "r",
                 },
             ],
             {"x": "v3", "y": "v4", "fields": {"v", "v3", "v4"}},
+            id="append_x_y_to_fields",
         ),
-        (
-            # find metric in nested structure
+        pytest.param(
             {
                 "some": "noise",
                 "very": {
@@ -184,7 +191,6 @@ def test_finding_lists(dictionary, expected_result):
                         "filename": "f",
                         "field": "v2",
                     },
-                    "rev": "r",
                 },
                 {
                     "v": 2,
@@ -194,13 +200,12 @@ def test_finding_lists(dictionary, expected_result):
                         "filename": "f",
                         "field": "v2",
                     },
-                    "rev": "r",
                 },
             ],
             {"x": "v", "y": "v2"},
+            id="find_in_nested_structure",
         ),
-        (
-            # provide list of fields in y def
+        pytest.param(
             {"metric": [{"v": 1, "v2": 0.1}, {"v": 2, "v2": 0.2}]},
             {"y": {"f": ["v", "v2"]}},
             [
@@ -212,7 +217,6 @@ def test_finding_lists(dictionary, expected_result):
                         "filename": "f",
                         "field": "v",
                     },
-                    "rev": "r",
                     "dvc_inferred_y_value": 1,
                     "step": 0,
                 },
@@ -224,7 +228,6 @@ def test_finding_lists(dictionary, expected_result):
                         "filename": "f",
                         "field": "v2",
                     },
-                    "rev": "r",
                     "dvc_inferred_y_value": 0.1,
                     "step": 0,
                 },
@@ -236,7 +239,6 @@ def test_finding_lists(dictionary, expected_result):
                         "filename": "f",
                         "field": "v",
                     },
-                    "rev": "r",
                     "dvc_inferred_y_value": 2,
                     "step": 1,
                 },
@@ -248,12 +250,12 @@ def test_finding_lists(dictionary, expected_result):
                         "filename": "f",
                         "field": "v2",
                     },
-                    "rev": "r",
                     "dvc_inferred_y_value": 0.2,
                     "step": 1,
                 },
             ],
             {"x": "step", "y": "dvc_inferred_y_value", "y_label": "y"},
+            id="y_def_list",
         ),
     ],
 )
@@ -275,6 +277,7 @@ def test_convert(
 def test_convert_skip_step():
     converter = VegaConverter()
     converter.skip_step("append_index")
+    converter.skip_step("generate_y")
 
     datapoints, resolved_properties = converter.convert(
         data={"a": "b", "metric": [{"v": 1}, {"v": 2}]},
@@ -285,13 +288,9 @@ def test_convert_skip_step():
     assert datapoints == [
         {
             "v": 1,
-            "rev": "r",
-            VERSION_FIELD: {"revision": "r", "filename": "f"},
         },
         {
             "v": 2,
-            "rev": "r",
-            VERSION_FIELD: {"revision": "r", "filename": "f"},
         },
     ]
     assert resolved_properties == {"x": "step", "y": "v"}

tests/unit/repo/experiments/conftest.py

@@ -1,8 +1,120 @@
 from contextlib import contextmanager
+from textwrap import dedent
 
 import pytest
 
 from dvc.repo.experiments.queue.celery import LocalCeleryQueue
+from tests.func.test_repro_multistage import COPY_SCRIPT
+
+DEFAULT_ITERATIONS = 2
+CHECKPOINT_SCRIPT_FORMAT = dedent(
+    """\
+    import os
+    import sys
+    import shutil
+
+    from dvc.api import make_checkpoint
+
+    checkpoint_file = {}
+    checkpoint_iterations = int({})
+    if os.path.exists(checkpoint_file):
+        with open(checkpoint_file) as fobj:
+            try:
+                value = int(fobj.read())
+            except ValueError:
+                value = 0
+    else:
+        with open(checkpoint_file, "w"):
+            pass
+        value = 0
+
+    shutil.copyfile({}, {})
+
+    if os.getenv("DVC_CHECKPOINT"):
+        for _ in range(checkpoint_iterations):
+            value += 1
+            with open(checkpoint_file, "w") as fobj:
+                fobj.write(str(value))
+            make_checkpoint()
+"""
+)
+CHECKPOINT_SCRIPT = CHECKPOINT_SCRIPT_FORMAT.format(
+    "sys.argv[1]", "sys.argv[2]", "sys.argv[3]", "sys.argv[4]"
+)
+
+
+@pytest.fixture
+def exp_stage(tmp_dir, scm, dvc):
+    tmp_dir.gen("copy.py", COPY_SCRIPT)
+    tmp_dir.gen("params.yaml", "foo: 1")
+    stage = dvc.run(
+        cmd="python copy.py params.yaml metrics.yaml",
+        metrics_no_cache=["metrics.yaml"],
+        params=["foo"],
+        name="copy-file",
+        deps=["copy.py"],
+    )
+    scm.add(
+        [
+            "dvc.yaml",
+            "dvc.lock",
+            "copy.py",
+            "params.yaml",
+            "metrics.yaml",
+            ".gitignore",
+        ]
+    )
+    scm.commit("init")
+    return stage
+
+
+@pytest.fixture
+def checkpoint_stage(tmp_dir, scm, dvc, mocker):
+    mocker.patch("dvc.stage.run.Monitor.AWAIT", 0.01)
+
+    tmp_dir.gen("checkpoint.py", CHECKPOINT_SCRIPT)
+    tmp_dir.gen("params.yaml", "foo: 1")
+    stage = dvc.run(
+        cmd=(
+            f"python checkpoint.py foo {DEFAULT_ITERATIONS} "
+            "params.yaml metrics.yaml"
+        ),
+        metrics_no_cache=["metrics.yaml"],
+        params=["foo"],
+        checkpoints=["foo"],
+        deps=["checkpoint.py"],
+        no_exec=True,
+        name="checkpoint-file",
+    )
+    scm.add(["dvc.yaml", "checkpoint.py", "params.yaml", ".gitignore"])
+    scm.commit("init")
+    stage.iterations = DEFAULT_ITERATIONS
+    return stage
+
+
+@pytest.fixture
+def failed_exp_stage(tmp_dir, scm, dvc):
+    tmp_dir.gen("copy.py", COPY_SCRIPT)
+    tmp_dir.gen("params.yaml", "foo: 1")
+    stage = dvc.stage.add(
+        cmd="python -c 'import sys; sys.exit(1)'",
+        metrics_no_cache=["failed-metrics.yaml"],
+        params=["foo"],
+        name="failed-copy-file",
+        deps=["copy.py"],
+    )
+    scm.add(
+        [
+            "dvc.yaml",
+            "dvc.lock",
+            "copy.py",
+            "params.yaml",
+            "failed-metrics.yaml",
+            ".gitignore",
+        ]
+    )
+    scm.commit("init")
+    return stage
 
 
 @contextmanager

tests/unit/repo/experiments/queue/test_local.py

@@ -6,9 +6,7 @@ from flaky.flaky_decorator import flaky
 
 from dvc.exceptions import DvcException
 from dvc.repo.experiments.exceptions import UnresolvedExpNamesError
-from dvc.repo.experiments.executor.local import TempDirExecutor
 from dvc.repo.experiments.queue.base import QueueDoneResult
-from dvc.repo.experiments.refs import EXEC_BASELINE, EXEC_HEAD, EXEC_MERGE
 
 
 def test_shutdown_no_tasks(test_queue, mocker):
@@ -110,20 +108,6 @@ def test_celery_queue_kill(test_queue, mocker):
     assert spy.called_once_with(mock_entry.stash_rev)
 
 
-def test_queue_clean_workspace_refs(git_dir, tmp_dir):
-
-    rev = git_dir.scm.get_rev()
-    exec_heads = (EXEC_MERGE, EXEC_BASELINE, EXEC_HEAD)
-    for ref in exec_heads:
-        git_dir.scm.set_ref(ref, rev)
-
-    temp_dir_executor = TempDirExecutor(str(tmp_dir), str(tmp_dir), rev)
-    temp_dir_executor.init_git(git_dir.scm)
-
-    for ref in exec_heads:
-        assert git_dir.scm.get_ref(ref) is None
-
-
 @pytest.mark.parametrize("status", ["FAILURE", "SUCCESS"])
 def test_queue_iter_done_task(test_queue, mocker, status):
 

tests/unit/repo/experiments/test_executor_status.py

@@ -0,0 +1,117 @@
+import os
+
+import pytest
+
+from dvc.exceptions import DvcException
+from dvc.repo.experiments.executor.base import ExecutorInfo, TaskStatus
+from dvc.repo.experiments.queue.tasks import (
+    cleanup_exp,
+    collect_exp,
+    setup_exp,
+)
+
+
+def test_celery_queue_success_status(dvc, scm, test_queue, exp_stage):
+    queue_entry = test_queue._stash_exp(
+        params={"params.yaml": ["foo=1"]},
+        targets=exp_stage.addressing,
+        name="success",
+    )
+    infofile = test_queue.get_infofile_path(queue_entry.stash_rev)
+    executor = setup_exp.s(queue_entry.asdict())()
+    executor_info = ExecutorInfo.load_json(infofile)
+    assert executor_info.status == TaskStatus.PREPARING
+
+    cmd = ["dvc", "exp", "exec-run", "--infofile", infofile]
+    proc_dict = test_queue.proc.run_signature(
+        cmd, name=queue_entry.stash_rev
+    )()
+
+    executor_info = ExecutorInfo.load_json(infofile)
+    assert executor_info.status == TaskStatus.SUCCESS
+
+    collect_exp.s(proc_dict, queue_entry.asdict())()
+    cleanup_exp.s(executor, infofile)()
+    executor_info = ExecutorInfo.load_json(infofile)
+    assert executor_info.status == TaskStatus.FINISHED
+
+
+def test_celery_queue_failure_status(dvc, scm, test_queue, failed_exp_stage):
+    queue_entry = test_queue._stash_exp(
+        params={"params.yaml": ["foo=1"]},
+        targets=failed_exp_stage.addressing,
+        name="failed",
+    )
+    infofile = test_queue.get_infofile_path(queue_entry.stash_rev)
+    setup_exp.s(queue_entry.asdict())()
+    cmd = ["dvc", "exp", "exec-run", "--infofile", infofile]
+    test_queue.proc.run_signature(cmd, name=queue_entry.stash_rev)()
+    executor_info = ExecutorInfo.load_json(infofile)
+    assert executor_info.status == TaskStatus.FAILED
+
+
+@pytest.mark.parametrize("queue_type", ["workspace_queue", "tempdir_queue"])
+def test_workspace_executor_success_status(dvc, scm, exp_stage, queue_type):
+    workspace_queue = getattr(dvc.experiments, queue_type)
+    queue_entry = workspace_queue.put(
+        params={"params.yaml": ["foo=1"]},
+        targets=exp_stage.addressing,
+        name="success",
+    )
+    name = workspace_queue._EXEC_NAME or queue_entry.stash_rev
+    infofile = workspace_queue.get_infofile_path(name)
+    entry, executor = workspace_queue.get()
+    rev = entry.stash_rev
+    exec_result = executor.reproduce(
+        info=executor.info,
+        rev=rev,
+        infofile=infofile,
+    )
+    executor_info = ExecutorInfo.load_json(infofile)
+    assert executor_info.status == TaskStatus.SUCCESS
+    if exec_result.ref_info:
+        workspace_queue.collect_executor(
+            dvc.experiments, executor, exec_result
+        )
+    executor.cleanup(infofile)
+
+    if queue_type == "tempdir_queue":
+        executor_info = ExecutorInfo.load_json(infofile)
+        assert executor_info.status == TaskStatus.FINISHED
+    else:
+        assert not os.path.exists(infofile)
+
+
+@pytest.mark.parametrize("queue_type", ["workspace_queue", "tempdir_queue"])
+def test_workspace_executor_failed_status(
+    dvc, scm, failed_exp_stage, queue_type
+):
+    workspace_queue = getattr(dvc.experiments, queue_type)
+    queue_entry = workspace_queue.put(
+        params={"params.yaml": ["foo=1"]},
+        targets=failed_exp_stage.addressing,
+        name="failed",
+    )
+    name = workspace_queue._EXEC_NAME or queue_entry.stash_rev
+
+    infofile = workspace_queue.get_infofile_path(name)
+    with pytest.raises(DvcException):
+        workspace_queue.reproduce()
+    if queue_type == "workspace_queue":
+        assert not os.path.exists(infofile)
+    else:
+        executor_info = ExecutorInfo.load_json(infofile)
+        assert executor_info.status == TaskStatus.FAILED
+
+
+def test_executor_status_compatiblity():
+    data = {
+        "git_url": "file:///Users/home",
+        "baseline_rev": "123",
+        "location": "dvc-task",
+        "root_dir": "/Users/home/8088/.dvc/tmp/exps/tmpx85892cx",
+        "dvc_dir": ".dvc",
+        "collected": True,
+    }
+    result = ExecutorInfo.from_dict(data)
+    assert result.status == TaskStatus.FINISHED