deps: bump dvc-data to 0.35.1

* add viztracer async option

* Update dvc/_debug.py

Co-authored-by: skshetry <[email protected]>

.flake8

@@ -1,12 +1,18 @@
 [flake8]
 ignore=
-    E203, # Whitespace before ':'
-    E266, # Too many leading '#' for block comment
-    W503, # Line break occurred before a binary operator
-    P1,  # unindexed parameters in the str.format, see:
+    # Whitespace before ':'
+    E203
+    # Too many leading '#' for block comment
+    E266
+    # Line break occurred before a binary operator
+    W503
+    # unindexed parameters in the str.format, see:
     # https://pypi.org/project/flake8-string-format/
+    P1
+    # use `!r` instead of surrounding by quotes
+    B028
 max_line_length = 79
-max-complexity = 15
+max-complexity = 10
 select = B,C,E,F,W,T4,B902,T,P
 show_source = true
 count = true

.github/workflows/tests.yaml

@@ -8,53 +8,22 @@ on:
     - cron: '5 1 * * *'  # every day at 01:05
   workflow_dispatch:
 
-env:
-  DVC_TEST: "true"
-  HOMEBREW_NO_AUTO_UPDATE: 1
-  SHELL: /bin/bash
-
 concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
   cancel-in-progress: true
 
-permissions:  # added using https://github.com/step-security/secure-workflows
+permissions:
   contents: read
 
 jobs:
-  lint:
-    timeout-minutes: 10
-    runs-on: ubuntu-latest
-    steps:
-    - uses: actions/checkout@v3
-      with:
-        fetch-depth: 0
-    - name: Set up Python 3.8
-      uses: actions/setup-python@v4
-      with:
-        python-version: 3.8
-        cache: pip
-        cache-dependency-path: pyproject.toml
-    - name: Install requirements
-      run: |
-        pip install --upgrade pip wheel
-        pip install ".[dev]"
-    - uses: pre-commit/[email protected]
   tests:
     timeout-minutes: 45
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
       matrix:
-        os: [ubuntu-20.04, windows-latest]
+        os: [ubuntu-20.04, windows-latest, macos-latest]
         pyv: ["3.8", "3.9", "3.10", "3.11"]
-        pytest-filter:
-        - "import or plot or live or experiment"
-        - "not (import or plot or live or experiment)"
-        include:
-        - {os: macos-latest, pyv: "3.8"}
-        - {os: macos-latest, pyv: "3.9"}
-        - {os: macos-latest, pyv: "3.10"}
-        - {os: macos-latest, pyv: "3.11"}
 
     steps:
     - uses: actions/checkout@v3
@@ -70,9 +39,12 @@ jobs:
       run: |
         pip install --upgrade pip wheel
         pip install -e ".[dev]"
+    - uses: pre-commit/[email protected]
     - name: run tests
-      timeout-minutes: 20
-      run: pytest -nauto --durations 100 --cov --cov-report=xml --cov-report=term -k "${{ matrix.pytest-filter }}"
+      timeout-minutes: 30
+      run: >
+        pytest -nauto --timeout=300 --durations=100
+        --cov --cov-report=xml --cov-report=term
     - name: upload coverage report
       uses: codecov/codecov-action@v3
       with:
@@ -80,7 +52,7 @@ jobs:
         fail_ci_if_error: false
   check:
     if: always()
-    needs: [lint, tests]
+    needs: [tests]
     runs-on: ubuntu-latest
     steps:
       - uses: re-actors/alls-green@release/v1
@@ -88,7 +60,7 @@ jobs:
           jobs: ${{ toJSON(needs) }}
   notify:
     if: github.ref == 'refs/heads/main' && failure()
-    needs: [lint, tests]
+    needs: [tests]
     runs-on: ubuntu-latest
     steps:
     - name: Slack Notification

.pre-commit-config.yaml

@@ -6,7 +6,7 @@ ci:
 
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.3.0
+    rev: v4.4.0
     hooks:
       - id: check-added-large-files
       - id: check-case-conflict
@@ -29,30 +29,34 @@ repos:
       - id: beautysh
         args: ["-i", "2"]
   - repo: https://github.com/asottile/pyupgrade
-    rev: v3.2.0
+    rev: v3.3.1
     hooks:
       - id: pyupgrade
         args:
           - --py38-plus
   - repo: https://github.com/psf/black
-    rev: 22.10.0
+    rev: 22.12.0
     hooks:
       - id: black
   - repo: https://github.com/PyCQA/isort
-    rev: 5.10.1
+    rev: 5.11.4
     hooks:
       - id: isort
+  - repo: https://github.com/charliermarsh/ruff-pre-commit
+    rev: 'v0.0.223'
+    hooks:
+      - id: ruff
+        args: ["--force-exclude"]
   - repo: https://github.com/codespell-project/codespell
     rev: v2.2.2
     hooks:
       - id: codespell
         additional_dependencies: ["tomli"]
   - repo: https://github.com/PyCQA/flake8
-    rev: 5.0.4
+    rev: 6.0.0
     hooks:
       - id: flake8
         additional_dependencies:
-          - flake8-broken-line
           - flake8-bugbear
           - flake8-comprehensions
           - flake8-debugger
@@ -65,11 +69,14 @@ repos:
       files: ^dvc/
       language: system
       types: [python]
+      require_serial: true
     - id: pylint
       name: pylint
       entry: pylint
+      files: ^dvc/
       language: system
       types: [python]
+      require_serial: true
   - hooks:
       - args:
           - git-hook

dvc/_debug.py

@@ -8,15 +8,21 @@ if TYPE_CHECKING:
 
 
 @contextmanager
-def viztracer_profile(path: Union[Callable[[], str], str], depth: int = -1):
+def viztracer_profile(
+    path: Union[Callable[[], str], str],
+    depth: int = -1,
+    log_async: bool = False,
+):
     try:
         import viztracer  # pylint: disable=import-error
     except ImportError:
-        print("Failed to run profiler, viztracer is not installed")
+        print(  # noqa: T201
+            "Failed to run profiler, viztracer is not installed"
+        )
         yield
         return
 
-    tracer = viztracer.VizTracer(max_stack_depth=depth)
+    tracer = viztracer.VizTracer(max_stack_depth=depth, log_async=log_async)
 
     tracer.start()
     yield
@@ -27,14 +33,14 @@ def viztracer_profile(path: Union[Callable[[], str], str], depth: int = -1):
 
 @contextmanager
 def yappi_profile(
-    path: Union[Callable[[], str], str] = None,
+    path: Optional[Union[Callable[[], str], str]] = None,
     wall_clock: Optional[bool] = True,
     separate_threads: Optional[bool] = False,
 ):
     try:
         import yappi  # pylint: disable=import-error
     except ImportError:
-        print("Failed to run profiler, yappi is not installed")
+        print("Failed to run profiler, yappi is not installed")  # noqa: T201
         yield
         return
 
@@ -60,7 +66,7 @@ def yappi_profile(
             st.save(out, type="callgrind")
         else:
             if ctx_id is not None:
-                print(f"\nThread {ctx_id}")
+                print(f"\nThread {ctx_id}")  # noqa: T201
             st.print_all()  # pylint:disable=no-member
             if ctx_id is None:
                 threads.print_all()  # pylint:disable=no-member
@@ -74,7 +80,9 @@ def instrument(html_output=False):
     try:
         from pyinstrument import Profiler  # pylint: disable=import-error
     except ImportError:
-        print("Failed to run profiler, pyinstrument is not installed")
+        print(  # noqa: T201
+            "Failed to run profiler, pyinstrument is not installed"
+        )
         yield
         return
 
@@ -87,11 +95,11 @@ def instrument(html_output=False):
     if html_output:
         profiler.open_in_browser()
         return
-    print(profiler.output_text(unicode=True, color=True))
+    print(profiler.output_text(unicode=True, color=True))  # noqa: T201
 
 
 @contextmanager
-def profile(dump_path: str = None):
+def profile(dump_path: Optional[str] = None):
     """Run a cprofile"""
     import cProfile
 
@@ -115,28 +123,31 @@ def debug():
         try:
             import ipdb as pdb  # noqa: T100, pylint: disable=import-error
         except ImportError:
-            import pdb  # noqa: T100
+            import pdb  # type: ignore[no-redef]  # noqa: T100
         pdb.post_mortem()
 
         raise  # prevent from jumping ahead
 
 
 def _sigshow(_, frame: Optional["FrameType"]) -> None:
+    import sys
     from shutil import get_terminal_size
     from traceback import format_stack
 
-    from dvc.ui import ui
-
     lines = "\u2015" * get_terminal_size().columns
-    ui.error_write(lines, "\n", *format_stack(frame), lines, sep="")
+    stack = format_stack(frame)
+    print(lines, "\n", *stack, lines, sep="", file=sys.stderr)  # noqa: T201
 
 
 @contextmanager
 def show_stack():
     r"""Show stack trace on SIGQUIT (Ctrl-\) or SIGINFO (Ctrl-T on macOS)."""
     import signal
+    import sys
+
+    if sys.platform != "win32":
+        signal.signal(signal.SIGQUIT, _sigshow)
 
-    signal.signal(signal.SIGQUIT, _sigshow)
     try:
         signal.signal(signal.SIGINFO, _sigshow)  # only available on macOS
     except AttributeError:
@@ -154,7 +165,7 @@ def _get_path_func(tool: str, ext: str):
 
 
 @contextmanager
-def debugtools(args: "Namespace" = None, **kwargs):
+def debugtools(args: Optional["Namespace"] = None, **kwargs):
     kw = vars(args) if args else {}
     kw.update(kwargs)
 
@@ -175,10 +186,17 @@ def debugtools(args: "Namespace" = None, **kwargs):
                     separate_threads=kw.get("yappi_separate_threads"),
                 )
             )
-        if kw.get("viztracer") or kw.get("viztracer_depth"):
+        if (
+            kw.get("viztracer")
+            or kw.get("viztracer_depth")
+            or kw.get("viztracer_async")
+        ):
             path_func = _get_path_func("viztracer", "json")
             depth = kw.get("viztracer_depth") or -1
-            prof = viztracer_profile(path=path_func, depth=depth)
+            log_async = kw.get("viztracer_async") or False
+            prof = viztracer_profile(
+                path=path_func, depth=depth, log_async=log_async
+            )
             stack.enter_context(prof)
         yield
 
@@ -202,6 +220,9 @@ def add_debugging_flags(parser):
         "--viztracer", action="store_true", default=False, help=SUPPRESS
     )
     parser.add_argument("--viztracer-depth", type=int, help=SUPPRESS)
+    parser.add_argument(
+        "--viztracer-async", action="store_true", default=False, help=SUPPRESS
+    )
     parser.add_argument("--cprofile-dump", help=SUPPRESS)
     parser.add_argument(
         "--pdb", action="store_true", default=False, help=SUPPRESS

dvc/analytics.py

@@ -125,7 +125,7 @@ def _system_info():
     system = platform.system()
 
     if system == "Windows":
-        version = sys.getwindowsversion()
+        version = sys.getwindowsversion()  # type: ignore[attr-defined]
 
         return {
             "os": "windows",

dvc/api/data.py

@@ -27,7 +27,7 @@ def get_url(path, repo=None, rev=None, remote=None):
         if not dvc_info:
             raise OutputNotFoundError(path, repo)
 
-        dvc_repo = info["repo"]
+        dvc_repo = info["repo"]  # pylint: disable=unsubscriptable-object
         md5 = dvc_info["md5"]
 
         return dvc_repo.cloud.get_url_for(remote, checksum=md5)
@@ -38,7 +38,11 @@ class _OpenContextManager(GCM):
         self, func, args, kwds
     ):  # pylint: disable=super-init-not-called
         self.gen = func(*args, **kwds)
-        self.func, self.args, self.kwds = func, args, kwds
+        self.func, self.args, self.kwds = (  # type: ignore[assignment]
+            func,
+            args,
+            kwds,
+        )
 
     def __getattr__(self, name):
         raise AttributeError(

dvc/api/params.py

@@ -1,3 +1,4 @@
+import typing
 from collections import Counter
 from typing import Dict, Iterable, Optional, Union
 
@@ -238,15 +239,15 @@ def params_show(
     if isinstance(stages, str):
         stages = [stages]
 
-    def _onerror_raise(result: Dict, exception: Exception, *args, **kwargs):
+    def _onerror_raise(_result: Dict, exception: Exception, *args, **kwargs):
         raise exception
 
     def _postprocess(params):
-        processed = {}
+        processed: Dict[str, Dict] = {}
         for rev, rev_data in params.items():
             processed[rev] = {}
 
-            counts = Counter()
+            counts: typing.Counter[str] = Counter()
             for file_data in rev_data["data"].values():
                 for k in file_data["data"]:
                     counts[k] += 1

dvc/cli/__init__.py

@@ -45,7 +45,6 @@ def _log_unknown_exceptions() -> None:
     if logger.isEnabledFor(logging.DEBUG):
         logger.debug("Version info for developers:\n%s", get_dvc_info())
     logger.info(FOOTER)
-    return None
 
 
 def _log_exceptions(exc: Exception) -> Optional[int]:
@@ -168,12 +167,12 @@ def main(argv=None):  # noqa: C901
         elif args.verbose == 1:
             level = logging.DEBUG
         elif args.verbose > 1:
-            level = logging.TRACE
+            level = logging.TRACE  # type: ignore[attr-defined]
 
         if level is not None:
             set_loggers_level(level)
 
-        logger.trace(args)
+        logger.trace(args)  # type: ignore[attr-defined]
 
         if not sys.stdout.closed and not args.quiet:
             from dvc.ui import ui

dvc/cli/command.py

@@ -1,6 +1,10 @@
 import logging
 import os
 from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from dvc.config import Config
 
 logger = logging.getLogger(__name__)
 
@@ -8,13 +12,13 @@ logger = logging.getLogger(__name__)
 class CmdBase(ABC):
     UNINITIALIZED = False
 
-    def __init__(self, args):
+    def __init__(self, args: Any):
         from dvc.repo import Repo
 
         os.chdir(args.cd)
 
         self.repo: "Repo" = Repo(uninitialized=self.UNINITIALIZED)
-        self.config = self.repo.config
+        self.config: "Config" = self.repo.config
         self.args = args
 
     def do_run(self):

dvc/commands/cache.py

@@ -42,7 +42,7 @@ def add_parser(subparsers, parent_parser):
 
     cache_subparsers = cache_parser.add_subparsers(
         dest="cmd",
-        help="Use `dvc cache CMD --help` for command-specific " "help.",
+        help="Use `dvc cache CMD --help` for command-specific help.",
     )
 
     fix_subparsers(cache_subparsers)

dvc/commands/check_ignore.py

@@ -103,7 +103,7 @@ def add_parser(subparsers, parent_parser):
         "--all",
         action="store_true",
         default=False,
-        help="Include the target paths which don’t match any pattern "
+        help="Include the target paths which don't match any pattern "
         "in the `--details` list.",
     )
     parser.add_argument(
@@ -111,7 +111,7 @@ def add_parser(subparsers, parent_parser):
         "--non-matching",
         action="store_true",
         default=False,
-        help="Include the target paths which don’t match any pattern "
+        help="Include the target paths which don't match any pattern "
         "in the `--details` list.",
     )
     parser.add_argument(

dvc/commands/config.py

@@ -37,13 +37,12 @@ class CmdConfig(CmdBaseNoRepo):
         self.config = Config(validate=False)
 
     def run(self):
-        if self.args.show_origin:
-            if any((self.args.value, self.args.unset)):
-                logger.error(
-                    "--show-origin can't be used together with any of these "
-                    "options: -u/--unset, value"
-                )
-                return 1
+        if self.args.show_origin and (self.args.value or self.args.unset):
+            logger.error(
+                "--show-origin can't be used together with any of these "
+                "options: -u/--unset, value"
+            )
+            return 1
 
         if self.args.list:
             return self._list()

dvc/commands/daemon.py

@@ -46,7 +46,7 @@ def add_parser(subparsers, parent_parser):
 
     daemon_subparsers = daemon_parser.add_subparsers(
         dest="cmd",
-        help="Use `dvc daemon CMD --help` for command-specific " "help.",
+        help="Use `dvc daemon CMD --help` for command-specific help.",
     )
 
     fix_subparsers(daemon_subparsers)

dvc/commands/data.py

@@ -72,7 +72,7 @@ class CmdDataStatus(CmdBase):
             yield stage, items
 
     @classmethod
-    def _show_status(cls, status: "DataStatus") -> int:
+    def _show_status(cls, status: "DataStatus") -> int:  # noqa: C901
         git_info = status.pop("git")  # type: ignore[misc]
         result = dict(cls._process_status(status))
         if not result:
@@ -114,7 +114,9 @@ class CmdDataStatus(CmdBase):
         return 0
 
     def run(self) -> int:
-        with log_durations(logger.trace, "in data_status"):  # type: ignore
+        with log_durations(
+            logger.trace, "in data_status"  # type: ignore[attr-defined]
+        ):
             status = self.repo.data_status(
                 granular=self.args.granular,
                 untracked_files=self.args.untracked_files,

dvc/commands/data_sync.py

@@ -122,7 +122,7 @@ def shared_parent_parser():
             "Limit command scope to these tracked files/directories, "
             ".dvc files and stage names."
         ),
-    ).complete = completion.DVC_FILE
+    ).complete = completion.DVC_FILE  # type: ignore[attr-defined]
 
     return parent_parser
 
@@ -302,7 +302,7 @@ def add_parser(subparsers, _parent_parser):
         "--with-deps",
         action="store_true",
         default=False,
-        help="Fetch cache for all dependencies of the " "specified target.",
+        help="Fetch cache for all dependencies of the specified target.",
     )
     fetch_parser.add_argument(
         "-R",

dvc/commands/diff.py

@@ -109,7 +109,7 @@ class CmdDiff(CmdBase):
             ui.write()
 
         if not sum(summary.values()):
-            return None
+            return
 
         states_summary = ", ".join(
             f"{summary[state]} {state}"

dvc/commands/experiments/__init__.py

@@ -18,6 +18,7 @@ from dvc.commands.experiments import (
     queue_worker,
     remove,
     run,
+    save,
     show,
 )
 
@@ -34,6 +35,7 @@ SUB_COMMANDS = [
     queue_worker,
     remove,
     run,
+    save,
     show,
 ]
 

dvc/commands/experiments/ls.py

@@ -3,6 +3,7 @@ import logging
 
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
+from dvc.ui import ui
 
 logger = logging.getLogger(__name__)
 
@@ -16,20 +17,13 @@ class CmdExperimentsList(CmdBase):
             num=self.args.num,
             git_remote=self.args.git_remote,
         )
-        tags = self.repo.scm.describe(exps)
-        remained = {baseline for baseline, tag in tags.items() if tag is None}
-        base = "refs/heads"
-        ref_heads = self.repo.scm.describe(remained, base=base)
 
         for baseline in exps:
-            name = baseline[:7]
-            if tags[baseline] or ref_heads[baseline]:
-                name = tags[baseline] or ref_heads[baseline][len(base) + 1 :]
             if not name_only:
-                print(f"{name}:")
+                ui.write(f"{baseline}:")
             for exp_name in exps[baseline]:
                 indent = "" if name_only else "\t"
-                print(f"{indent}{exp_name}")
+                ui.write(f"{indent}{exp_name}")
 
         return 0
 

dvc/commands/experiments/pull.py

@@ -68,7 +68,7 @@ def add_parser(experiments_subparsers, parent_parser):
         "-f",
         "--force",
         action="store_true",
-        help="Replace local experiment already exists.",
+        help="Replace local experiment if it already exists.",
     )
     experiments_pull_parser.add_argument(
         "--no-cache",

dvc/commands/experiments/save.py

@@ -0,0 +1,79 @@
+import argparse
+import logging
+
+from dvc.cli.command import CmdBase
+from dvc.cli.utils import append_doc_link
+from dvc.exceptions import DvcException
+from dvc.ui import ui
+
+logger = logging.getLogger(__name__)
+
+
+class CmdExperimentsSave(CmdBase):
+    def run(self):
+
+        try:
+            ref = self.repo.experiments.save(
+                name=self.args.name,
+                force=self.args.force,
+                include_untracked=self.args.include_untracked,
+            )
+        except DvcException:
+            logger.exception("failed to save experiment")
+            return 1
+
+        if self.args.json:
+            ui.write_json({"ref": ref})
+        else:
+            name = self.repo.experiments.get_exact_name([ref])[ref]
+            ui.write(f"Experiment has been saved as: {name}")
+            ui.write(
+                "\nTo promote an experiment to a Git branch run:\n\n"
+                "\tdvc exp branch <exp> <branch>\n"
+            )
+
+        return 0
+
+
+def add_parser(experiments_subparsers, parent_parser):
+    EXPERIMENTS_SAVE_HELP = "Save current workspace as a dvc experiment."
+    save_parser = experiments_subparsers.add_parser(
+        "save",
+        parents=[parent_parser],
+        description=append_doc_link(EXPERIMENTS_SAVE_HELP, "exp/save"),
+        help=EXPERIMENTS_SAVE_HELP,
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    save_parser.add_argument(
+        "-f",
+        "--force",
+        action="store_true",
+        default=False,
+        help="Replace experiment if it already exists.",
+    )
+    save_parser.add_argument(
+        "--json",
+        "--show-json",
+        action="store_true",
+        default=False,
+        help="Show output in JSON format.",
+    )
+    save_parser.add_argument(
+        "-n",
+        "--name",
+        default=None,
+        help=(
+            "Human-readable experiment name. If not specified, a name will "
+            "be auto-generated."
+        ),
+        metavar="<name>",
+    )
+    save_parser.add_argument(
+        "-I",
+        "--include-untracked",
+        action="append",
+        default=[],
+        help="List of untracked paths to include in the experiment.",
+        metavar="<path>",
+    )
+    save_parser.set_defaults(func=CmdExperimentsSave)

dvc/commands/experiments/show.py

@@ -4,7 +4,7 @@ import os
 import re
 from collections import Counter, OrderedDict, defaultdict
 from datetime import date, datetime
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any, Dict, Set
 
 from funcy import lmap
 
@@ -37,9 +37,9 @@ def _update_names(names, items):
 
 
 def _collect_names(all_experiments, **kwargs):
-    metric_names = defaultdict(dict)
-    param_names = defaultdict(dict)
-    deps_names = set()
+    metric_names: Dict[str, Dict[str, Any]] = defaultdict(dict)
+    param_names: Dict[str, Dict[str, Any]] = defaultdict(dict)
+    deps_names: Set[str] = set()
 
     for _, experiments in all_experiments.items():
         for exp_data in experiments.values():
@@ -62,7 +62,7 @@ experiment_types = {
 }
 
 
-def _collect_rows(
+def _collect_rows(  # noqa: C901
     base_rev,
     experiments,
     all_headers,
@@ -314,7 +314,7 @@ def baseline_styler(typ):
     return {"style": "bold"} if typ == "baseline" else {}
 
 
-def show_experiments(
+def show_experiments(  # noqa: C901
     all_experiments,
     keep=None,
     drop=None,
@@ -328,7 +328,7 @@ def show_experiments(
 
     metric_names, param_names, deps_names = _collect_names(all_experiments)
 
-    headers = [
+    headers_list = [
         "Experiment",
         "rev",
         "typ",
@@ -340,13 +340,13 @@ def show_experiments(
 
     names = {**metric_names, **param_names}
     counter = Counter(flatten_list([list(a.keys()) for a in names.values()]))
-    counter.update(headers)
+    counter.update(headers_list)
     metric_headers = _normalize_headers(metric_names, counter)
     param_headers = _normalize_headers(param_names, counter)
 
     td = experiments_table(
         all_experiments,
-        headers,
+        headers_list,
         metric_headers,
         metric_names,
         param_headers,
@@ -359,7 +359,7 @@ def show_experiments(
         kwargs.get("iso"),
     )
     if keep:
-        for col in td.keys():
+        for col in td.keys():  # noqa: SIM118
             if re.match(keep, col):
                 td.protect(col)
 
@@ -410,13 +410,15 @@ def show_experiments(
 
     cols_to_drop = set()
     if drop is not None:
-        cols_to_drop = {col for col in td.keys() if re.match(drop, col)}
+        cols_to_drop = {
+            col for col in td.keys() if re.match(drop, col)  # noqa: SIM118
+        }
     if pcp:
         cols_to_drop.add("Created")
     td.drop(*cols_to_drop)
 
     if pcp:
-        subset = {x for x in td.keys() if x != "Experiment"}
+        subset = {x for x in td.keys() if x != "Experiment"}  # noqa: SIM118
         td.dropna(
             "rows",
             how="all",
@@ -494,7 +496,7 @@ class CmdExperimentsShow(CmdBase):
                 else DEFAULT_PRECISION
             )
             fill_value = "" if self.args.csv else FILL_VALUE
-            iso = True if self.args.csv else False
+            iso = self.args.csv
 
             show_experiments(
                 all_experiments,

dvc/commands/gc.py

@@ -10,7 +10,7 @@ logger = logging.getLogger(__name__)
 
 
 class CmdGC(CmdBase):
-    def run(self):
+    def run(self):  # noqa: C901
         from dvc.repo.gc import _validate_args
 
         _validate_args(

dvc/commands/get.py

@@ -3,9 +3,9 @@ import logging
 
 from dvc.cli import completion
 from dvc.cli.command import CmdBaseNoRepo
+from dvc.cli.utils import append_doc_link
 from dvc.exceptions import DvcException
-
-from ..cli.utils import append_doc_link
+from dvc.scm import CloneError
 
 logger = logging.getLogger(__name__)
 
@@ -44,6 +44,9 @@ class CmdGet(CmdBaseNoRepo):
                 jobs=self.args.jobs,
             )
             return 0
+        except CloneError:
+            logger.exception(f"failed to get '{self.args.path}'")
+            return 1
         except DvcException:
             logger.exception(
                 "failed to get '{}' from '{}'".format(

dvc/commands/get_url.py

@@ -3,10 +3,9 @@ import logging
 
 from dvc.cli import completion
 from dvc.cli.command import CmdBaseNoRepo
+from dvc.cli.utils import append_doc_link
 from dvc.exceptions import DvcException
 
-from ..cli.utils import append_doc_link
-
 logger = logging.getLogger(__name__)
 
 

dvc/commands/git_hook.py

@@ -64,15 +64,15 @@ class CmdPrePush(CmdHookBase):
 
 class CmdMergeDriver(CmdHookBase):
     def _run(self):
-        from dvc.dvcfile import Dvcfile
+        from dvc.dvcfile import load_file
         from dvc.repo import Repo
 
         dvc = Repo()
 
         try:
-            ancestor = Dvcfile(dvc, self.args.ancestor, verify=False)
-            our = Dvcfile(dvc, self.args.our, verify=False)
-            their = Dvcfile(dvc, self.args.their, verify=False)
+            ancestor = load_file(dvc, self.args.ancestor, verify=False)
+            our = load_file(dvc, self.args.our, verify=False)
+            their = load_file(dvc, self.args.their, verify=False)
 
             our.merge(ancestor, their, allowed=["add", "remove", "change"])
 

dvc/commands/imp.py

@@ -5,6 +5,7 @@ from dvc.cli import completion
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
 from dvc.exceptions import DvcException
+from dvc.scm import CloneError
 
 logger = logging.getLogger(__name__)
 
@@ -26,6 +27,9 @@ class CmdImport(CmdBase):
                 meta=self.args.meta,
                 jobs=self.args.jobs,
             )
+        except CloneError:
+            logger.exception(f"failed to import '{self.args.path}'")
+            return 1
         except DvcException:
             logger.exception(
                 "failed to import '{}' from '{}'.".format(

dvc/commands/imp_url.py

@@ -73,6 +73,18 @@ def add_parser(subparsers, parent_parser):
         help="Specify name of the .dvc file this command will generate.",
         metavar="<filename>",
     ).complete = completion.DIR
+    import_parser.add_argument(
+        "--to-remote",
+        action="store_true",
+        default=False,
+        help="Download it directly to the remote",
+    )
+    import_parser.add_argument(
+        "-r",
+        "--remote",
+        help="Remote storage to download to",
+        metavar="<name>",
+    )
     no_download_exec_group = import_parser.add_mutually_exclusive_group()
     no_download_exec_group.add_argument(
         "--no-exec",
@@ -87,18 +99,6 @@ def add_parser(subparsers, parent_parser):
         help="Create .dvc file including target data hash value(s)"
         " but do not actually download the file(s).",
     )
-    import_parser.add_argument(
-        "--to-remote",
-        action="store_true",
-        default=False,
-        help="Download it directly to the remote",
-    )
-    import_parser.add_argument(
-        "-r",
-        "--remote",
-        help="Remote storage to download to",
-        metavar="<name>",
-    )
     import_parser.add_argument(
         "-j",
         "--jobs",

dvc/commands/machine.py

@@ -1,4 +1,5 @@
 import argparse
+from typing import Dict, List
 
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link, fix_subparsers
@@ -6,7 +7,6 @@ from dvc.commands.config import CmdConfig
 from dvc.compare import TabularData
 from dvc.config import ConfigError
 from dvc.exceptions import DvcException
-from dvc.types import Dict, List
 from dvc.ui import ui
 from dvc.utils import format_link
 
@@ -177,6 +177,7 @@ class CmdMachineRename(CmdBase):
             self._check_exists(conf)
             conf["machine"][self.args.new] = conf["machine"][self.args.name]
             try:
+                assert self.repo.machine
                 self.repo.machine.rename(self.args.name, self.args.new)
             except DvcException as error:
                 del conf["machine"][self.args.new]
@@ -199,7 +200,7 @@ class CmdMachineDefault(CmdMachineConfig):
         if self.args.name is None and not self.args.unset:
             conf = self.config.read(self.args.level)
             try:
-                print(conf["core"]["machine"])
+                ui.write(conf["core"]["machine"])
             except KeyError:
                 ui.write("No default machine set")
                 return 1
@@ -324,7 +325,7 @@ def add_parser(subparsers, parent_parser):
 
     machine_subparsers = machine_parser.add_subparsers(
         dest="cmd",
-        help="Use `dvc machine CMD --help` for " "command-specific help.",
+        help="Use `dvc machine CMD --help` for command-specific help.",
     )
 
     fix_subparsers(machine_subparsers)

dvc/commands/plots.py

@@ -79,10 +79,8 @@ class CmdPlots(CmdBase):
         return {k: v for k, v in props.items() if v is not None}
 
     def _config_files(self):
-        config_files = None
         if self.args.from_config:
-            config_files = {self.args.from_config}
-        return config_files
+            return {self.args.from_config}
 
     def _html_template_path(self):
         html_template_path = self.args.html_template
@@ -91,12 +89,13 @@ class CmdPlots(CmdBase):
                 "html_template", None
             )
             if html_template_path and not os.path.isabs(html_template_path):
+                assert self.repo.dvc_dir
                 html_template_path = os.path.join(
                     self.repo.dvc_dir, html_template_path
                 )
         return html_template_path
 
-    def run(self):
+    def run(self) -> int:  # noqa: C901
         from pathlib import Path
 
         from dvc.render.match import match_defs_renderers
@@ -223,9 +222,8 @@ class CmdPlotsTemplates(CmdBase):
                         return 0
                 raise InvalidArgumentError(f"Unexpected template: {target}.")
 
-            else:
-                for template in TEMPLATES:
-                    ui.write(template.DEFAULT_NAME)
+            for template in TEMPLATES:
+                ui.write(template.DEFAULT_NAME)
 
             return 0
         except DvcException:

dvc/commands/queue/__init__.py

@@ -26,7 +26,7 @@ def add_parser(subparsers, parent_parser):
 
     queue_subparsers = queue_parser.add_subparsers(
         dest="cmd",
-        help="Use `dvc queue CMD --help` to display " "command-specific help.",
+        help="Use `dvc queue CMD --help` to display command-specific help.",
     )
 
     fix_subparsers(queue_subparsers)

dvc/commands/queue/kill.py

@@ -11,13 +11,18 @@ class CmdQueueKill(CmdBase):
     """Kill exp task in queue."""
 
     def run(self):
-        self.repo.experiments.celery_queue.kill(revs=self.args.task)
+        self.repo.experiments.celery_queue.kill(
+            revs=self.args.task, force=self.args.force
+        )
 
         return 0
 
 
 def add_parser(queue_subparsers, parent_parser):
-    QUEUE_KILL_HELP = "Kill actively running experiment queue tasks."
+    QUEUE_KILL_HELP = (
+        "Gracefully interrupt running experiment queue tasks "
+        "(equivalent to Ctrl-C)"
+    )
     queue_kill_parser = queue_subparsers.add_parser(
         "kill",
         parents=[parent_parser],
@@ -25,6 +30,13 @@ def add_parser(queue_subparsers, parent_parser):
         help=QUEUE_KILL_HELP,
         formatter_class=argparse.RawDescriptionHelpFormatter,
     )
+    queue_kill_parser.add_argument(
+        "-f",
+        "--force",
+        action="store_true",
+        default=False,
+        help="Forcefully and immediately kill running experiment queue tasks",
+    )
     queue_kill_parser.add_argument(
         "task",
         nargs="*",

dvc/commands/queue/status.py

@@ -1,33 +1,34 @@
 import argparse
 import logging
-from typing import List, Mapping, Optional
+from typing import Dict, List, Optional
 
 from dvc.cli.command import CmdBase
 from dvc.cli.utils import append_doc_link
+from dvc.commands.experiments.show import format_time
 from dvc.compare import TabularData
 from dvc.ui import ui
 
-from ..experiments.show import format_time
-
 logger = logging.getLogger(__name__)
 
 
 class CmdQueueStatus(CmdBase):
     """Show queue task and worker status."""
 
-    def run(self):
+    def run(self) -> int:
         result: List[
-            Mapping[str, Optional[str]]
+            Dict[str, Optional[str]]
         ] = self.repo.experiments.celery_queue.status()
         if result:
             all_headers = ["Task", "Name", "Created", "Status"]
             td = TabularData(all_headers)
             for exp in result:
                 created = format_time(exp.get("timestamp"))
+                assert exp["rev"]
+                assert exp["status"]
                 td.append(
                     [
                         exp["rev"][:7],
-                        exp.get("name", ""),
+                        exp.get("name") or "",
                         created,
                         exp["status"],
                     ]

dvc/commands/remote.py

@@ -164,7 +164,7 @@ def add_parser(subparsers, parent_parser):
 
     remote_subparsers = remote_parser.add_subparsers(
         dest="cmd",
-        help="Use `dvc remote CMD --help` for " "command-specific help.",
+        help="Use `dvc remote CMD --help` for command-specific help.",
     )
 
     fix_subparsers(remote_subparsers)

dvc/commands/stage.py

@@ -67,7 +67,7 @@ def prepare_stages_data(
 class CmdStageList(CmdBase):
     def _get_stages(self) -> Iterable["Stage"]:
         if self.args.all:
-            stages: List["Stage"] = self.repo.index.stages  # type: ignore
+            stages: List["Stage"] = self.repo.index.stages
             logger.trace(  # type: ignore[attr-defined]
                 "%d no. of stages found", len(stages)
             )
@@ -95,7 +95,7 @@ class CmdStageList(CmdBase):
 
         stages = self._get_stages()
         data = prepare_stages_data(stages, description=not self.args.name_only)
-        ui.table(data.items())
+        ui.table(list(data.items()))
 
         return 0
 

dvc/commands/update.py

@@ -31,7 +31,8 @@ class CmdUpdate(CmdBase):
 def add_parser(subparsers, parent_parser):
     UPDATE_HELP = (
         "Update data artifact imported (via dvc import or dvc import-url) "
-        "from an external DVC repository or URL."
+        "from an external DVC repository or URL, or update data from a "
+        "DVC worktree remote."
     )
     update_parser = subparsers.add_parser(
         "update",

dvc/compare.py

@@ -58,7 +58,7 @@ class TabularData(MutableSequence[Sequence["CellT"]]):
         return self._columns[name]
 
     def items(self) -> ItemsView[str, Column]:
-        projection = {k: self.column(k) for k in self.keys()}
+        projection = {k: self.column(k) for k in self.keys()}  # noqa: SIM118
         return projection.items()
 
     def keys(self) -> List[str]:
@@ -167,7 +167,7 @@ class TabularData(MutableSequence[Sequence["CellT"]]):
         return buff.getvalue()
 
     def to_parallel_coordinates(
-        self, output_path: "StrPath", color_by: str = None
+        self, output_path: "StrPath", color_by: Optional[str] = None
     ) -> "StrPath":
         from dvc_render.html import render_html
         from dvc_render.plotly import ParallelCoordinatesRenderer
@@ -195,7 +195,8 @@ class TabularData(MutableSequence[Sequence["CellT"]]):
                 self.add_column(key)
 
         row: List["CellT"] = [
-            with_value(d.get(key), self._fill_value) for key in self.keys()
+            with_value(d.get(key), self._fill_value)
+            for key in self.keys()  # noqa: SIM118
         ]
         self.append(row)
 
@@ -208,14 +209,14 @@ class TabularData(MutableSequence[Sequence["CellT"]]):
             ui.table(self, headers=self.keys(), **kwargs)
 
     def as_dict(
-        self, cols: Iterable[str] = None
+        self, cols: Optional[Iterable[str]] = None
     ) -> Iterable[Dict[str, "CellT"]]:
         keys = self.keys() if cols is None else set(cols)
         return [
             {k: self._columns[k][i] for k in keys} for i in range(len(self))
         ]
 
-    def dropna(
+    def dropna(  # noqa: C901
         self,
         axis: str = "rows",
         how="any",
@@ -228,7 +229,7 @@ class TabularData(MutableSequence[Sequence["CellT"]]):
             )
         if how not in ["any", "all"]:
             raise ValueError(
-                f"Invalid 'how' value {how}." "Choose one of ['any', 'all']"
+                f"Invalid 'how' value {how}. Choose one of ['any', 'all']"
             )
 
         match_line: Set = set()
@@ -244,8 +245,7 @@ class TabularData(MutableSequence[Sequence["CellT"]]):
                     if axis == "rows":
                         match_line.add(n_row)
                         break
-                    else:
-                        match_line.add(self.keys()[n_col])
+                    match_line.add(self.keys()[n_col])
 
         to_drop = match_line
         if how == "all":
@@ -256,7 +256,7 @@ class TabularData(MutableSequence[Sequence["CellT"]]):
             to_drop -= match_line
 
         if axis == "rows":
-            for name in self.keys():
+            for name in self.keys():  # noqa: SIM118
                 self._columns[name] = Column(
                     [
                         x
@@ -267,7 +267,7 @@ class TabularData(MutableSequence[Sequence["CellT"]]):
         else:
             self.drop(*to_drop)
 
-    def drop_duplicates(
+    def drop_duplicates(  # noqa: C901
         self,
         axis: str = "rows",
         subset: Optional[Iterable[str]] = None,
@@ -309,7 +309,7 @@ class TabularData(MutableSequence[Sequence["CellT"]]):
                 else:
                     unique_rows.append(tuple_row)
 
-            for name in self.keys():
+            for name in self.keys():  # noqa: SIM118
                 self._columns[name] = Column(
                     [
                         x
@@ -324,12 +324,13 @@ def _normalize_float(val: float, precision: int):
 
 
 def _format_field(
-    val: Any, precision: int = None, round_digits: bool = False
+    val: Any, precision: Optional[int] = None, round_digits: bool = False
 ) -> str:
     def _format(_val):
         if isinstance(_val, float) and precision:
-            func = round if round_digits else _normalize_float
-            return func(_val, precision)
+            if round_digits:
+                return round(val, precision)
+            return _normalize_float(val, precision)
         if isinstance(_val, abc.Mapping):
             return {k: _format(v) for k, v in _val.items()}
         if isinstance(_val, list):
@@ -345,11 +346,11 @@ def diff_table(
     old: bool = True,
     no_path: bool = False,
     show_changes: bool = True,
-    precision: int = None,
+    precision: Optional[int] = None,
     round_digits: bool = False,
-    on_empty_diff: str = None,
-    a_rev: str = None,
-    b_rev: str = None,
+    on_empty_diff: Optional[str] = None,
+    a_rev: Optional[str] = None,
+    b_rev: Optional[str] = None,
 ) -> TabularData:
     a_rev = a_rev or "HEAD"
     b_rev = b_rev or "workspace"
@@ -393,12 +394,12 @@ def show_diff(
     old: bool = True,
     no_path: bool = False,
     show_changes: bool = True,
-    precision: int = None,
+    precision: Optional[int] = None,
     round_digits: bool = False,
-    on_empty_diff: str = None,
+    on_empty_diff: Optional[str] = None,
     markdown: bool = False,
-    a_rev: str = None,
-    b_rev: str = None,
+    a_rev: Optional[str] = None,
+    b_rev: Optional[str] = None,
 ) -> None:
     td = diff_table(
         diff,
@@ -420,7 +421,7 @@ def metrics_table(
     all_branches: bool = False,
     all_tags: bool = False,
     all_commits: bool = False,
-    precision: int = None,
+    precision: Optional[int] = None,
     round_digits: bool = False,
 ):
     from dvc.utils.diff import format_dict
@@ -460,7 +461,7 @@ def show_metrics(
     all_branches: bool = False,
     all_tags: bool = False,
     all_commits: bool = False,
-    precision: int = None,
+    precision: Optional[int] = None,
     round_digits: bool = False,
 ) -> None:
     td = metrics_table(

dvc/config.py

@@ -4,10 +4,17 @@ import os
 import re
 from contextlib import contextmanager
 from functools import partial
+from typing import TYPE_CHECKING, Dict, Optional
 
-from funcy import cached_property, compact, memoize, re_find
+from funcy import compact, memoize, re_find
 
 from dvc.exceptions import DvcException, NotDvcRepoError
+from dvc.types import DictStrAny, StrPath
+
+from .utils.objects import cached_property
+
+if TYPE_CHECKING:
+    from dvc.fs import FileSystem
 
 logger = logging.getLogger(__name__)
 
@@ -84,7 +91,11 @@ class Config(dict):
     CONFIG_LOCAL = "config.local"
 
     def __init__(
-        self, dvc_dir=None, validate=True, fs=None, config=None
+        self,
+        dvc_dir: Optional[StrPath] = None,
+        validate: bool = True,
+        fs: Optional["FileSystem"] = None,
+        config: Optional[DictStrAny] = None,
     ):  # pylint: disable=super-init-not-called
         from dvc.fs import LocalFileSystem
 
@@ -116,7 +127,7 @@ class Config(dict):
             return site_config_dir(cls.APPNAME, cls.APPAUTHOR)
 
     @cached_property
-    def files(self):
+    def files(self) -> Dict[str, str]:
         files = {
             level: os.path.join(self.get_dir(level), self.CONFIG)
             for level in ("system", "global")
@@ -139,10 +150,10 @@ class Config(dict):
             dvc.config.Config: config object.
         """
         config_file = os.path.join(dvc_dir, Config.CONFIG)
-        open(config_file, "w+", encoding="utf-8").close()
-        return Config(dvc_dir)
+        with open(config_file, "w+", encoding="utf-8"):
+            return Config(dvc_dir)
 
-    def load(self, validate=True, config=None):
+    def load(self, validate: bool = True, config: Optional[DictStrAny] = None):
         """Loads config from all the config files.
 
         Raises:
@@ -169,7 +180,7 @@ class Config(dict):
         return self.fs if level == "repo" else self.wfs
 
     def _load_config(self, level):
-        from configobj import ConfigObj
+        from configobj import ConfigObj, ConfigObjError
 
         filename = self.files[level]
         fs = self._get_fs(level)
@@ -180,6 +191,8 @@ class Config(dict):
                     conf_obj = ConfigObj(fobj)
                 except UnicodeDecodeError as exc:
                     raise ConfigError(str(exc)) from exc
+                except ConfigObjError as exc:
+                    raise ConfigError(str(exc)) from exc
         else:
             conf_obj = ConfigObj()
         return _parse_named(_lower_keys(conf_obj.dict()))
@@ -278,7 +291,7 @@ class Config(dict):
         return Schema(dirs_schema, extra=ALLOW_EXTRA)(conf)
 
     def load_config_to_level(self, level=None):
-        merged_conf = {}
+        merged_conf: Dict = {}
         for merge_level in self.LEVELS:
             if merge_level == level:
                 break
@@ -324,7 +337,7 @@ class Config(dict):
 
 
 def _parse_named(conf):
-    result = {"remote": {}, "machine": {}}
+    result: Dict[str, Dict] = {"remote": {}, "machine": {}}
 
     for section, val in conf.items():
         match = re_find(r'^\s*(remote|machine)\s*"(.*)"\s*$', section)

dvc/config_schema.py

@@ -97,6 +97,8 @@ HTTP_COMMON = {
     "ask_password": Bool,
     "ssl_verify": Any(Bool, str),
     "method": str,
+    "connect_timeout": All(Coerce(float), Range(0, min_included=True)),
+    "read_timeout": All(Coerce(float), Range(0, min_included=True)),
     Optional("verify", default=False): Bool,
 }
 WEBDAV_COMMON = {

dvc/daemon.py

@@ -25,21 +25,24 @@ def _popen(cmd, **kwargs):
 
 
 def _spawn_windows(cmd, env):
-    from subprocess import (
-        CREATE_NEW_PROCESS_GROUP,
-        CREATE_NO_WINDOW,
-        STARTF_USESHOWWINDOW,
-        STARTUPINFO,
-    )
+    if sys.platform == "win32":
+        from subprocess import (
+            CREATE_NEW_PROCESS_GROUP,
+            CREATE_NO_WINDOW,
+            STARTF_USESHOWWINDOW,
+            STARTUPINFO,
+        )
 
-    # https://stackoverflow.com/a/7006424
-    # https://bugs.python.org/issue41619
-    creationflags = CREATE_NEW_PROCESS_GROUP | CREATE_NO_WINDOW
+        # https://stackoverflow.com/a/7006424
+        # https://bugs.python.org/issue41619
+        creationflags = CREATE_NEW_PROCESS_GROUP | CREATE_NO_WINDOW
 
-    startupinfo = STARTUPINFO()
-    startupinfo.dwFlags |= STARTF_USESHOWWINDOW
+        startupinfo = STARTUPINFO()
+        startupinfo.dwFlags |= STARTF_USESHOWWINDOW
 
-    _popen(cmd, env=env, creationflags=creationflags, startupinfo=startupinfo)
+        _popen(
+            cmd, env=env, creationflags=creationflags, startupinfo=startupinfo
+        )
 
 
 def _spawn_posix(cmd, env):
@@ -49,17 +52,19 @@ def _spawn_posix(cmd, env):
     # with PyInstaller has trouble with SystemExit exception and throws
     # errors such as "[26338] Failed to execute script __main__"
     try:
-        pid = os.fork()  # pylint: disable=no-member
+        # pylint: disable-next=no-member
+        pid = os.fork()  # type: ignore[attr-defined]
         if pid > 0:
             return
     except OSError:
         logger.exception("failed at first fork")
         os._exit(1)  # pylint: disable=protected-access
 
-    os.setsid()  # pylint: disable=no-member
+    os.setsid()  # type: ignore[attr-defined]  # pylint: disable=no-member
 
     try:
-        pid = os.fork()  # pylint: disable=no-member
+        # pylint: disable-next=no-member
+        pid = os.fork()  # type: ignore[attr-defined]
         if pid > 0:
             os._exit(0)  # pylint: disable=protected-access
     except OSError:

dvc/dagascii.py

@@ -59,8 +59,7 @@ class AsciiCanvas:
     def draw(self):
         """Draws ASCII canvas on the screen."""
         lines = map("".join, self.canvas)
-        joined_lines = os.linesep.join(lines)
-        return joined_lines
+        return os.linesep.join(lines)
 
     def point(self, x, y, char):
         """Create a point on ASCII canvas.
@@ -81,7 +80,7 @@ class AsciiCanvas:
 
         self.canvas[y][x] = char
 
-    def line(self, x0, y0, x1, y1, char):
+    def line(self, x0, y0, x1, y1, char):  # noqa: C901
         """Create a line on ASCII canvas.
 
         Args:

dvc/data_cloud.py

@@ -3,8 +3,8 @@
 import logging
 from typing import TYPE_CHECKING, Iterable, Optional
 
-from funcy import cached_property
-
+from dvc.config import NoRemoteError, RemoteConfigError
+from dvc.utils.objects import cached_property
 from dvc_data.hashfile.db import get_index
 
 if TYPE_CHECKING:
@@ -23,9 +23,17 @@ class Remote:
 
         self.worktree = config.pop("worktree", False)
         self.config = config
+        if self.worktree:
+            version_aware = self.config.get("version_aware")
+            if version_aware is False:
+                raise RemoteConfigError(
+                    "worktree remotes require version_aware cloud"
+                )
+            if version_aware is None:
+                self.fs.version_aware = True
 
     @cached_property
-    def odb(self):
+    def odb(self) -> "HashFileDB":
         from dvc_data.hashfile.db import get_odb
 
         path = self.path
@@ -54,8 +62,6 @@ class DataCloud:
         name: Optional[str] = None,
         command: str = "<command>",
     ) -> "Remote":
-        from dvc.config import NoRemoteError
-
         if not name:
             name = self.repo.config["core"].get("remote")
 
@@ -88,6 +94,10 @@ class DataCloud:
         command: str = "<command>",
     ) -> "HashFileDB":
         remote = self.get_remote(name=name, command=command)
+        if remote.fs.version_aware or remote.worktree:
+            raise NoRemoteError(
+                f"'{command}' is unsupported for cloud versioned remotes"
+            )
         return remote.odb
 
     def _log_missing(self, status: "CompareStatusResult"):
@@ -173,7 +183,6 @@ class DataCloud:
         jobs: Optional[int] = None,
         remote: Optional[str] = None,
         odb: Optional["HashFileDB"] = None,
-        log_missing: bool = True,
     ):
         """Check status of data items in a cloud-agnostic way.
 
@@ -184,8 +193,6 @@ class DataCloud:
                 cache to. By default remote from core.remote config option
                 is used.
             odb: optional ODB to check status from. Overrides remote.
-            log_missing: log warning messages if file doesn't exist
-                neither in cache, neither in cloud.
         """
         from dvc_data.hashfile.status import compare_status
 
@@ -196,7 +203,6 @@ class DataCloud:
             odb,
             objs,
             jobs=jobs,
-            log_missing=log_missing,
             dest_index=get_index(odb),
             cache_odb=self.repo.odb.local,
         )

dvc/dependency/__init__.py

@@ -1,5 +1,5 @@
 from collections import defaultdict
-from typing import Any, Mapping
+from typing import Any, Mapping, Set
 
 from dvc.output import ARTIFACT_SCHEMA, DIR_FILES_SCHEMA, Output
 
@@ -59,7 +59,7 @@ def _merge_params(s_list):
     default_file = ParamsDependency.DEFAULT_PARAMS_FILE
 
     # figure out completely tracked params file, and ignore specific keys
-    wholly_tracked = set()
+    wholly_tracked: Set[str] = set()
     for key in s_list:
         if not isinstance(key, dict):
             continue

dvc/dependency/base.py

@@ -25,11 +25,9 @@ class DependencyIsStageFileError(DvcException):
 class Dependency(Output):
     IS_DEPENDENCY = True
 
-    DoesNotExistError = DependencyDoesNotExistError  # type: Type[DvcException]
-    IsNotFileOrDirError = (
-        DependencyIsNotFileOrDirError
-    )  # type: Type[DvcException]
-    IsStageFileError = DependencyIsStageFileError  # type: Type[DvcException]
+    DoesNotExistError: Type[DvcException] = DependencyDoesNotExistError
+    IsNotFileOrDirError: Type[DvcException] = DependencyIsNotFileOrDirError
+    IsStageFileError: Type[DvcException] = DependencyIsStageFileError
 
     def workspace_status(self) -> Dict[str, str]:
         if self.fs.version_aware:

dvc/dependency/param.py

@@ -2,7 +2,7 @@ import logging
 import os
 import typing
 from collections import defaultdict
-from typing import Dict
+from typing import Dict, cast
 
 import dpath.util
 from voluptuous import Any
@@ -41,7 +41,10 @@ class ParamsDependency(Dependency):
         self.params = list(params) if params else []
         hash_info = HashInfo()
         if isinstance(params, dict):
-            hash_info = HashInfo(self.PARAM_PARAMS, params)
+            hash_info = HashInfo(
+                self.PARAM_PARAMS,
+                params,  # type: ignore[arg-type]
+            )
         repo = repo or stage.repo
         path = path or os.path.join(repo.root_dir, self.DEFAULT_PARAMS_FILE)
         super().__init__(stage, path, repo=repo)
@@ -64,7 +67,10 @@ class ParamsDependency(Dependency):
         for param in self.params:
             if param in values:
                 info[param] = values[param]
-        self.hash_info = HashInfo(self.PARAM_PARAMS, info)
+        self.hash_info = HashInfo(
+            self.PARAM_PARAMS,
+            info,  # type: ignore[arg-type]
+        )
 
     def read_params(
         self, flatten: bool = True, **kwargs: typing.Any
@@ -104,8 +110,8 @@ class ParamsDependency(Dependency):
 
         from funcy import ldistinct
 
-        status = defaultdict(dict)
-        info = self.hash_info.value if self.hash_info else {}
+        status: Dict[str, Any] = defaultdict(dict)
+        info = cast(dict, self.hash_info.value) if self.hash_info else {}
         actual = self.read_params()
 
         # NOTE: we want to preserve the order of params as specified in the
@@ -119,9 +125,14 @@ class ParamsDependency(Dependency):
             elif param not in info:
                 st = "new"
             elif actual[param] != info[param]:
-                st = "modified"
+                if (
+                    isinstance(actual[param], tuple)
+                    and list(actual[param]) == info[param]
+                ):
+                    continue
+                else:
+                    st = "modified"
             else:
-                assert actual[param] == info[param]
                 continue
 
             status[str(self)][param] = st
@@ -159,13 +170,13 @@ class ParamsDependency(Dependency):
                 )
             )
 
-        return HashInfo(self.PARAM_PARAMS, info)
+        return HashInfo(self.PARAM_PARAMS, info)  # type: ignore[arg-type]
 
     def save(self):
         if not self.exists:
             raise self.DoesNotExistError(self)
 
-        if not self.isfile and not self.isdir:
+        if not self.isfile() and not self.isdir():
             raise self.IsNotFileOrDirError(self)
 
         self.ignore()

dvc/dependency/repo.py

@@ -1,7 +1,15 @@
 import os
 from collections import defaultdict
 from copy import copy
-from typing import TYPE_CHECKING, Dict, Optional, Set, Tuple
+from typing import (
+    TYPE_CHECKING,
+    ContextManager,
+    Dict,
+    Optional,
+    Set,
+    Tuple,
+    Union,
+)
 
 from voluptuous import Required
 
@@ -11,8 +19,11 @@ from dvc_data.hashfile.meta import Meta
 from .base import Dependency
 
 if TYPE_CHECKING:
-    from dvc_data.hashfile.file import HashFile
+    from dvc.output import Output
+    from dvc.repo import Repo
+    from dvc.stage import Stage
     from dvc_data.hashfile.hash_info import HashInfo
+    from dvc_data.hashfile.obj import HashFile
     from dvc_objects.db import ObjectDB
 
 
@@ -30,7 +41,9 @@ class RepoDependency(Dependency):
         }
     }
 
-    def __init__(self, def_repo, stage, *args, **kwargs):
+    def __init__(
+        self, def_repo: Dict[str, str], stage: "Stage", *args, **kwargs
+    ):
         self.def_repo = def_repo
         self._objs: Dict[str, "HashFile"] = {}
         self._meta: Dict[str, Meta] = {}
@@ -61,10 +74,10 @@ class RepoDependency(Dependency):
     def save(self):
         pass
 
-    def dumpd(self, **kwargs):
+    def dumpd(self, **kwargs) -> Dict[str, Union[str, Dict[str, str]]]:
         return {self.PARAM_PATH: self.def_path, self.PARAM_REPO: self.def_repo}
 
-    def download(self, to, jobs=None):
+    def download(self, to: "Output", jobs: Optional[int] = None):
         from dvc_data.hashfile.checkout import checkout
 
         for odb, objs in self.get_used_objs().items():
@@ -81,13 +94,13 @@ class RepoDependency(Dependency):
             prompt=confirm,
         )
 
-    def update(self, rev=None):
+    def update(self, rev: Optional[str] = None):
         if rev:
             self.def_repo[self.PARAM_REV] = rev
         with self._make_repo(locked=False) as repo:
             self.def_repo[self.PARAM_REV_LOCK] = repo.get_rev()
 
-    def changed_checksum(self):
+    def changed_checksum(self) -> bool:
         # From current repo point of view what describes RepoDependency is its
         # origin project url and rev_lock, and it makes RepoDependency
         # immutable, hence its impossible for checksum to change.
@@ -100,7 +113,7 @@ class RepoDependency(Dependency):
         return used
 
     def _get_used_and_obj(
-        self, obj_only=False, **kwargs
+        self, obj_only: bool = False, **kwargs
     ) -> Tuple[Dict[Optional["ObjectDB"], Set["HashInfo"]], Meta, "HashFile"]:
         from dvc.config import NoRemoteError
         from dvc.exceptions import NoOutputOrStageError, PathMissingError
@@ -154,7 +167,9 @@ class RepoDependency(Dependency):
                 used_obj_ids[object_store].update(oid for _, _, oid in obj)
             return used_obj_ids, meta, obj
 
-    def _check_circular_import(self, odb, obj_ids):
+    def _check_circular_import(
+        self, odb: "ObjectDB", obj_ids: Set["HashInfo"]
+    ) -> None:
         from dvc.exceptions import CircularImportError
         from dvc.fs.dvc import DVCFileSystem
         from dvc_data.hashfile.db.reference import ReferenceHashFileDB
@@ -169,6 +184,7 @@ class RepoDependency(Dependency):
                     tree = Tree.load(odb, hash_info)
                     yield from (odb.get(hi.value) for _, _, hi in tree)
                 else:
+                    assert hash_info.value
                     yield odb.get(hash_info.value)
 
         checked_urls = set()
@@ -209,14 +225,16 @@ class RepoDependency(Dependency):
         )
         return meta
 
-    def _make_repo(self, locked=True, **kwargs):
+    def _make_repo(
+        self, locked: bool = True, **kwargs
+    ) -> "ContextManager[Repo]":
         from dvc.external_repo import external_repo
 
         d = self.def_repo
         rev = self._get_rev(locked=locked)
         return external_repo(d[self.PARAM_URL], rev=rev, **kwargs)
 
-    def _get_rev(self, locked=True):
+    def _get_rev(self, locked: bool = True):
         d = self.def_repo
         return (d.get(self.PARAM_REV_LOCK) if locked else None) or d.get(
             self.PARAM_REV

dvc/dvcfile.py

@@ -1,7 +1,17 @@
 import contextlib
 import logging
 import os
-from typing import TYPE_CHECKING, Any, Callable, Tuple, TypeVar, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Tuple,
+    TypeVar,
+    Union,
+)
 
 from dvc.exceptions import DvcException
 from dvc.parsing.versions import LOCKFILE_VERSION, SCHEMA_KWD
@@ -11,21 +21,25 @@ from dvc.stage.exceptions import (
     StageFileDoesNotExistError,
     StageFileIsNotDvcFileError,
 )
-from dvc.types import AnyPath
+from dvc.types import StrOrBytesPath
 from dvc.utils import relpath
 from dvc.utils.collections import apply_diff
+from dvc.utils.objects import cached_property
 from dvc.utils.serialize import dump_yaml, modify_yaml
 
 if TYPE_CHECKING:
     from dvc.repo import Repo
 
+    from .parsing import DataResolver
+    from .stage import Stage
+
 logger = logging.getLogger(__name__)
 _T = TypeVar("_T")
 
 DVC_FILE = "Dvcfile"
 DVC_FILE_SUFFIX = ".dvc"
-PIPELINE_FILE = "dvc.yaml"
-PIPELINE_LOCK = "dvc.lock"
+PROJECT_FILE = "dvc.yaml"
+LOCK_FILE = "dvc.lock"
 
 
 class FileIsGitIgnored(DvcException):
@@ -44,7 +58,7 @@ class ParametrizedDumpError(DvcException):
 def is_valid_filename(path):
     return path.endswith(DVC_FILE_SUFFIX) or os.path.basename(path) in [
         DVC_FILE,
-        PIPELINE_FILE,
+        PROJECT_FILE,
     ]
 
 
@@ -55,7 +69,7 @@ def is_dvc_file(path):
 
 
 def is_lock_file(path):
-    return os.path.basename(path) == PIPELINE_LOCK
+    return os.path.basename(path) == LOCK_FILE
 
 
 def is_git_ignored(repo, path):
@@ -75,7 +89,7 @@ def check_dvcfile_path(repo, path):
         raise StageFileBadNameError(
             "bad DVC file name '{}'. DVC files should be named "
             "'{}' or have a '.dvc' suffix (e.g. '{}.dvc').".format(
-                relpath(path), PIPELINE_FILE, os.path.basename(path)
+                relpath(path), PROJECT_FILE, os.path.basename(path)
             )
         )
 
@@ -150,7 +164,7 @@ class FileMixin:
         return self._load_yaml(**kwargs)
 
     @classmethod
-    def validate(cls, d: _T, fname: str = None) -> _T:
+    def validate(cls, d: _T, fname: Optional[str] = None) -> _T:
         from dvc.utils.strictyaml import validate
 
         return validate(d, cls.SCHEMA, path=fname)  # type: ignore[arg-type]
@@ -180,17 +194,21 @@ class SingleStageFile(FileMixin):
     from dvc.schema import COMPILED_SINGLE_STAGE_SCHEMA as SCHEMA
     from dvc.stage.loader import SingleStageLoader as LOADER
 
+    metrics: List[str] = []
+    plots: Any = {}
+    params: List[str] = []
+
     @property
-    def stage(self):
+    def stage(self) -> "Stage":
         data, raw = self._load()
         return self.LOADER.load_stage(self, data, raw)
 
     @property
-    def stages(self):
+    def stages(self) -> LOADER:
         data, raw = self._load()
         return self.LOADER(self, data, raw)
 
-    def dump(self, stage, **kwargs):
+    def dump(self, stage, **kwargs) -> None:
         """Dumps given stage appropriately in the dvcfile."""
         from dvc.stage import PipelineStage
 
@@ -213,7 +231,7 @@ class SingleStageFile(FileMixin):
         self.dump(stage)
 
 
-class PipelineFile(FileMixin):
+class ProjectFile(FileMixin):
     """Abstraction for pipelines file, .yaml + .lock combined."""
 
     from dvc.schema import COMPILED_MULTI_STAGE_SCHEMA as SCHEMA
@@ -223,6 +241,12 @@ class PipelineFile(FileMixin):
     def _lockfile(self):
         return Lockfile(self.repo, os.path.splitext(self.path)[0] + ".lock")
 
+    def _reset(self):
+        self.__dict__.pop("contents", None)
+        self.__dict__.pop("lockfile_contents", None)
+        self.__dict__.pop("resolver", None)
+        self.__dict__.pop("stages", None)
+
     def dump(
         self, stage, update_pipeline=True, update_lock=True, **kwargs
     ):  # pylint: disable=arguments-differ
@@ -275,14 +299,39 @@ class PipelineFile(FileMixin):
     @property
     def stage(self):
         raise DvcException(
-            "PipelineFile has multiple stages. Please specify it's name."
+            "ProjectFile has multiple stages. Please specify it's name."
         )
 
+    @cached_property
+    def contents(self) -> Dict[str, Any]:
+        return self._load()[0]
+
+    @cached_property
+    def lockfile_contents(self) -> Dict[str, Any]:
+        return self._lockfile.load()
+
+    @cached_property
+    def resolver(self) -> "DataResolver":
+        from .parsing import DataResolver
+
+        wdir = self.repo.fs.path.parent(self.path)
+        return DataResolver(self.repo, wdir, self.contents)
+
+    @cached_property
+    def stages(self) -> LOADER:
+        return self.LOADER(self, self.contents, self.lockfile_contents)
+
     @property
-    def stages(self):
-        data, _ = self._load()
-        lockfile_data = self._lockfile.load()
-        return self.LOADER(self, data, lockfile_data)
+    def metrics(self) -> List[str]:
+        return self.contents.get("metrics", [])
+
+    @property
+    def plots(self) -> Any:
+        return self.contents.get("plots", {})
+
+    @property
+    def params(self) -> List[str]:
+        return self.contents.get("params", [])
 
     def remove(self, force=False):
         if not force:
@@ -415,20 +464,10 @@ class Lockfile(FileMixin):
         raise NotImplementedError
 
 
-class Dvcfile:
-    def __new__(cls, repo: "Repo", path: AnyPath, **kwargs: Any):
-        assert path
-        assert repo
-
-        return make_dvcfile(repo, path, **kwargs)
-
-
-DVCFile = Union["PipelineFile", "SingleStageFile"]
-
-
-def make_dvcfile(repo: "Repo", path: AnyPath, **kwargs: Any) -> DVCFile:
-    _, ext = os.path.splitext(str(path))
-    if ext in [".yaml", ".yml"]:
-        return PipelineFile(repo, path, **kwargs)
-    # fallback to single stage file for better error messages
+def load_file(
+    repo: "Repo", path: StrOrBytesPath, **kwargs: Any
+) -> Union[ProjectFile, SingleStageFile]:
+    _, ext = os.path.splitext(path)
+    if ext in (".yaml", ".yml"):
+        return ProjectFile(repo, path, **kwargs)
     return SingleStageFile(repo, path, **kwargs)

dvc/env.py

@@ -4,6 +4,8 @@ DVC_PAGER = "DVC_PAGER"
 DVC_ROOT = "DVC_ROOT"
 DVCLIVE_RESUME = "DVCLIVE_RESUME"
 DVC_IGNORE_ISATTY = "DVC_IGNORE_ISATTY"
+DVC_EXP_BASELINE_REV = "DVC_EXP_BASELINE_REV"
+DVC_EXP_NAME = "DVC_EXP_NAME"
 DVC_EXP_GIT_REMOTE = "DVC_EXP_GIT_REMOTE"
 DVC_EXP_AUTO_PUSH = "DVC_EXP_AUTO_PUSH"
 DVC_NO_ANALYTICS = "DVC_NO_ANALYTICS"

dvc/exceptions.py

@@ -1,4 +1,6 @@
 """Exceptions raised by the dvc."""
+from typing import Dict, List
+
 from dvc.utils import format_link
 
 
@@ -227,7 +229,7 @@ class UploadError(FileTransferError):
 
 
 class CheckoutError(DvcException):
-    def __init__(self, target_infos, stats=None):
+    def __init__(self, target_infos: List[str], stats: Dict[str, List[str]]):
         from dvc.utils import error_link
 
         self.target_infos = target_infos

dvc/external_repo.py

@@ -3,7 +3,7 @@ import os
 import tempfile
 import threading
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Dict
+from typing import TYPE_CHECKING, Collection, Dict, Iterator, Optional, Tuple
 
 from funcy import retry, wrap_with
 
@@ -17,6 +17,7 @@ from dvc.exceptions import (
 )
 from dvc.repo import Repo
 from dvc.scm import CloneError, map_scm_exception
+from dvc.types import StrPath
 from dvc.utils import relpath
 
 if TYPE_CHECKING:
@@ -28,8 +29,13 @@ logger = logging.getLogger(__name__)
 @contextmanager
 @map_scm_exception()
 def external_repo(
-    url, rev=None, for_write=False, cache_dir=None, cache_types=None, **kwargs
-):
+    url,
+    rev: Optional[str] = None,
+    for_write: bool = False,
+    cache_dir: Optional[StrPath] = None,
+    cache_types: Optional[Collection[str]] = None,
+    **kwargs,
+) -> Iterator["Repo"]:
     from scmrepo.git import Git
 
     from dvc.config import NoRemoteError
@@ -108,7 +114,7 @@ def erepo_factory(url, root_dir, cache_config):
     return make_repo
 
 
-CLONES: Dict[str, str] = {}
+CLONES: Dict[str, Tuple[str, bool]] = {}
 CACHE_DIRS: Dict[str, str] = {}
 
 
@@ -185,14 +191,14 @@ def _cached_clone(url, rev, for_write=False):
 
 
 @wrap_with(threading.Lock())
-def _clone_default_branch(url, rev, for_write=False):
+def _clone_default_branch(url, rev, for_write=False):  # noqa: C901
     """Get or create a clean clone of the url.
 
     The cloned is reactualized with git pull unless rev is a known sha.
     """
     from scmrepo.git import Git
 
-    clone_path, shallow = CLONES.get(url, (None, False))
+    clone_path, shallow = CLONES.get(url) or (None, False)
 
     git = None
     try:

dvc/fs/__init__.py

@@ -72,7 +72,9 @@ def get_fs_config(config, **kwargs):
         except KeyError:
             from dvc.config import RemoteNotFoundError
 
-            raise RemoteNotFoundError(f"remote '{name}' doesn't exist")
+            raise RemoteNotFoundError(  # noqa: B904
+                f"remote '{name}' doesn't exist"
+            )
     else:
         remote_conf = kwargs
     return _resolve_remote_refs(config, remote_conf)

dvc/fs/callbacks.py

@@ -1,9 +1,8 @@
 # pylint: disable=unused-import
 from contextlib import ExitStack
-from typing import TYPE_CHECKING, Any, BinaryIO, Dict, Optional, Union
-
-from funcy import cached_property
+from typing import TYPE_CHECKING, Any, BinaryIO, Dict, Optional, Union, cast
 
+from dvc.utils.objects import cached_property
 from dvc_objects.fs.callbacks import (  # noqa: F401
     DEFAULT_CALLBACK,
     Callback,
@@ -11,6 +10,8 @@ from dvc_objects.fs.callbacks import (  # noqa: F401
 )
 
 if TYPE_CHECKING:
+    from rich.progress import TaskID
+
     from dvc.ui._rich_progress import RichTransferProgress
 
 
@@ -19,10 +20,10 @@ class RichCallback(Callback):
         self,
         size: Optional[int] = None,
         value: int = 0,
-        progress: "RichTransferProgress" = None,
-        desc: str = None,
-        bytes: bool = False,  # pylint: disable=redefined-builtin
-        unit: str = None,
+        progress: Optional["RichTransferProgress"] = None,
+        desc: Optional[str] = None,
+        bytes: bool = False,  # noqa: A002, pylint: disable=redefined-builtin
+        unit: Optional[str] = None,
         disable: bool = False,
         transient: bool = True,
     ) -> None:
@@ -41,7 +42,7 @@ class RichCallback(Callback):
         super().__init__(size=size, value=value)
 
     @cached_property
-    def progress(self):
+    def progress(self) -> "RichTransferProgress":
         from dvc.ui import ui
         from dvc.ui._rich_progress import RichTransferProgress
 
@@ -53,11 +54,13 @@ class RichCallback(Callback):
             disable=self.disable,
             console=ui.error_console,
         )
-        return self._stack.enter_context(progress)
+        return cast(RichTransferProgress, self._stack.enter_context(progress))
 
     @cached_property
-    def task(self):
-        return self.progress.add_task(**self._task_kwargs)
+    def task(self) -> "TaskID":
+        return self.progress.add_task(
+            **self._task_kwargs  # type: ignore[arg-type]
+        )
 
     def __enter__(self):
         return self
@@ -80,7 +83,7 @@ class RichCallback(Callback):
         path_1: Union[str, BinaryIO],
         path_2: str,
         kwargs: Dict[str, Any],
-        child: "Callback" = None,
+        child: Optional["Callback"] = None,
     ):
         child = child or RichCallback(
             progress=self.progress,

dvc/fs/data.py

@@ -1,13 +1,15 @@
+import functools
 import logging
 import os
-import threading
-
-from funcy import cached_property, wrap_prop
+from typing import TYPE_CHECKING
 
 from dvc.utils import as_posix
-from dvc_data.fs import DataFileSystem as _DataFileSystem
 from dvc_objects.fs.base import FileSystem
 
+if TYPE_CHECKING:
+    from dvc_data.fs import DataFileSystem as _DataFileSystem
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -19,9 +21,12 @@ class DataFileSystem(FileSystem):
     def _prepare_credentials(self, **config):
         return config
 
-    @wrap_prop(threading.Lock())
-    @cached_property
-    def fs(self):
+    @functools.cached_property
+    def fs(  # pylint: disable=invalid-overridden-method
+        self,
+    ) -> "_DataFileSystem":
+        from dvc_data.fs import DataFileSystem as _DataFileSystem
+
         return _DataFileSystem(**self.fs_args)
 
     def isdvc(self, path, **kwargs):

dvc/fs/dvc.py

@@ -1,15 +1,26 @@
 import errno
+import functools
 import logging
 import ntpath
 import os
 import posixpath
 import threading
 from contextlib import suppress
-from typing import TYPE_CHECKING, Any, Callable, Optional, Tuple, Type, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Optional,
+    Tuple,
+    Type,
+    Union,
+)
 
 from fsspec.spec import AbstractFileSystem
-from funcy import cached_property, wrap_prop, wrap_with
+from funcy import wrap_with
 
+from dvc.types import StrPath
 from dvc_objects.fs.base import FileSystem
 from dvc_objects.fs.path import Path
 
@@ -20,7 +31,7 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
-RepoFactory = Union[Callable[[str], "Repo"], Type["Repo"]]
+RepoFactory = Union[Callable[..., "Repo"], Type["Repo"]]
 Key = Tuple[str, ...]
 
 
@@ -118,7 +129,7 @@ class _DVCFileSystem(AbstractFileSystem):  # pylint:disable=abstract-method
             repo = self._make_repo(
                 url=url, rev=rev, subrepos=subrepos, **repo_kwargs
             )
-            assert repo
+            assert repo is not None
             # pylint: disable=protected-access
             repo_factory = repo._fs_conf["repo_factory"]
 
@@ -131,6 +142,7 @@ class _DVCFileSystem(AbstractFileSystem):  # pylint:disable=abstract-method
 
         def _getcwd():
             relparts = ()
+            assert repo is not None
             if repo.fs.path.isin(repo.fs.path.getcwd(), repo.root_dir):
                 relparts = repo.fs.path.relparts(
                     repo.fs.path.getcwd(), repo.root_dir
@@ -154,16 +166,16 @@ class _DVCFileSystem(AbstractFileSystem):  # pylint:disable=abstract-method
         if hasattr(repo, "dvc_dir"):
             self._datafss[key] = DataFileSystem(index=repo.index.data["repo"])
 
-    def _get_key(self, path) -> Key:
+    def _get_key(self, path: StrPath) -> Key:
         parts = self.repo.fs.path.relparts(path, self.repo.root_dir)
-        if parts == (".",):
-            parts = ()
+        if parts == (os.curdir,):
+            return ()
         return parts
 
     def _get_key_from_relative(self, path) -> Key:
         parts = self.path.relparts(path, self.root_marker)
         if parts and parts[0] == os.curdir:
-            parts = parts[1:]
+            return parts[1:]
         return parts
 
     def _from_key(self, parts: Key) -> str:
@@ -171,8 +183,6 @@ class _DVCFileSystem(AbstractFileSystem):  # pylint:disable=abstract-method
 
     @property
     def repo_url(self):
-        if self.repo is None:
-            return None
         return self.repo.url
 
     @classmethod
@@ -276,43 +286,38 @@ class _DVCFileSystem(AbstractFileSystem):  # pylint:disable=abstract-method
         key = self._get_key_from_relative(path)
         repo, dvc_fs, subkey = self._get_subrepo_info(key)
 
-        names = set()
+        dvc_infos = {}
         if dvc_fs:
             with suppress(FileNotFoundError):
                 dvc_path = _get_dvc_path(dvc_fs, subkey)
-                for entry in dvc_fs.ls(dvc_path, detail=False):
-                    names.add(dvc_fs.path.name(entry))
+                for info in dvc_fs.ls(dvc_path, detail=True):
+                    dvc_infos[dvc_fs.path.name(info["name"])] = info
 
+        fs_infos = {}
         ignore_subrepos = kwargs.get("ignore_subrepos", True)
         if not dvc_only:
             fs = self.repo.fs
             fs_path = self._from_key(key)
             try:
-                for entry in repo.dvcignore.ls(
-                    fs, fs_path, detail=False, ignore_subrepos=ignore_subrepos
+                for info in repo.dvcignore.ls(
+                    fs, fs_path, detail=True, ignore_subrepos=ignore_subrepos
                 ):
-                    names.add(fs.path.name(entry))
+                    fs_infos[fs.path.name(info["name"])] = info
             except (FileNotFoundError, NotADirectoryError):
                 pass
 
         dvcfiles = kwargs.get("dvcfiles", False)
-        if not dvcfiles:
-            names = (name for name in names if not _is_dvc_file(name))
 
         infos = []
         paths = []
+        names = set(dvc_infos.keys()) | set(fs_infos.keys())
         for name in names:
-            entry_path = self.path.join(path, name)
-            entry_key = key + (name,)
-            try:
-                info = self._info(
-                    entry_key,
-                    entry_path,
-                    ignore_subrepos=ignore_subrepos,
-                    check_ignored=False,
-                )
-            except FileNotFoundError:
+            if not dvcfiles and _is_dvc_file(name):
                 continue
+
+            entry_path = self.path.join(path, name)
+            info = _merge_info(repo, fs_infos.get(name), dvc_infos.get(name))
+            info["name"] = entry_path
             infos.append(info)
             paths.append(entry_path)
 
@@ -326,7 +331,9 @@ class _DVCFileSystem(AbstractFileSystem):  # pylint:disable=abstract-method
         ignore_subrepos = kwargs.get("ignore_subrepos", True)
         return self._info(key, path, ignore_subrepos=ignore_subrepos)
 
-    def _info(self, key, path, ignore_subrepos=True, check_ignored=True):
+    def _info(  # noqa: C901
+        self, key, path, ignore_subrepos=True, check_ignored=True
+    ):
         repo, dvc_fs, subkey = self._get_subrepo_info(key)
 
         dvc_info = None
@@ -374,30 +381,30 @@ class DVCFileSystem(FileSystem):
     protocol = "local"
     PARAM_CHECKSUM = "md5"
 
-    def _prepare_credentials(self, **config):
+    def _prepare_credentials(self, **config) -> Dict[str, Any]:
         return config
 
-    @wrap_prop(threading.Lock())
-    @cached_property
-    def fs(self):
+    @functools.cached_property
+    # pylint: disable-next=invalid-overridden-method
+    def fs(self) -> "DVCFileSystem":
         return _DVCFileSystem(**self.fs_args)
 
-    def isdvc(self, path, **kwargs):
+    def isdvc(self, path, **kwargs) -> bool:
         return self.fs.isdvc(path, **kwargs)
 
     @property
-    def path(self):  # pylint: disable=invalid-overridden-method
+    def path(self) -> Path:  # pylint: disable=invalid-overridden-method
         return self.fs.path
 
     @property
-    def repo(self):
+    def repo(self) -> "Repo":
         return self.fs.repo
 
     @property
-    def repo_url(self):
+    def repo_url(self) -> str:
         return self.fs.repo_url
 
-    def from_os_path(self, path):
+    def from_os_path(self, path: str) -> str:
         if os.path.isabs(path):
             path = os.path.relpath(path, self.repo.root_dir)
 

dvc/fs/git.py

@@ -1,7 +1,5 @@
-import threading
-from typing import TYPE_CHECKING, Any
-
-from funcy import cached_property, wrap_prop
+import functools
+from typing import TYPE_CHECKING, Any, Optional
 
 from . import FileSystem
 
@@ -19,10 +17,10 @@ class GitFileSystem(FileSystem):  # pylint:disable=abstract-method
 
     def __init__(
         self,
-        path: str = None,
-        rev: str = None,
-        scm: "Git" = None,
-        trie: "GitTrie" = None,
+        path: Optional[str] = None,
+        rev: Optional[str] = None,
+        scm: Optional["Git"] = None,
+        trie: Optional["GitTrie"] = None,
         **kwargs: Any,
     ) -> None:
         from dvc.scm import resolve_rev
@@ -39,15 +37,16 @@ class GitFileSystem(FileSystem):  # pylint:disable=abstract-method
             }
         )
 
-    @wrap_prop(threading.Lock())
-    @cached_property
-    def fs(self) -> "FsspecGitFileSystem":
+    @functools.cached_property
+    def fs(  # pylint: disable=invalid-overridden-method
+        self,
+    ) -> "FsspecGitFileSystem":
         from scmrepo.fs import GitFileSystem as FsspecGitFileSystem
 
         return FsspecGitFileSystem(**self.fs_args)
 
-    @cached_property
-    def path(self):
+    @functools.cached_property
+    def path(self):  # pylint: disable=invalid-overridden-method
         return self.fs.path
 
     @property

dvc/ignore.py

@@ -3,6 +3,7 @@ import os
 import re
 from collections import namedtuple
 from itertools import chain, groupby, takewhile
+from typing import List, Optional
 
 from pathspec.patterns import GitWildMatchPattern
 from pathspec.util import normalize_file
@@ -10,7 +11,6 @@ from pygtrie import Trie
 
 from dvc.fs import AnyFSPath, FileSystem, Schemes, localfs
 from dvc.pathspec_math import PatternInfo, merge_patterns
-from dvc.types import List, Optional
 
 logger = logging.getLogger(__name__)
 
@@ -24,20 +24,27 @@ class DvcIgnore:
 
 class DvcIgnorePatterns(DvcIgnore):
     def __init__(self, pattern_list, dirname, sep):
-        if pattern_list:
-            if isinstance(pattern_list[0], str):
-                pattern_list = [
-                    PatternInfo(pattern, "") for pattern in pattern_list
-                ]
+        from pathspec.patterns.gitwildmatch import _DIR_MARK
+
+        if pattern_list and isinstance(pattern_list[0], str):
+            pattern_list = [
+                PatternInfo(pattern, "") for pattern in pattern_list
+            ]
 
         self.sep = sep
         self.pattern_list = pattern_list
         self.dirname = dirname
 
-        self.regex_pattern_list = [
-            GitWildMatchPattern.pattern_to_regex(pattern_info.patterns)
-            for pattern_info in pattern_list
-        ]
+        self.regex_pattern_list = []
+        for count, pattern in enumerate(pattern_list):
+            pattern, group = GitWildMatchPattern.pattern_to_regex(
+                pattern.patterns
+            )
+            if pattern:
+                pattern = pattern.replace(
+                    f"<{_DIR_MARK}>", f"<{_DIR_MARK}{count}>"
+                )
+                self.regex_pattern_list.append((pattern, group))
 
         self.ignore_spec = [
             (ignore, re.compile("|".join(item[0] for item in group)))
@@ -84,7 +91,7 @@ class DvcIgnorePatterns(DvcIgnore):
             return False
 
         if os.name == "nt":
-            path = normalize_file(path)
+            return normalize_file(path)
         return path
 
     def matches(self, dirname, basename, is_dir=False, details: bool = False):
@@ -195,8 +202,8 @@ class DvcIgnoreFilter:
 
     def _get_key(self, path):
         parts = self.fs.path.relparts(path, self.root_dir)
-        if parts == (".",):
-            parts = ()
+        if parts == (os.curdir,):
+            return ()
         return parts
 
     def _update_trie(self, dirname: str, trie: Trie) -> None:
@@ -379,7 +386,7 @@ class DvcIgnoreFilter:
         return False
 
     def is_ignored_dir(self, path: str, ignore_subrepos: bool = True) -> bool:
-        "Only used in LocalFileSystem"
+        # only used in LocalFileSystem
         path = self.fs.path.abspath(path)
         if path == self.root_dir:
             return False
@@ -387,7 +394,7 @@ class DvcIgnoreFilter:
         return self._is_ignored(path, True, ignore_subrepos=ignore_subrepos)
 
     def is_ignored_file(self, path: str, ignore_subrepos: bool = True) -> bool:
-        "Only used in LocalFileSystem"
+        # only used in LocalFileSystem
         path = self.fs.path.abspath(path)
         return self._is_ignored(path, False, ignore_subrepos=ignore_subrepos)
 

dvc/lock.py

@@ -116,7 +116,7 @@ class Lock(LockBase):
                 self._lock = zc.lockfile.LockFile(self._lockfile)
         except zc.lockfile.LockError:
             self._lock_failed = True
-            raise LockError(FAILED_TO_LOCK_MESSAGE)
+            raise LockError(FAILED_TO_LOCK_MESSAGE)  # noqa: B904
 
     def lock(self):
         retries = 6
@@ -131,6 +131,7 @@ class Lock(LockBase):
 
         if not self.is_locked:
             raise DvcException("Unlock called on an unlocked lock")
+        assert self._lock
         self._lock.close()
         self._lock = None
 
@@ -185,14 +186,15 @@ class HardlinkLock(flufl.lock.Lock, LockBase):
         try:
             super().lock(timedelta(seconds=DEFAULT_TIMEOUT))
         except flufl.lock.TimeOutError:
-            raise LockError(FAILED_TO_LOCK_MESSAGE)
+            raise LockError(FAILED_TO_LOCK_MESSAGE)  # noqa: B904
 
     def _set_claimfile(self):
         super()._set_claimfile()
 
         if self._tmp_dir is not None:
             # Under Windows file path length is limited so we hash it
-            filename = hashlib.md5(self._claimfile.encode()).hexdigest()
+            hasher = hashlib.md5(self._claimfile.encode())  # noqa: S324
+            filename = hasher.hexdigest()
             self._claimfile = os.path.join(self._tmp_dir, filename + ".lock")
 
 

dvc/logger.py

@@ -204,7 +204,7 @@ def disable_other_loggers():
     loggerDict = logging.root.manager.loggerDict  # pylint: disable=no-member
     for logger_name, logger in loggerDict.items():
         if logger_name != "dvc" and not logger_name.startswith("dvc."):
-            logger.disabled = True
+            logger.disabled = True  # type: ignore[union-attr]
 
 
 def set_loggers_level(level: int = logging.INFO) -> None:

dvc/machine/__init__.py

@@ -127,6 +127,7 @@ class MachineManager:
         self, repo: "Repo", backends: Optional[Iterable[str]] = None, **kwargs
     ):
         self.repo = repo
+        assert self.repo.tmp_dir
         tmp_dir = os.path.join(self.repo.tmp_dir, "machine")
         self.backends = MachineBackends(backends, tmp_dir=tmp_dir, **kwargs)
 
@@ -167,7 +168,9 @@ class MachineManager:
             except KeyError:
                 from dvc.config import MachineNotFoundError
 
-                raise MachineNotFoundError(f"machine '{name}' doesn't exist")
+                raise MachineNotFoundError(  # noqa: B904
+                    f"machine '{name}' doesn't exist"
+                )
         else:
             conf = kwargs
         return conf
@@ -179,7 +182,9 @@ class MachineManager:
             backend = self.CLOUD_BACKENDS[cloud]
             return self.backends[backend]
         except KeyError:
-            raise NoMachineError(f"Machine platform '{cloud}' unsupported")
+            raise NoMachineError(  # noqa: B904
+                f"Machine platform '{cloud}' unsupported"
+            )
 
     def create(self, name: Optional[str]):
         """Create and start the specified machine instance."""

dvc/machine/backend/terraform.py

@@ -6,7 +6,6 @@ from typing import TYPE_CHECKING, Iterator, Optional
 from dvc_ssh import DEFAULT_PORT, SSHFileSystem
 
 from dvc.exceptions import DvcException
-from dvc.utils.fs import makedirs
 
 from .base import BaseMachineBackend
 
@@ -30,7 +29,7 @@ def _sshfs(resource: dict):
 class TerraformBackend(BaseMachineBackend):
     def __init__(self, tmp_dir: "StrPath", **kwargs):
         super().__init__(tmp_dir, **kwargs)
-        makedirs(tmp_dir, exist_ok=True)
+        os.makedirs(tmp_dir, exist_ok=True)
 
     @contextmanager
     def make_tpi(self, name: str):
@@ -39,7 +38,7 @@ class TerraformBackend(BaseMachineBackend):
 
         try:
             working_dir = os.path.join(self.tmp_dir, name)
-            makedirs(working_dir, exist_ok=True)
+            os.makedirs(working_dir, exist_ok=True)
             yield TPIBackend(working_dir=working_dir)
         except TPIError as exc:
             raise DvcException("TPI operation failed") from exc

dvc/odbmgr.py

@@ -39,7 +39,7 @@ class ODBManager:
             from dvc.config_schema import LOCAL_COMMON
 
             settings = {"url": config["dir"]}
-            for opt in LOCAL_COMMON.keys():
+            for opt in LOCAL_COMMON:
                 if opt in config:
                     settings[str(opt)] = config.get(opt)
 

dvc/output.py

@@ -1,10 +1,10 @@
 import logging
 import os
 from collections import defaultdict
-from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple, Type
+from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple, Type, cast
 from urllib.parse import urlparse
 
-from funcy import cached_property, collecting, project
+from funcy import collecting, project
 from voluptuous import And, Any, Coerce, Length, Lower, Required, SetTo
 
 from dvc import prompt
@@ -17,6 +17,7 @@ from dvc.exceptions import (
     MergeError,
     RemoteCacheRequiredError,
 )
+from dvc.utils.objects import cached_property
 from dvc_data.hashfile import Tree
 from dvc_data.hashfile import check as ocheck
 from dvc_data.hashfile import load as oload
@@ -31,7 +32,7 @@ from dvc_objects.errors import ObjectFormatError
 
 from .annotations import ANNOTATION_FIELDS, ANNOTATION_SCHEMA, Annotation
 from .fs import LocalFileSystem, RemoteMissingDepsError, Schemes, get_cloud_fs
-from .fs.callbacks import DEFAULT_CALLBACK
+from .fs.callbacks import DEFAULT_CALLBACK, Callback
 from .utils import relpath
 from .utils.fs import path_isin
 
@@ -40,7 +41,7 @@ if TYPE_CHECKING:
     from dvc_data.index import DataIndexKey
     from dvc_objects.db import ObjectDB
 
-    from .fs.callbacks import Callback
+    from .ignore import DvcIgnoreFilter
 
 logger = logging.getLogger(__name__)
 
@@ -143,7 +144,7 @@ def _split_dict(d, keys):
 
 
 def _merge_data(s_list):
-    d = defaultdict(dict)
+    d: Dict[str, Dict] = defaultdict(dict)
     for key in s_list:
         if isinstance(key, str):
             d[key].update({})
@@ -267,10 +268,10 @@ class Output:
         },
     )
 
-    DoesNotExistError = OutputDoesNotExistError  # type: Type[DvcException]
-    IsNotFileOrDirError = OutputIsNotFileOrDirError  # type: Type[DvcException]
-    IsStageFileError = OutputIsStageFileError  # type: Type[DvcException]
-    IsIgnoredError = OutputIsIgnoredError  # type: Type[DvcException]
+    DoesNotExistError: Type[DvcException] = OutputDoesNotExistError
+    IsNotFileOrDirError: Type[DvcException] = OutputIsNotFileOrDirError
+    IsStageFileError: Type[DvcException] = OutputIsStageFileError
+    IsIgnoredError: Type[DvcException] = OutputIsIgnoredError
 
     def __init__(
         self,
@@ -283,13 +284,13 @@ class Output:
         persist=False,
         checkpoint=False,
         desc=None,
-        type=None,  # pylint: disable=redefined-builtin
+        type=None,  # noqa: A002, pylint: disable=redefined-builtin
         labels=None,
         meta=None,
         remote=None,
         repo=None,
         fs_config=None,
-        files: List[Dict[str, Any]] = None,
+        files: Optional[List[Dict[str, Any]]] = None,
         push: bool = True,
     ):
         self.annot = Annotation(
@@ -315,7 +316,7 @@ class Output:
             and isinstance(stage.repo.fs, LocalFileSystem)
             and path_isin(path, stage.repo.root_dir)
         ):
-            self.def_path = relpath(path, stage.wdir)
+            self.def_path: str = relpath(path, stage.wdir)
             self.fs = stage.repo.fs
         else:
             self.def_path = path
@@ -486,10 +487,9 @@ class Output:
         return self.hash_info.isdir
 
     def _is_path_dvcignore(self, path) -> bool:
-        if not self.IS_DEPENDENCY and self.dvcignore:
-            if self.dvcignore.is_ignored(self.fs, path, ignore_subrepos=False):
-                return True
-        return False
+        if self.IS_DEPENDENCY or not self.dvcignore:
+            return False
+        return self.dvcignore.is_ignored(self.fs, path, ignore_subrepos=False)
 
     @property
     def exists(self):
@@ -561,7 +561,7 @@ class Output:
             return self.meta.version_id != self.get_meta().version_id
         return False
 
-    def workspace_status(self):
+    def workspace_status(self) -> Dict[str, str]:
         if not self.exists:
             return {str(self): "deleted"}
 
@@ -573,40 +573,40 @@ class Output:
 
         return {}
 
-    def status(self):
+    def status(self) -> Dict[str, str]:
         if self.hash_info and self.use_cache and self.changed_cache():
             return {str(self): "not in cache"}
 
         return self.workspace_status()
 
-    def changed(self):
+    def changed(self) -> bool:
         status = self.status()
         logger.debug(str(status))
         return bool(status)
 
     @property
-    def dvcignore(self):
+    def dvcignore(self) -> Optional["DvcIgnoreFilter"]:
         if self.fs.protocol == "local":
             return self.repo.dvcignore
         return None
 
     @property
-    def is_empty(self):
+    def is_empty(self) -> bool:
         return self.fs.is_empty(self.fs_path)
 
-    def isdir(self):
+    def isdir(self) -> bool:
         if self._is_path_dvcignore(self.fs_path):
             return False
         return self.fs.isdir(self.fs_path)
 
-    def isfile(self):
+    def isfile(self) -> bool:
         if self._is_path_dvcignore(self.fs_path):
             return False
         return self.fs.isfile(self.fs_path)
 
     # pylint: disable=no-member
 
-    def ignore(self):
+    def ignore(self) -> None:
         if not self.use_scm_ignore:
             return
 
@@ -615,7 +615,7 @@ class Output:
 
         self.repo.scm_context.ignore(self.fspath)
 
-    def ignore_remove(self):
+    def ignore_remove(self) -> None:
         if not self.use_scm_ignore:
             return
 
@@ -623,11 +623,11 @@ class Output:
 
     # pylint: enable=no-member
 
-    def save(self):
+    def save(self) -> None:
         if not self.exists:
             raise self.DoesNotExistError(self)
 
-        if not self.isfile and not self.isdir:
+        if not self.isfile() and not self.isdir():
             raise self.IsNotFileOrDirError(self)
 
         if self.is_empty:
@@ -663,11 +663,11 @@ class Output:
         self.hash_info = self.obj.hash_info
         self.files = None
 
-    def set_exec(self):
+    def set_exec(self) -> None:
         if self.isfile() and self.meta.isexec:
             self.odb.set_exec(self.fs_path)
 
-    def _checkout(self, *args, **kwargs):
+    def _checkout(self, *args, **kwargs) -> Optional[bool]:
         from dvc_data.hashfile.checkout import CheckoutError as _CheckoutError
         from dvc_data.hashfile.checkout import LinkError, PromptError
 
@@ -675,13 +675,13 @@ class Output:
         try:
             return checkout(*args, **kwargs)
         except PromptError as exc:
-            raise ConfirmRemoveError(exc.path)
+            raise ConfirmRemoveError(exc.path)  # noqa: B904
         except LinkError as exc:
-            raise CacheLinkError([exc.path])
+            raise CacheLinkError([exc.path])  # noqa: B904
         except _CheckoutError as exc:
-            raise CheckoutError(exc.paths)
+            raise CheckoutError(exc.paths, {})  # noqa: B904
 
-    def commit(self, filter_info=None):
+    def commit(self, filter_info=None) -> None:
         if not self.exists:
             raise self.DoesNotExistError(self)
 
@@ -721,7 +721,7 @@ class Output:
             )
             self.set_exec()
 
-    def _commit_granular_dir(self, filter_info):
+    def _commit_granular_dir(self, filter_info) -> Optional["HashFile"]:
         prefix = self.fs.path.parts(
             self.fs.path.relpath(filter_info, self.fs_path)
         )
@@ -732,7 +732,8 @@ class Output:
             self.hash_name,
             ignore=self.dvcignore,
         )
-        save_obj = save_obj.filter(prefix)
+        save_obj = cast(Tree, save_obj)
+        save_obj = cast(Tree, save_obj.filter(prefix))
         checkout_obj = save_obj.get_obj(self.odb, prefix)
         otransfer(
             staging,
@@ -743,10 +744,10 @@ class Output:
         )
         return checkout_obj
 
-    def dumpd(self, **kwargs):
+    def dumpd(self, **kwargs):  # noqa: C901
         meta = self.meta.to_dict()
         meta.pop("isdir", None)
-        ret = {**self.hash_info.to_dict(), **meta}
+        ret: Dict[str, Any] = {**self.hash_info.to_dict(), **meta}
 
         if self.is_in_repo:
             path = self.fs.path.as_posix(
@@ -762,12 +763,12 @@ class Output:
             if not self.use_cache:
                 ret[self.PARAM_CACHE] = self.use_cache
 
-            if isinstance(self.metric, dict):
-                if (
-                    self.PARAM_METRIC_XPATH in self.metric
-                    and not self.metric[self.PARAM_METRIC_XPATH]
-                ):
-                    del self.metric[self.PARAM_METRIC_XPATH]
+            if (
+                isinstance(self.metric, dict)
+                and self.PARAM_METRIC_XPATH in self.metric
+                and not self.metric[self.PARAM_METRIC_XPATH]
+            ):
+                del self.metric[self.PARAM_METRIC_XPATH]
 
             if self.metric:
                 ret[self.PARAM_METRIC] = self.metric
@@ -792,11 +793,14 @@ class Output:
             and self.hash_info.isdir
             and (kwargs.get("with_files") or self.files is not None)
         ):
+            obj: Optional["HashFile"]
             if self.obj:
                 obj = self.obj
             else:
                 obj = self.get_obj()
-            ret[self.PARAM_FILES] = obj.as_list(with_meta=True)
+            if obj:
+                obj = cast(Tree, obj)
+                ret[self.PARAM_FILES] = obj.as_list(with_meta=True)
 
         return ret
 
@@ -844,20 +848,21 @@ class Output:
         fs_path = self.fs.path
         if filter_info and filter_info != self.fs_path:
             prefix = fs_path.relparts(filter_info, self.fs_path)
+            obj = cast(Tree, obj)
             obj = obj.get_obj(self.odb, prefix)
 
         return obj
 
     def checkout(
         self,
-        force=False,
+        force: bool = False,
         progress_callback: "Callback" = DEFAULT_CALLBACK,
-        relink=False,
-        filter_info=None,
-        allow_missing=False,
-        checkpoint_reset=False,
+        relink: bool = False,
+        filter_info: Optional[str] = None,
+        allow_missing: bool = False,
+        checkpoint_reset: bool = False,
         **kwargs,
-    ):
+    ) -> Optional[Tuple[bool, Optional[bool]]]:
         if not self.use_cache:
             if progress_callback != DEFAULT_CALLBACK:
                 progress_callback.relative_update(
@@ -898,7 +903,7 @@ class Output:
         return added, False if added else modified
 
     def remove(self, ignore_remove=False):
-        self.fs.remove(self.fs_path)
+        self.fs.remove(self.fs_path, recursive=True)
         if self.protocol != Schemes.LOCAL:
             return
 
@@ -1020,7 +1025,7 @@ class Output:
                 "Would you like to continue? Use '-f' to force."
             )
             if not force and not prompt.confirm(msg.format(self.fs_path)):
-                raise CollectCacheError(
+                raise CollectCacheError(  # noqa: B904
                     "unable to fully collect used cache"
                     " without cache for directory '{}'".format(self)
                 )
@@ -1032,10 +1037,11 @@ class Output:
             prefix = self.fs.path.parts(
                 self.fs.path.relpath(filter_info, self.fs_path)
             )
-            obj = obj.filter(prefix)
-        return obj
+            obj = cast(Tree, obj)
+            return obj.filter(prefix)
+        return cast(Tree, obj)
 
-    def get_used_objs(
+    def get_used_objs(  # noqa: C901
         self, **kwargs
     ) -> Dict[Optional["ObjectDB"], Set["HashInfo"]]:
         """Return filtered set of used object IDs for this out."""
@@ -1071,6 +1077,7 @@ class Output:
             logger.warning(msg)
             return {}
 
+        obj: Optional["HashFile"]
         if self.is_dir_checksum:
             obj = self._collect_used_dir_cache(**kwargs)
         else:

dvc/parsing/__init__.py

@@ -14,10 +14,11 @@ from typing import (
     Union,
 )
 
-from funcy import cached_property, collecting, first, isa, join, reraise
+from funcy import collecting, first, isa, join, reraise
 
 from dvc.exceptions import DvcException
 from dvc.parsing.interpolate import ParseError
+from dvc.utils.objects import cached_property
 
 from .context import (
     Context,
@@ -40,6 +41,8 @@ if TYPE_CHECKING:
 
     from dvc.repo import Repo
 
+from dvc.types import DictStrAny
+
 logger = logging.getLogger(__name__)
 
 STAGES_KWD = "stages"
@@ -52,7 +55,6 @@ DO_KWD = "do"
 DEFAULT_PARAMS_FILE = "params.yaml"
 
 JOIN = "@"
-DictStr = Dict[str, Any]
 
 
 class ResolveError(DvcException):
@@ -81,7 +83,7 @@ def format_and_raise(exc: Exception, msg: str, path: str) -> "NoReturn":
 
 
 def _reraise_err(
-    exc_cls: Type[Exception], *args, from_exc: Exception = None
+    exc_cls: Type[Exception], *args, from_exc: Optional[Exception] = None
 ) -> "NoReturn":
     err = exc_cls(*args)
     if from_exc and logger.isEnabledFor(logging.DEBUG):
@@ -90,7 +92,7 @@ def _reraise_err(
 
 
 def check_syntax_errors(
-    definition: DictStr, name: str, path: str, where: str = "stages"
+    definition: DictStrAny, name: str, path: str, where: str = "stages"
 ):
     for key, d in definition.items():
         try:
@@ -109,8 +111,8 @@ def split_foreach_name(name: str) -> Tuple[str, Optional[str]]:
     return group, first(keys)
 
 
-def check_interpolations(data: DictStr, where: str, path: str):
-    def func(s: DictStr) -> None:
+def check_interpolations(data: DictStrAny, where: str, path: str):
+    def func(s: DictStrAny) -> None:
         if is_interpolated_string(s):
             raise ResolveError(
                 _format_preamble(f"'{where}'", path)
@@ -124,7 +126,7 @@ Definition = Union["ForeachDefinition", "EntryDefinition"]
 
 
 def make_definition(
-    resolver: "DataResolver", name: str, definition: DictStr, **kwargs
+    resolver: "DataResolver", name: str, definition: DictStrAny, **kwargs
 ) -> Definition:
     args = resolver, resolver.context, name, definition
     if FOREACH_KWD in definition:
@@ -193,7 +195,7 @@ class DataResolver:
     def has_key(self, key: str):
         return self._has_group_and_key(*split_foreach_name(key))
 
-    def _has_group_and_key(self, group: str, key: str = None):
+    def _has_group_and_key(self, group: str, key: Optional[str] = None):
         try:
             definition = self.definitions[group]
         except KeyError:
@@ -223,7 +225,7 @@ class EntryDefinition:
         resolver: DataResolver,
         context: Context,
         name: str,
-        definition: DictStr,
+        definition: DictStrAny,
         where: str = STAGES_KWD,
     ):
         self.resolver = resolver
@@ -235,7 +237,7 @@ class EntryDefinition:
         self.where = where
 
     def _resolve_wdir(
-        self, context: Context, name: str, wdir: str = None
+        self, context: Context, name: str, wdir: Optional[str] = None
     ) -> str:
         if not wdir:
             return self.wdir
@@ -252,7 +254,7 @@ class EntryDefinition:
         except ContextError as exc:
             format_and_raise(exc, f"stage '{self.name}'", self.relpath)
 
-    def resolve_stage(self, skip_checks: bool = False) -> DictStr:
+    def resolve_stage(self, skip_checks: bool = False) -> DictStrAny:
         context = self.context
         name = self.name
         if not skip_checks:
@@ -300,7 +302,7 @@ class EntryDefinition:
 
     def _resolve(
         self, context: "Context", value: Any, key: str, skip_checks: bool
-    ) -> DictStr:
+    ) -> DictStrAny:
         try:
             return context.resolve(
                 value, skip_interpolation_checks=skip_checks, key=key
@@ -322,7 +324,7 @@ class ForeachDefinition:
         resolver: DataResolver,
         context: Context,
         name: str,
-        definition: DictStr,
+        definition: DictStrAny,
         where: str = STAGES_KWD,
     ):
         self.resolver = resolver
@@ -415,13 +417,13 @@ class ForeachDefinition:
     def _generate_name(self, key: str) -> str:
         return f"{self.name}{JOIN}{key}"
 
-    def resolve_all(self) -> DictStr:
+    def resolve_all(self) -> DictStrAny:
         return join(map(self.resolve_one, self.normalized_iterable))
 
-    def resolve_one(self, key: str) -> DictStr:
+    def resolve_one(self, key: str) -> DictStrAny:
         return self._each_iter(key)
 
-    def _each_iter(self, key: str) -> DictStr:
+    def _each_iter(self, key: str) -> DictStrAny:
         err_message = f"Could not find '{key}' in foreach group '{self.name}'"
         with reraise(KeyError, EntryNotFound(err_message)):
             value = self.normalized_iterable[key]

dvc/parsing/context.py

@@ -172,7 +172,7 @@ class Container(Node, ABC):  # noqa: B024
         return self._convert_with_meta(value, meta)
 
     @staticmethod
-    def _convert_with_meta(value, meta: Meta = None):
+    def _convert_with_meta(value, meta: Optional[Meta] = None):
         if value is None or isinstance(value, PRIMITIVES):
             assert meta
             return Value(value, meta=meta)
@@ -240,7 +240,7 @@ class Container(Node, ABC):  # noqa: B024
 class CtxList(Container, MutableSequence):
     _key_transform = staticmethod(int)
 
-    def __init__(self, values: Sequence, meta: Meta = None):
+    def __init__(self, values: Sequence, meta: Optional[Meta] = None):
         super().__init__(meta=meta)
         self.data: list = []
         self.extend(values)
@@ -263,7 +263,12 @@ class CtxList(Container, MutableSequence):
 
 
 class CtxDict(Container, MutableMapping):
-    def __init__(self, mapping: Mapping = None, meta: Meta = None, **kwargs):
+    def __init__(
+        self,
+        mapping: Optional[Mapping] = None,
+        meta: Optional[Meta] = None,
+        **kwargs,
+    ):
         super().__init__(meta=meta)
 
         self.data: dict = {}
@@ -301,7 +306,7 @@ class Context(CtxDict):
         """
         super().__init__(*args, **kwargs)
         self._track = False
-        self._tracked_data = defaultdict(dict)
+        self._tracked_data: Dict[str, Dict] = defaultdict(dict)
         self.imports = {}
         self._reserved_keys = {}
 
@@ -354,7 +359,7 @@ class Context(CtxDict):
 
     @classmethod
     def load_from(
-        cls, fs, path: str, select_keys: List[str] = None
+        cls, fs, path: str, select_keys: Optional[List[str]] = None
     ) -> "Context":
         from dvc.utils.serialize import load_path
 
@@ -414,27 +419,27 @@ class Context(CtxDict):
             self.imports[path].extend(cp)
 
     def check_loaded(self, path, item, keys):
-        if not keys and isinstance(self.imports[path], list):
+        imported = self.imports[path]
+        if not keys and isinstance(imported, list):
             raise VarsAlreadyLoaded(
                 f"cannot load '{item}' as it's partially loaded already"
             )
-        elif keys and self.imports[path] is None:
+        if keys and imported is None:
             raise VarsAlreadyLoaded(
                 f"cannot partially load '{item}' as it's already loaded."
             )
-        elif isinstance(self.imports[path], list):
-            if not set(keys).isdisjoint(set(self.imports[path])):
-                raise VarsAlreadyLoaded(
-                    f"cannot load '{item}' as it's partially loaded already"
-                )
+        if isinstance(imported, list) and set(keys) & set(imported):
+            raise VarsAlreadyLoaded(
+                f"cannot load '{item}' as it's partially loaded already"
+            )
 
     def load_from_vars(
         self,
         fs,
         vars_: List,
         wdir: str,
-        stage_name: str = None,
-        default: str = None,
+        stage_name: Optional[str] = None,
+        default: Optional[str] = None,
     ):
         if default:
             to_import = fs.path.join(wdir, default)
@@ -482,7 +487,7 @@ class Context(CtxDict):
         try:
             yield
         finally:
-            for key in new.keys():
+            for key in new:
                 self._reserved_keys.pop(key)
 
     @contextmanager

dvc/parsing/interpolate.py

@@ -92,7 +92,7 @@ def _(obj: bool):
 
 
 @to_str.register(dict)
-def _(obj: dict):
+def _(obj: dict):  # noqa: C901
     from dvc.config import Config
 
     config = Config().get("parsing", {})
@@ -186,7 +186,7 @@ def parse_expr(s: str):
         result = get_parser().parseString(s, parseAll=True)
     except ParseException as exc:
         format_and_raise_parse_error(exc)
-        raise AssertionError("unreachable")
+        raise AssertionError("unreachable")  # noqa: B904
 
     joined = result.asList()
     assert len(joined) == 1
@@ -204,9 +204,8 @@ def validate_value(value, key):
     not_primitive = value is not None and not isinstance(value, PRIMITIVES)
     not_foreach = key is not None and "foreach" not in key
     if not_primitive and not_foreach:
-        if isinstance(value, dict):
-            if key == "cmd":
-                return True
+        if isinstance(value, dict) and key == "cmd":
+            return True
         raise ParseError(
             f"Cannot interpolate data of type '{type(value).__name__}'"
         )

dvc/pathspec_math.py

@@ -50,8 +50,7 @@ def change_rule(rule, rel):
         rule = f"!/{rel}{rule}"
     else:
         rule = f"/{rel}{rule}"
-    rule = normalize_file(rule)
-    return rule
+    return normalize_file(rule)
 
 
 def _change_dirname(dirname, pattern_list, new_dirname):

dvc/proc/manager.py

@@ -38,7 +38,7 @@ class ProcessManager:
             with open(info_path, encoding="utf-8") as fobj:
                 return ProcessInfo.from_dict(json.load(fobj))
         except FileNotFoundError:
-            raise KeyError
+            raise KeyError  # noqa: B904
 
     @reraise(FileNotFoundError, KeyError)
     def __setitem__(self, key: str, value: "ProcessInfo"):
@@ -83,13 +83,12 @@ class ProcessManager:
     def send_signal(self, name: str, sig: int):
         """Send `signal` to the specified named process."""
         process_info = self[name]
-        if sys.platform == "win32":
-            if sig not in (
-                signal.SIGTERM,
-                signal.CTRL_C_EVENT,
-                signal.CTRL_BREAK_EVENT,
-            ):
-                raise UnsupportedSignalError(sig)
+        if sys.platform == "win32" and sig not in (
+            signal.SIGTERM,
+            signal.CTRL_C_EVENT,  # pylint: disable=no-member
+            signal.CTRL_BREAK_EVENT,  # pylint: disable=no-member
+        ):
+            raise UnsupportedSignalError(sig)
 
         def handle_closed_process():
             logging.warning(
@@ -105,10 +104,9 @@ class ProcessManager:
                 handle_closed_process()
                 raise
             except OSError as exc:
-                if sys.platform == "win32":
-                    if exc.winerror == 87:
-                        handle_closed_process()
-                        raise ProcessLookupError from exc
+                if getattr(exc, "winerror", None) == 87:
+                    handle_closed_process()
+                    raise ProcessLookupError from exc
                 raise
 
     def terminate(self, name: str):
@@ -120,7 +118,7 @@ class ProcessManager:
         if sys.platform == "win32":
             self.send_signal(name, signal.SIGTERM)
         else:
-            self.send_signal(name, signal.SIGKILL)
+            self.send_signal(name, signal.SIGKILL)  # pylint: disable=no-member
 
     def remove(self, name: str, force: bool = False):
         """Remove the specified named process from this manager.

dvc/proc/process.py

@@ -5,12 +5,11 @@ import shlex
 import subprocess
 from contextlib import AbstractContextManager
 from dataclasses import asdict, dataclass
-from typing import List, Optional, TextIO, Union
+from typing import BinaryIO, List, Optional, Union
 
-from funcy import cached_property
 from shortuuid import uuid
 
-from dvc.utils.fs import makedirs
+from dvc.utils.objects import cached_property
 
 from .exceptions import TimeoutExpired
 
@@ -60,8 +59,8 @@ class ManagedProcess(AbstractContextManager):
         self.wdir = wdir
         self.name = name or uuid()
         self.returncode: Optional[int] = None
-        self._stdout: Optional[TextIO] = None
-        self._stderr: Optional[TextIO] = None
+        self._stdout: Optional[BinaryIO] = None
+        self._stderr: Optional[BinaryIO] = None
 
         self._run()
 
@@ -103,7 +102,7 @@ class ManagedProcess(AbstractContextManager):
 
     def _make_wdir(self):
         if self.wdir:
-            makedirs(self.wdir, exist_ok=True)
+            os.makedirs(self.wdir, exist_ok=True)
 
     def _dump(self):
         self._make_wdir()
@@ -112,13 +111,13 @@ class ManagedProcess(AbstractContextManager):
         with open(self.pidfile_path, "w", encoding="utf-8") as fobj:
             fobj.write(str(self.pid))
 
-    def _run(self):
+    def _run(self) -> None:
         self._make_wdir()
         logger.debug(
             "Appending output to '%s'",
             self.stdout_path,
         )
-        self._stdout = open(self.stdout_path, "ab")
+        self._stdout = open(self.stdout_path, "ab")  # noqa: SIM115
         try:
             self._proc = subprocess.Popen(
                 self.args,

dvc/progress.py

@@ -46,7 +46,7 @@ class Tqdm(tqdm):
         desc=None,
         leave=False,
         bar_format=None,
-        bytes=False,  # pylint: disable=redefined-builtin
+        bytes=False,  # noqa: A002, pylint: disable=redefined-builtin
         file=None,
         total=None,
         postfix=None,
@@ -147,10 +147,16 @@ class Tqdm(tqdm):
     def format_dict(self):
         """inject `ncols_desc` to fill the display width (`ncols`)"""
         d = super().format_dict
-        ncols = d["ncols"] or 80
+        ncols: int = d["ncols"] or 80
         # assumes `bar_format` has max one of ("ncols_desc" & "ncols_info")
         ncols_left = (
-            ncols - len(self.format_meter(ncols_desc=1, ncols_info=1, **d)) + 1
+            ncols
+            - len(
+                self.format_meter(  # type: ignore[call-arg]
+                    ncols_desc=1, ncols_info=1, **d
+                )
+            )
+            + 1
         )
         ncols_left = max(ncols_left, 0)
         if ncols_left:

dvc/prompt.py

@@ -3,11 +3,12 @@
 import logging
 import sys
 from getpass import getpass
+from typing import Collection, Optional
 
 logger = logging.getLogger(__name__)
 
 
-def ask(prompt, limited_to=None):
+def ask(prompt: str, limited_to: Optional[Collection[str]] = None):
     if not sys.stdout.isatty():
         return None
 
@@ -29,7 +30,7 @@ def ask(prompt, limited_to=None):
         )
 
 
-def confirm(statement):
+def confirm(statement: str) -> bool:
     """Ask the user for confirmation about the specified statement.
 
     Args:
@@ -43,7 +44,7 @@ def confirm(statement):
     return answer and answer.startswith("y")
 
 
-def password(statement):
+def password(statement: str) -> str:
     """Ask the user for a password.
 
     Args:

dvc/render/converter/__init__.py

@@ -5,7 +5,7 @@ class Converter:
     def __init__(
         self,
         plot_id: str,
-        data: Dict[str, Any] = None,
+        data: Optional[Dict[str, Any]] = None,
         properties: Optional[Dict] = None,
     ):
         self.plot_id = plot_id

dvc/render/converter/vega.py

@@ -1,5 +1,5 @@
 from copy import deepcopy
-from typing import Any, Dict, Iterable, List, Tuple, Union
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
 from funcy import first, last
 
@@ -59,13 +59,12 @@ def _verify_field(file2datapoints: Dict[str, List], filename: str, field: str):
     return
 
 
-def _get_x(properties: Dict, file2datapoints: Dict[str, List[Dict]]):
+def _get_xs(properties: Dict, file2datapoints: Dict[str, List[Dict]]):
     x = properties.get("x", None)
     if x is not None and isinstance(x, dict):
-        filename, field = first(_file_field(x))
-        _verify_field(file2datapoints, filename, field)
-        return filename, field
-    return None
+        for filename, field in _file_field(x):
+            _verify_field(file2datapoints, filename, field)
+            yield filename, field
 
 
 def _get_ys(properties, file2datapoints: Dict[str, List[Dict]]):
@@ -107,7 +106,10 @@ class VegaConverter(Converter):
     """
 
     def __init__(
-        self, plot_id: str, data: Dict = None, properties: Dict = None
+        self,
+        plot_id: str,
+        data: Optional[Dict] = None,
+        properties: Optional[Dict] = None,
     ):
         super().__init__(plot_id, data, properties)
         self.plot_id = plot_id
@@ -127,25 +129,30 @@ class VegaConverter(Converter):
                     break
 
     def _infer_x_y(self):
-        def _infer_files(from_name, to_name):
-            from_value = self.properties.get(from_name, None)
-            to_value = self.properties.get(to_name, None)
+        x = self.properties.get("x", None)
+        y = self.properties.get("y", None)
 
-            if isinstance(to_value, str):
-                self.inferred_properties[to_name] = {}
-                if isinstance(from_value, dict):
-                    for file in from_value.keys():
-                        self.inferred_properties[to_name][file] = to_value
-                else:
-                    self.inferred_properties[to_name][self.plot_id] = to_value
+        # Infer x.
+        if isinstance(x, str):
+            self.inferred_properties["x"] = {}
+            # If multiple y files, duplicate x for each file.
+            if isinstance(y, dict):
+                for file, fields in y.items():
+                    # Duplicate x for each y.
+                    if isinstance(fields, list):
+                        self.inferred_properties["x"][file] = [x] * len(fields)
+                    else:
+                        self.inferred_properties["x"][file] = x
+            # Otherwise use plot ID as file.
             else:
-                self.inferred_properties[to_name] = to_value
+                self.inferred_properties["x"][self.plot_id] = x
 
-        _infer_files("y", "x")
-        _infer_files("x", "y")
-
-        if self.inferred_properties.get("y", None) is None:
+        # Infer y.
+        if y is None:
             self._infer_y_from_data()
+        # If y files not provided, use plot ID as file.
+        elif not isinstance(y, dict):
+            self.inferred_properties["y"] = {self.plot_id: y}
 
     def _find_datapoints(self):
         result = {}
@@ -157,56 +164,64 @@ class VegaConverter(Converter):
     @staticmethod
     def infer_y_label(properties):
         y_label = properties.get("y_label", None)
-        if y_label is None:
-            y = properties.get("y", None)
-            if isinstance(y, dict):
-                fields = {field for _, field in _file_field(y)}
-                if len(fields) == 1:
-                    y_label = first(fields)
-                else:
-                    y_label = "y"
-            elif isinstance(y, list):
-                y_label = "y"
-            elif isinstance(y, str):
-                y_label = y
+        if y_label is not None:
+            return y_label
+        y = properties.get("y", None)
+        if isinstance(y, str):
+            return y
+        if isinstance(y, list):
+            return "y"
+        if not isinstance(y, dict):
+            return
 
-        return y_label
+        fields = {field for _, field in _file_field(y)}
+        if len(fields) == 1:
+            return first(fields)
+        return "y"
 
     @staticmethod
     def infer_x_label(properties):
         x_label = properties.get("x_label", None)
-        if x_label is None:
-            x = properties.get("x", None)
-            if isinstance(x, dict):
-                fields = {field for _, field in _file_field(x)}
-                if len(fields) == 1:
-                    x_label = first(fields)
-                else:
-                    x_label = "x"
-            else:
-                x_label = INDEX_FIELD
-        return x_label
+        if x_label is not None:
+            return x_label
 
-    def flat_datapoints(self, revision):
+        x = properties.get("x", None)
+        if not isinstance(x, dict):
+            return INDEX_FIELD
+
+        fields = {field for _, field in _file_field(x)}
+        if len(fields) == 1:
+            return first(fields)
+        return "x"
+
+    def flat_datapoints(self, revision):  # noqa: C901
 
         file2datapoints, properties = self.convert()
 
         props_update = {}
 
-        x = _get_x(properties, file2datapoints)
+        xs = list(_get_xs(properties, file2datapoints))
 
         # assign "step" if no x provided
-        if not x:
+        if not xs:
             x_file, x_field = (
                 None,
                 INDEX_FIELD,
             )
         else:
-            x_file, x_field = x
+            x_file, x_field = xs[0]
         props_update["x"] = x_field
 
         ys = list(_get_ys(properties, file2datapoints))
 
+        num_xs = len(xs)
+        num_ys = len(ys)
+        if num_xs > 1 and num_xs != num_ys:
+            raise DvcException(
+                f"Cannot have different number of x and y data sources. Found "
+                f"{num_xs} x and {num_ys} y data sources."
+            )
+
         all_datapoints = []
         all_y_fields = {y_field for _, y_field in ys}
 
@@ -216,7 +231,9 @@ class VegaConverter(Converter):
         else:
             props_update["y"] = first(all_y_fields)
 
-        for y_file, y_field in ys:
+        for i, (y_file, y_field) in enumerate(ys):
+            if num_xs > 1:
+                x_file, x_field = xs[i]
             datapoints = deepcopy(file2datapoints.get(y_file, []))
 
             if props_update.get("y", None) == "dvc_inferred_y_value":
@@ -237,9 +254,9 @@ class VegaConverter(Converter):
                         source_datapoints=x_datapoints,
                     )
                 except IndexError:
-                    raise DvcException(
+                    raise DvcException(  # noqa: B904
                         f"Cannot join '{x_field}' from '{x_file}' and "
-                        "'{y_field}' from '{y_file}'. "
+                        f"'{y_field}' from '{y_file}'. "
                         "They have to have same length."
                     )
 
@@ -290,8 +307,8 @@ class VegaConverter(Converter):
 def _update_from_field(
     target_datapoints: List[Dict],
     field: str,
-    source_datapoints: List[Dict] = None,
-    source_field: str = None,
+    source_datapoints: Optional[List[Dict]] = None,
+    source_field: Optional[str] = None,
 ):
     if source_datapoints is None:
         source_datapoints = target_datapoints

dvc/repo/__init__.py

@@ -3,9 +3,16 @@ import os
 from collections import defaultdict
 from contextlib import contextmanager
 from functools import wraps
-from typing import TYPE_CHECKING, Callable, Optional
-
-from funcy import cached_property
+from typing import (
+    TYPE_CHECKING,
+    Callable,
+    Iterable,
+    List,
+    Optional,
+    Tuple,
+    Union,
+    cast,
+)
 
 from dvc.exceptions import FileMissingError
 from dvc.exceptions import IsADirectoryError as DvcIsADirectoryError
@@ -13,11 +20,22 @@ from dvc.exceptions import NotDvcRepoError, OutputNotFoundError
 from dvc.ignore import DvcIgnoreFilter
 from dvc.utils import env2bool
 from dvc.utils.fs import path_isin
+from dvc.utils.objects import cached_property
 
 if TYPE_CHECKING:
-    from dvc.fs import FileSystem
-    from dvc.repo.scm_context import SCMContext
-    from dvc.scm import Base
+    from dvc.fs import FileSystem, GitFileSystem
+    from dvc.fs.data import DataFileSystem
+    from dvc.fs.dvc import DVCFileSystem
+    from dvc.lock import LockBase
+    from dvc.machine import MachineManager
+    from dvc.scm import Base, Git, NoSCM
+    from dvc.stage import Stage
+    from dvc.types import DictStrAny
+    from dvc_data.hashfile.state import StateBase
+
+    from .experiments import Experiments
+    from .index import Index
+    from .scm_context import SCMContext
 
 logger = logging.getLogger(__name__)
 
@@ -25,7 +43,7 @@ logger = logging.getLogger(__name__)
 @contextmanager
 def lock_repo(repo: "Repo"):
     # pylint: disable=protected-access
-    depth = repo._lock_depth
+    depth: int = repo._lock_depth
     repo._lock_depth += 1
 
     try:
@@ -86,16 +104,16 @@ class Repo:
 
     def _get_repo_dirs(
         self,
-        root_dir: str = None,
-        fs: "FileSystem" = None,
+        root_dir: Optional[str] = None,
+        fs: Optional["FileSystem"] = None,
         uninitialized: bool = False,
-        scm: "Base" = None,
-    ):
+        scm: Optional[Union["Git", "NoSCM"]] = None,
+    ) -> Tuple[str, Optional[str], Optional[str]]:
         from dvc.fs import localfs
         from dvc.scm import SCM, SCMError
 
-        dvc_dir = None
-        tmp_dir = None
+        dvc_dir: Optional[str] = None
+        tmp_dir: Optional[str] = None
         try:
             root_dir = self.find_root(root_dir, fs)
             fs = fs or localfs
@@ -117,7 +135,7 @@ class Repo:
         assert root_dir
         return root_dir, dvc_dir, tmp_dir
 
-    def _get_database_dir(self, db_name):
+    def _get_database_dir(self, db_name: str) -> Optional[str]:
         # NOTE: by default, store SQLite-based remote indexes and state's
         # `links` and `md5s` caches in the repository itself to avoid any
         # possible state corruption in 'shared cache dir' scenario, but allow
@@ -145,15 +163,15 @@ class Repo:
 
     def __init__(
         self,
-        root_dir=None,
-        fs=None,
-        rev=None,
-        subrepos=False,
-        uninitialized=False,
-        config=None,
-        url=None,
-        repo_factory=None,
-        scm=None,
+        root_dir: Optional[str] = None,
+        fs: Optional["FileSystem"] = None,
+        rev: Optional[str] = None,
+        subrepos: bool = False,
+        uninitialized: bool = False,
+        config: Optional["DictStrAny"] = None,
+        url: Optional[str] = None,
+        repo_factory: Optional[Callable] = None,
+        scm: Optional["Base"] = None,
     ):
         from dvc.config import Config
         from dvc.data_cloud import DataCloud
@@ -178,6 +196,9 @@ class Repo:
             root_dir = "/"
             self._fs = GitFileSystem(scm=self._scm, rev=rev)
 
+        self.root_dir: str
+        self.dvc_dir: Optional[str]
+        self.tmp_dir: Optional[str]
         self.root_dir, self.dvc_dir, self.tmp_dir = self._get_repo_dirs(
             root_dir=root_dir,
             fs=self.fs,
@@ -185,22 +206,25 @@ class Repo:
             scm=scm,
         )
 
-        self.config = Config(self.dvc_dir, fs=self.fs, config=config)
+        self.config: Config = Config(self.dvc_dir, fs=self.fs, config=config)
         self._uninitialized = uninitialized
 
         # used by DVCFileSystem to determine if it should traverse subrepos
         self.subrepos = subrepos
 
         self.cloud = DataCloud(self)
-        self.stage = StageLoad(self)
+        self.stage: "StageLoad" = StageLoad(self)
 
+        self.lock: "LockBase"
+        self.odb: ODBManager
+        self.state: "StateBase"
         if isinstance(self.fs, GitFileSystem) or not self.dvc_dir:
             self.lock = LockNoop()
             self.state = StateNoop()
             self.odb = ODBManager(self)
             self.tmp_dir = None
         else:
-            self.fs.makedirs(self.tmp_dir, exist_ok=True)
+            self.fs.makedirs(cast(str, self.tmp_dir), exist_ok=True)
 
             if isinstance(self.fs, LocalFileSystem):
                 self.lock = make_lock(
@@ -223,23 +247,32 @@ class Repo:
 
             self._ignore()
 
-        self.metrics = Metrics(self)
-        self.plots = Plots(self)
-        self.params = Params(self)
+        self.metrics: Metrics = Metrics(self)
+        self.plots: Plots = Plots(self)
+        self.params: Params = Params(self)
 
         self.stage_collection_error_handler: Optional[
             Callable[[str, Exception], None]
         ] = None
-        self._lock_depth = 0
+        self._lock_depth: int = 0
 
     def __str__(self):
         return self.url or self.root_dir
 
     @cached_property
-    def index(self):
+    def index(self) -> "Index":
         from dvc.repo.index import Index
 
-        return Index(self)
+        return Index.from_repo(self)
+
+    def check_graph(
+        self, stages: Iterable["Stage"], callback: Optional[Callable] = None
+    ) -> None:
+        if not getattr(self, "_skip_graph_checks", False):
+            new = self.index.update(stages)
+            if callable(callback):
+                callback()
+            new.check_graph()
 
     @staticmethod
     def open(url, *args, **kwargs):
@@ -257,7 +290,7 @@ class Repo:
         return external_repo(url, *args, **kwargs)
 
     @cached_property
-    def scm(self):
+    def scm(self) -> Union["Git", "NoSCM"]:
         from dvc.scm import SCM, SCMError
 
         if self._scm:
@@ -281,7 +314,6 @@ class Repo:
 
     @cached_property
     def dvcignore(self) -> DvcIgnoreFilter:
-
         return DvcIgnoreFilter(self.fs, self.root_dir)
 
     def get_rev(self):
@@ -293,16 +325,16 @@ class Repo:
 
             with map_scm_exception():
                 return self.scm.get_rev()
-        return self.fs.rev
+        return cast("GitFileSystem", self.fs).rev
 
     @cached_property
-    def experiments(self):
+    def experiments(self) -> "Experiments":
         from dvc.repo.experiments import Experiments
 
         return Experiments(self)
 
     @cached_property
-    def machine(self):
+    def machine(self) -> Optional["MachineManager"]:
         from dvc.machine import MachineManager
 
         if self.tmp_dir and (
@@ -356,12 +388,14 @@ class Repo:
         raise NotDvcRepoError(msg)
 
     @classmethod
-    def find_dvc_dir(cls, root=None):
+    def find_dvc_dir(cls, root=None) -> str:
         root_dir = cls.find_root(root)
         return os.path.join(root_dir, cls.DVC_DIR)
 
     @staticmethod
-    def init(root_dir=os.curdir, no_scm=False, force=False, subdir=False):
+    def init(
+        root_dir=os.curdir, no_scm=False, force=False, subdir=False
+    ) -> "Repo":
         from dvc.repo.init import init
 
         return init(
@@ -372,10 +406,11 @@ class Repo:
         return self.odb.repo.unprotect(target)
 
     def _ignore(self):
-        flist = [self.config.files["local"], self.tmp_dir]
-
+        flist = [self.config.files["local"]]
+        if tmp_dir := self.tmp_dir:
+            flist.append(tmp_dir)
         if path_isin(self.odb.repo.path, self.root_dir):
-            flist += [self.odb.repo.path]
+            flist.append(self.odb.repo.path)
 
         for file in flist:
             self.scm_context.ignore(file)
@@ -449,7 +484,9 @@ class Repo:
         return used
 
     @property
-    def stages(self):  # obsolete, only for backward-compatibility
+    def stages(
+        self,
+    ) -> List["Stage"]:  # obsolete, only for backward-compatibility
         return self.index.stages
 
     def find_outs_by_path(self, path, outs=None, recursive=False, strict=True):
@@ -484,13 +521,13 @@ class Repo:
         return self.DVC_DIR in path_parts
 
     @cached_property
-    def datafs(self):
+    def datafs(self) -> "DataFileSystem":
         from dvc.fs.data import DataFileSystem
 
         return DataFileSystem(index=self.index.data["repo"])
 
     @cached_property
-    def dvcfs(self):
+    def dvcfs(self) -> "DVCFileSystem":
         from dvc.fs.dvc import DVCFileSystem
 
         return DVCFileSystem(
@@ -507,6 +544,7 @@ class Repo:
         from dvc.fs.data import DataFileSystem
         from dvc.fs.dvc import DVCFileSystem
 
+        fs: Union["FileSystem", DataFileSystem, DVCFileSystem]
         if os.path.isabs(path):
             fs = DataFileSystem(index=self.index.data["local"])
             fs_path = path
@@ -536,11 +574,16 @@ class Repo:
         self.scm.close()
         self.state.close()
 
+    def _reset_cached_indecies(self):
+        self.__dict__.pop("index", None)
+        self.__dict__.pop("dvcignore", None)
+        self.__dict__.pop("dvcfs", None)
+        self.__dict__.pop("datafs", None)
+
     def _reset(self):
         self.state.close()
         self.scm._reset()  # pylint: disable=protected-access
-        self.__dict__.pop("index", None)
-        self.__dict__.pop("dvcignore", None)
+        self._reset_cached_indecies()
 
     def __enter__(self):
         return self

dvc/repo/add.py

@@ -2,22 +2,27 @@ import logging
 import os
 from contextlib import contextmanager
 from itertools import tee
-from typing import TYPE_CHECKING, Any, Iterator, List
+from typing import TYPE_CHECKING, Any, Iterator, List, Optional
 
 import colorama
 
-from dvc.ui import ui
-
-from ..exceptions import (
+from dvc.exceptions import (
     CacheLinkError,
     InvalidArgumentError,
     OutputDuplicationError,
     OverlappingOutputPathsError,
     RecursiveAddingWhileUsingFilename,
 )
-from ..repo.scm_context import scm_context
-from ..utils import LARGE_DIR_SIZE, glob_targets, resolve_output, resolve_paths
-from ..utils.collections import ensure_list, validate
+from dvc.repo.scm_context import scm_context
+from dvc.ui import ui
+from dvc.utils import (
+    LARGE_DIR_SIZE,
+    glob_targets,
+    resolve_output,
+    resolve_paths,
+)
+from dvc.utils.collections import ensure_list, validate
+
 from . import locked
 
 if TYPE_CHECKING:
@@ -91,9 +96,11 @@ def translate_graph_error(stages: Stages) -> Iterator[None]:
             parent=exc.parent.fs_path,
             parent_stage=exc.parent.stage.addressing,
         )
-        raise OverlappingOutputPathsError(exc.parent, exc.overlapping_out, msg)
+        raise OverlappingOutputPathsError(  # noqa: B904
+            exc.parent, exc.overlapping_out, msg
+        )
     except OutputDuplicationError as exc:
-        raise OutputDuplicationError(
+        raise OutputDuplicationError(  # noqa: B904
             exc.output, list(set(exc.stages) - set(stages))
         )
 
@@ -153,7 +160,7 @@ def add(  # noqa: C901
     targets: "TargetType",
     recursive: bool = False,
     no_commit: bool = False,
-    fname: str = None,
+    fname: Optional[str] = None,
     to_remote: bool = False,
     **kwargs: Any,
 ):
@@ -173,9 +180,9 @@ def add(  # noqa: C901
     with translate_graph_error(stages), ui.status(msg) as status:
         # remove existing stages that are to-be replaced with these
         # new stages for the graph checks.
-        new_index = repo.index.update(stages)
-        status.update("Checking graph")
-        new_index.check_graph()
+        repo.check_graph(
+            stages=stages, callback=lambda: status.update("Checking graph")
+        )
 
     odb = None
     if to_remote:
@@ -244,7 +251,7 @@ def _find_all_targets(
 def create_stages(
     repo: "Repo",
     targets: Iterator[str],
-    fname: str = None,
+    fname: Optional[str] = None,
     transfer: bool = False,
     external: bool = False,
     **kwargs: Any,

dvc/repo/brancher.py

@@ -77,7 +77,7 @@ def brancher(  # noqa: E302
         from dvc.fs import GitFileSystem
 
         for sha, names in found_revs.items():
-            self.__dict__.pop("index", None)
+            self._reset_cached_indecies()  # pylint: disable=W0212
             self.fs = GitFileSystem(scm=scm, rev=sha)
             self.root_dir = self.fs.path.join("/", *repo_root_parts)
 

dvc/repo/checkout.py

@@ -1,6 +1,6 @@
 import logging
 import os
-from typing import TYPE_CHECKING, Set
+from typing import TYPE_CHECKING, Dict, List, Set
 
 from dvc.exceptions import (
     CheckoutError,
@@ -81,7 +81,12 @@ def checkout(
 ):
     from dvc.fs.callbacks import Callback
 
-    stats = {"added": [], "deleted": [], "modified": [], "failed": []}
+    stats: Dict[str, List[str]] = {
+        "added": [],
+        "deleted": [],
+        "modified": [],
+        "failed": [],
+    }
     if not targets:
         targets = [None]
         stats["deleted"] = _remove_unused_links(self)

dvc/repo/collect.py

@@ -1,7 +1,5 @@
 import logging
-from typing import TYPE_CHECKING, Callable, Iterable, List, Tuple
-
-from dvc.types import AnyPath
+from typing import TYPE_CHECKING, Callable, Iterable, List, Optional, Tuple
 
 if TYPE_CHECKING:
     from dvc.output import Output
@@ -12,12 +10,11 @@ logger = logging.getLogger(__name__)
 
 FilterFn = Callable[["Output"], bool]
 Outputs = List["Output"]
-AnyPaths = List[AnyPath]
 StrPaths = List[str]
 
 
 def _collect_outs(
-    repo: "Repo", output_filter: FilterFn = None, deps: bool = False
+    repo: "Repo", output_filter: Optional[FilterFn] = None, deps: bool = False
 ) -> Outputs:
     index = repo.index
     index.check_graph()  # ensure graph is correct
@@ -28,7 +25,6 @@ def _collect_paths(
     repo: "Repo",
     targets: Iterable[str],
     recursive: bool = False,
-    rev: str = None,
 ) -> StrPaths:
     from dvc.fs.dvc import DVCFileSystem
 
@@ -39,13 +35,6 @@ def _collect_paths(
     for fs_path in fs_paths:
         if recursive and fs.isdir(fs_path):
             target_paths.extend(fs.find(fs_path))
-
-        rel = fs.path.relpath(fs_path)
-        if not fs.exists(fs_path):
-            if rev == "workspace" or rev == "":
-                logger.warning("'%s' was not found in current workspace.", rel)
-            else:
-                logger.warning("'%s' was not found at: '%s'.", rel, rev)
         target_paths.append(fs_path)
 
     return target_paths
@@ -71,9 +60,8 @@ def _filter_outs(
 def collect(
     repo: "Repo",
     deps: bool = False,
-    targets: Iterable[str] = None,
-    output_filter: FilterFn = None,
-    rev: str = None,
+    targets: Optional[Iterable[str]] = None,
+    output_filter: Optional[FilterFn] = None,
     recursive: bool = False,
     duplicates: bool = False,
 ) -> Tuple[Outputs, StrPaths]:
@@ -85,6 +73,6 @@ def collect(
         fs_paths: StrPaths = []
         return outs, fs_paths
 
-    target_paths = _collect_paths(repo, targets, recursive=recursive, rev=rev)
+    target_paths = _collect_paths(repo, targets, recursive=recursive)
 
     return _filter_outs(outs, target_paths, duplicates=duplicates)

dvc/repo/commit.py

@@ -45,8 +45,6 @@ def commit(
     allow_missing=False,
     data_only=False,
 ):
-    from dvc.dvcfile import Dvcfile
-
     stages_info = [
         info
         for info in self.stage.collect_granular(
@@ -66,6 +64,5 @@ def commit(
         stage.commit(
             filter_info=stage_info.filter_info, allow_missing=allow_missing
         )
-
-        Dvcfile(self, stage.path).dump(stage, update_pipeline=False)
+        stage.dump(update_pipeline=False)
     return [s.stage for s in stages_info]

dvc/repo/data.py

@@ -13,14 +13,12 @@ from typing import (
     Tuple,
     TypedDict,
     Union,
-    cast,
 )
 
 from dvc.fs.git import GitFileSystem
 from dvc.ui import ui
 
 if TYPE_CHECKING:
-    from dvc.output import Output
     from dvc.repo import Repo
     from dvc.scm import Git, NoSCM
     from dvc_data.hashfile.db import HashFileDB
@@ -175,7 +173,6 @@ def _diff_index_to_wtree(repo: "Repo", **kwargs: Any) -> Dict[str, List[str]]:
 
     unstaged_diff = defaultdict(list)
     for out in repo.index.outs:
-        out = cast("Output", out)
         if not out.use_cache:
             continue
 
@@ -217,13 +214,12 @@ def _diff_head_to_index(
     repo: "Repo", head: str = "HEAD", **kwargs: Any
 ) -> Dict[str, List[str]]:
     # we need to store objects from index and the HEAD to diff later
-    objs: Dict[str, Dict[str, Tuple["HashFile", "HashInfo"]]]
+    objs: Dict[str, Dict[str, Tuple[Optional["HashFile"], "HashInfo"]]]
     objs = defaultdict(dict)
 
     staged_diff = defaultdict(list)
     for rev in repo.brancher(revs=[head]):
         for out in repo.index.outs:
-            out = cast("Output", out)
             if not out.use_cache:
                 continue
 
@@ -313,7 +309,7 @@ def status(repo: "Repo", untracked_files: str = "no", **kwargs: Any) -> Status:
 
 def ls(
     repo: "Repo",
-    targets: List[Optional[str]] = None,
+    targets: Optional[List[Optional[str]]] = None,
     recursive: bool = False,
 ) -> Iterator[Dict[str, Any]]:
     targets = targets or [None]

dvc/repo/diff.py

@@ -1,17 +1,40 @@
 import logging
 import os
 from collections import defaultdict
-from typing import Dict, List
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Iterable,
+    Iterator,
+    List,
+    Optional,
+    Set,
+    Tuple,
+    cast,
+)
 
 from dvc.exceptions import PathMissingError
 from dvc.repo import locked
 from dvc.utils.collections import ensure_list
 
+if TYPE_CHECKING:
+    from dvc.fs import DVCFileSystem, FileSystem
+    from dvc.output import Output
+    from dvc.repo import Repo
+    from dvc_data.hashfile.obj import HashFile
+    from dvc_data.hashfile.tree import Tree
+
 logger = logging.getLogger(__name__)
 
 
 @locked
-def diff(self, a_rev="HEAD", b_rev=None, targets=None):
+def diff(
+    self,
+    a_rev: str = "HEAD",
+    b_rev: Optional[str] = None,
+    targets: Optional[List[str]] = None,
+):
     """
     By default, it compares the workspace with the last commit's fs.
 
@@ -26,8 +49,7 @@ def diff(self, a_rev="HEAD", b_rev=None, targets=None):
     from dvc.fs.dvc import DVCFileSystem
 
     dvcfs = DVCFileSystem(repo=self)
-    targets = ensure_list(targets)
-    targets = [dvcfs.from_os_path(target) for target in targets]
+    targets = [dvcfs.from_os_path(target) for target in ensure_list(targets)]
 
     b_rev = b_rev if b_rev else "workspace"
     results = {}
@@ -98,7 +120,9 @@ def diff(self, a_rev="HEAD", b_rev=None, targets=None):
     return ret if any(ret.values()) else {}
 
 
-def _paths_checksums(repo, targets):
+def _paths_checksums(
+    repo: "Repo", targets: Optional[List[str]]
+) -> Dict[str, Optional[str]]:
     """
     A dictionary of checksums addressed by relpaths collected from
     the current fs outputs.
@@ -113,18 +137,20 @@ def _paths_checksums(repo, targets):
     return dict(_output_paths(repo, targets))
 
 
-def _output_paths(repo, targets):
+def _output_paths(  # noqa: C901
+    repo: "Repo", targets: Optional[List[str]]
+) -> Iterable[Tuple[str, Optional[str]]]:
     from dvc.fs import LocalFileSystem
     from dvc_data.hashfile.build import build
 
     on_working_fs = isinstance(repo.fs, LocalFileSystem)
 
-    def _exists(output):
+    def _exists(output: "Output") -> bool:
         if on_working_fs:
             return output.exists
         return True
 
-    def _to_path(output):
+    def _to_path(output: "Output") -> str:
         relparts = output.fs.path.relparts(output.fs_path)
         base = os.path.join(*relparts)
         if output.is_dir_checksum:
@@ -132,45 +158,57 @@ def _output_paths(repo, targets):
         return base
 
     for output in repo.index.outs:
-        if _exists(output):
-            yield_output = targets is None or any(
-                output.fs.path.isin_or_eq(output.fs_path, target)
-                for target in targets
+        if not _exists(output):
+            continue
+
+        yield_output = targets is None or any(
+            output.fs.path.isin_or_eq(output.fs_path, target)
+            for target in targets
+        )
+
+        obj: Optional["HashFile"]
+        if on_working_fs:
+            _, _, obj = build(
+                repo.odb.repo,
+                output.fs_path,
+                repo.odb.repo.fs,
+                "md5",
+                dry_run=True,
+                ignore=output.dvcignore,
             )
+            hash_info = obj.hash_info
+        else:
+            hash_info = output.hash_info
+            obj = output.get_obj()
 
-            if on_working_fs:
-                _, _, obj = build(
-                    repo.odb.repo,
-                    output.fs_path,
-                    repo.odb.repo.fs,
-                    "md5",
-                    dry_run=True,
-                    ignore=output.dvcignore,
-                )
-                hash_info = obj.hash_info
-            else:
-                hash_info = output.hash_info
-                obj = output.get_obj()
+        if yield_output:
+            yield _to_path(output), hash_info.value
 
-            if yield_output:
-                yield _to_path(output), hash_info.value
+        if not obj:
+            continue
 
-            if not obj:
-                continue
-
-            if output.is_dir_checksum and (
-                yield_output
-                or any(
+        if output.is_dir_checksum and (
+            yield_output
+            or (
+                isinstance(targets, list)
+                and any(
                     output.fs.path.isin(target, output.fs_path)
                     for target in targets
                 )
-            ):
-                yield from _dir_output_paths(
-                    output.fs, output.fs_path, obj, targets
-                )
+            )
+        ):
+            obj = cast("Tree", obj)
+            yield from _dir_output_paths(
+                output.fs, output.fs_path, obj, targets
+            )
 
 
-def _dir_output_paths(fs, fs_path, obj, targets=None):
+def _dir_output_paths(
+    fs: "FileSystem",
+    fs_path: str,
+    obj: "Tree",
+    targets: Optional[Iterable[str]] = None,
+) -> Iterator[Tuple[str, str]]:
     base = os.path.join(*fs.path.relparts(fs_path))
     for key, _, oid in obj:
         fname = fs.path.join(fs_path, *key)
@@ -181,7 +219,7 @@ def _dir_output_paths(fs, fs_path, obj, targets=None):
             yield os.path.join(base, *key), oid.value
 
 
-def _filter_missing(dvcfs, paths):
+def _filter_missing(dvcfs: "DVCFileSystem", paths: Set[str]) -> Iterator[str]:
     for path in paths:
         fs_path = dvcfs.from_os_path(path)
         try:
@@ -198,9 +236,14 @@ def _filter_missing(dvcfs, paths):
             pass
 
 
-def _targets_to_paths(dvcfs, targets):
-    paths = []
-    missing = []
+def _targets_to_paths(
+    dvcfs: "DVCFileSystem", targets: List[str]
+) -> Tuple[List[str], List[str]]:
+    paths: List[str] = []
+    missing: List[str] = []
+
+    if not targets:
+        return paths, missing
 
     for target in targets:
         if dvcfs.exists(target):
@@ -211,7 +254,12 @@ def _targets_to_paths(dvcfs, targets):
     return paths, missing
 
 
-def _calculate_renamed(new, old, added, deleted):
+def _calculate_renamed(
+    new: Dict[str, Any],
+    old: Dict[str, Any],
+    added: Iterable[str],
+    deleted: Iterable[str],
+) -> List[Dict[str, Dict[str, str]]]:
     old_inverted: Dict[str, List[str]] = defaultdict(list)
     # It is needed to be dict of lists to cover cases
     # when repo has paths with same hash

dvc/repo/experiments/__init__.py

@@ -4,11 +4,12 @@ import re
 import time
 from typing import Dict, Iterable, Optional
 
-from funcy import cached_property, chain, first
+from funcy import chain, first
 
 from dvc.exceptions import DvcException
 from dvc.ui import ui
 from dvc.utils import relpath
+from dvc.utils.objects import cached_property
 
 from .exceptions import (
     BaselineMismatchError,
@@ -32,7 +33,7 @@ from .refs import (
     ExpRefInfo,
 )
 from .stash import ExpStashEntry
-from .utils import exp_refs_by_rev, scm_locked, unlocked_repo
+from .utils import check_ref_format, exp_refs_by_rev, unlocked_repo
 
 logger = logging.getLogger(__name__)
 
@@ -50,29 +51,23 @@ class Experiments:
     )
 
     def __init__(self, repo):
-        from dvc.lock import make_lock
         from dvc.scm import NoSCMError
 
         if repo.config["core"].get("no_scm", False):
             raise NoSCMError
 
         self.repo = repo
-        self.scm_lock = make_lock(
-            os.path.join(self.repo.tmp_dir, "exp_scm_lock"),
-            tmp_dir=self.repo.tmp_dir,
-            hardlink_lock=repo.config["core"].get("hardlink_lock", False),
-        )
 
     @property
     def scm(self):
         return self.repo.scm
 
     @cached_property
-    def dvc_dir(self):
+    def dvc_dir(self) -> str:
         return relpath(self.repo.dvc_dir, self.repo.scm.root_dir)
 
     @cached_property
-    def args_file(self):
+    def args_file(self) -> str:
         return os.path.join(self.repo.tmp_dir, BaseExecutor.PACKED_ARGS_FILE)
 
     @cached_property
@@ -166,7 +161,7 @@ class Experiments:
             return last_applied
         return None
 
-    def reproduce_celery(
+    def reproduce_celery(  # noqa: C901
         self, entries: Optional[Iterable[QueueEntry]] = None, **kwargs
     ) -> Dict[str, str]:
         results: Dict[str, str] = {}
@@ -186,8 +181,7 @@ class Experiments:
         if not entries:
             return results
 
-        # TODO: re-enable --jobs concurrency
-        self.celery_queue.spawn_worker()
+        self.celery_queue.start_workers(count=kwargs.get("jobs", 1))
         failed = []
         try:
             ui.write(
@@ -242,18 +236,6 @@ class Experiments:
             "\tdvc exp branch <exp> <branch>\n"
         )
 
-    def _validate_new_ref(self, exp_ref: ExpRefInfo):
-        from .utils import check_ref_format
-
-        if not exp_ref.name:
-            return
-
-        check_ref_format(self.scm, exp_ref)
-
-        if self.scm.get_ref(str(exp_ref)):
-            raise ExperimentExistsError(exp_ref.name)
-
-    @scm_locked
     def new(
         self,
         queue: BaseStashQueue,
@@ -272,13 +254,13 @@ class Experiments:
 
         name = kwargs.get("name", None)
         baseline_sha = kwargs.get("baseline_rev") or self.repo.scm.get_rev()
-        exp_ref = ExpRefInfo(baseline_sha=baseline_sha, name=name)
 
-        try:
-            self._validate_new_ref(exp_ref)
-        except ExperimentExistsError as err:
-            if not (kwargs.get("force", False) or kwargs.get("reset", False)):
-                raise err
+        if name:
+            exp_ref = ExpRefInfo(baseline_sha=baseline_sha, name=name)
+            check_ref_format(self.scm, exp_ref)
+            force = kwargs.get("force", False)
+            if self.scm.get_ref(str(exp_ref)) and not force:
+                raise ExperimentExistsError(exp_ref.name)
 
         return queue.put(*args, **kwargs)
 
@@ -428,7 +410,7 @@ class Experiments:
     def get_exact_name(self, revs: Iterable[str]) -> Dict[str, Optional[str]]:
         """Returns preferred name for the specified revision.
 
-        Prefers tags, branches (heads), experiments in that orer.
+        Prefers tags, branches (heads), experiments in that order.
         """
         result: Dict[str, Optional[str]] = {}
         exclude = f"{EXEC_NAMESPACE}/*"
@@ -446,8 +428,10 @@ class Experiments:
             if not name:
                 if rev in self.stash_revs:
                     name = self.stash_revs[rev].name
-                elif rev in self.celery_queue.failed_stash.stash_revs:
-                    name = self.celery_queue.failed_stash.stash_revs[rev].name
+                else:
+                    failed_stash = self.celery_queue.failed_stash
+                    if failed_stash and rev in failed_stash.stash_revs:
+                        name = failed_stash.stash_revs[rev].name
             result[rev] = name
         return result
 
@@ -487,6 +471,11 @@ class Experiments:
 
         return run(self.repo, *args, **kwargs)
 
+    def save(self, *args, **kwargs):
+        from dvc.repo.experiments.save import save
+
+        return save(self.repo, *args, **kwargs)
+
     def gc(self, *args, **kwargs):
         from dvc.repo.experiments.gc import gc
 

dvc/repo/experiments/apply.py

@@ -16,46 +16,61 @@ from .executor.base import BaseExecutor
 from .refs import EXEC_APPLY
 
 if TYPE_CHECKING:
-    from scmrepo import Git
+    from scmrepo.git import Git
 
     from dvc.repo import Repo
+    from dvc.repo.experiments import Experiments
 
 logger = logging.getLogger(__name__)
 
 
 @locked
 @scm_context
-def apply(repo: "Repo", rev: str, force: bool = True, **kwargs):
+def apply(repo: "Repo", rev: str, force: bool = True, **kwargs):  # noqa: C901
     from scmrepo.exceptions import SCMError as _SCMError
 
     from dvc.repo.checkout import checkout as dvc_checkout
     from dvc.scm import GitMergeError, RevError, resolve_rev
 
-    exps = repo.experiments
+    exps: "Experiments" = repo.experiments
+
+    is_stash: bool = False
 
     try:
         exp_rev = resolve_rev(repo.scm, rev)
         exps.check_baseline(exp_rev)
-    except (RevError, BaselineMismatchError) as exc:
+    except RevError as exc:
+        (
+            exp_ref_info,
+            queue_entry,
+        ) = exps.celery_queue.get_ref_and_entry_by_names(rev)[rev]
+        baseline_sha = repo.scm.get_rev()
+        if exp_ref_info:
+            if baseline_sha != exp_ref_info.baseline_sha:
+                raise InvalidExpRevError(rev)  # noqa: B904
+            exp_rev = repo.scm.get_ref(str(exp_ref_info))
+        elif queue_entry:
+            if queue_entry.baseline_rev != baseline_sha:
+                raise InvalidExpRevError(rev)
+            exp_rev = queue_entry.stash_rev
+            is_stash = True
+        else:
+            raise InvalidExpRevError(rev) from exc
+    except BaselineMismatchError as exc:
         raise InvalidExpRevError(rev) from exc
 
-    stash_rev = exp_rev in exps.stash_revs
-    if not stash_rev and not exps.get_branch_by_rev(
-        exp_rev, allow_multiple=True
-    ):
-        raise InvalidExpRevError(exp_rev)
-
     # NOTE: we don't use scmrepo's stash_workspace() here since we need
     # finer control over the merge behavior when we unstash everything
     with _apply_workspace(repo, rev, force):
         try:
             repo.scm.merge(exp_rev, commit=False, squash=True)
         except _SCMError as exc:
-            raise GitMergeError(str(exc), scm=repo.scm)
+            raise GitMergeError(str(exc), scm=repo.scm)  # noqa: B904
 
     repo.scm.reset()
 
-    if stash_rev:
+    if is_stash:
+        assert repo.tmp_dir is not None
         args_path = os.path.join(repo.tmp_dir, BaseExecutor.PACKED_ARGS_FILE)
         if os.path.exists(args_path):
             remove(args_path)

dvc/repo/experiments/branch.py

@@ -19,7 +19,7 @@ def branch(repo, exp_rev, branch_name, *args, **kwargs):
     try:
         rev = resolve_rev(repo.scm, exp_rev)
     except RevError:
-        raise InvalidArgumentError(exp_rev)
+        raise InvalidArgumentError(exp_rev)  # noqa: B904
     ref_info = None
 
     ref_infos = list(exp_refs_by_rev(repo.scm, rev))

dvc/repo/experiments/exceptions.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Collection, Iterable
+from typing import TYPE_CHECKING, Collection, Iterable, Optional
 
 from dvc.exceptions import DvcException, InvalidArgumentError
 
@@ -6,12 +6,6 @@ if TYPE_CHECKING:
     from .refs import ExpRefInfo
 
 
-class UnchangedExperimentError(DvcException):
-    def __init__(self, rev):
-        super().__init__(f"Experiment unchanged from '{rev[:7]}'.")
-        self.rev = rev
-
-
 class BaselineMismatchError(DvcException):
     def __init__(self, rev, expected):
         if hasattr(rev, "hexsha"):
@@ -25,13 +19,13 @@ class BaselineMismatchError(DvcException):
 
 
 class ExperimentExistsError(DvcException):
-    def __init__(self, name: str):
+    def __init__(self, name: str, command: str = "run"):
         msg = (
             "Experiment conflicts with existing experiment "
             f"'{name}'. To overwrite the existing experiment run:\n\n"
-            "\tdvc exp run -f ...\n\n"
-            "To run this experiment with a different name run:\n\n"
-            f"\tdvc exp run -n <new_name> ...\n"
+            f"\tdvc exp {command} -f ...\n\n"
+            f"To {command} this experiment with a different name run:\n\n"
+            f"\tdvc exp {command} -n <new_name> ...\n"
         )
         super().__init__(msg)
         self.name = name
@@ -108,7 +102,10 @@ class UnresolvedExpNamesError(InvalidArgumentError):
     NAME = "experiment name"
 
     def __init__(
-        self, unresolved_list: Collection[str], *args, git_remote: str = None
+        self,
+        unresolved_list: Collection[str],
+        *args,
+        git_remote: Optional[str] = None,
     ):
         unresolved_names = "; ".join(unresolved_list)
         if not git_remote:

dvc/repo/experiments/executor/base.py

@@ -12,6 +12,7 @@ from typing import (
     Callable,
     Dict,
     Iterable,
+    List,
     NamedTuple,
     Optional,
     Tuple,
@@ -24,28 +25,21 @@ from scmrepo.exceptions import SCMError
 
 from dvc.env import DVC_EXP_AUTO_PUSH, DVC_EXP_GIT_REMOTE
 from dvc.exceptions import DvcException
+from dvc.repo.experiments.exceptions import (
+    CheckpointExistsError,
+    ExperimentExistsError,
+)
+from dvc.repo.experiments.refs import (
+    EXEC_BASELINE,
+    EXEC_BRANCH,
+    EXEC_CHECKPOINT,
+    ExpRefInfo,
+)
 from dvc.stage.serialize import to_lockfile
 from dvc.ui import ui
 from dvc.utils import dict_sha256, env2bool, relpath
 from dvc.utils.fs import remove
 
-from ..exceptions import (
-    CheckpointExistsError,
-    ExperimentExistsError,
-    UnchangedExperimentError,
-)
-from ..refs import (
-    EXEC_BASELINE,
-    EXEC_BRANCH,
-    EXEC_CHECKPOINT,
-    EXEC_HEAD,
-    EXEC_MERGE,
-    EXEC_NAMESPACE,
-    EXPS_NAMESPACE,
-    EXPS_STASH,
-    ExpRefInfo,
-)
-
 if TYPE_CHECKING:
     from queue import Queue
 
@@ -58,10 +52,6 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
-EXEC_TMP_DIR = "exps"
-EXEC_PID_DIR = "run"
-
-
 class ExecutorResult(NamedTuple):
     exp_hash: Optional[str]
     ref_info: Optional["ExpRefInfo"]
@@ -94,9 +84,8 @@ class ExecutorInfo:
 
     @classmethod
     def from_dict(cls, d):
-        if "collected" in d:
-            if d.pop("collected"):
-                d["status"] = TaskStatus.FINISHED
+        if d.pop("collected", None):
+            d["status"] = TaskStatus.FINISHED
         return cls(**d)
 
     def asdict(self):
@@ -113,10 +102,9 @@ class ExecutorInfo:
         )
 
     def dump_json(self, filename: str):
-        from dvc.utils.fs import makedirs
         from dvc.utils.serialize import modify_json
 
-        makedirs(os.path.dirname(filename), exist_ok=True)
+        os.makedirs(os.path.dirname(filename), exist_ok=True)
         with modify_json(filename) as d:
             d.update(self.asdict())
 
@@ -172,6 +160,7 @@ class BaseExecutor(ABC):
     @abstractmethod
     def init_git(
         self,
+        repo: "Repo",
         scm: "Git",
         stash_rev: str,
         entry: "ExpStashEntry",
@@ -263,7 +252,7 @@ class BaseExecutor(ABC):
         root_dir: str,
         **kwargs,
     ) -> _T:
-        executor = cls(
+        return cls(
             root_dir=root_dir,
             dvc_dir=relpath(repo.dvc_dir, repo.scm.root_dir),
             baseline_rev=entry.baseline_rev,
@@ -272,7 +261,6 @@ class BaseExecutor(ABC):
             wdir=relpath(os.getcwd(), repo.scm.root_dir),
             **kwargs,
         )
-        return executor
 
     @staticmethod
     def hash_exp(stages: Iterable["PipelineStage"]) -> str:
@@ -299,10 +287,9 @@ class BaseExecutor(ABC):
             open_func = fs.open
             fs.makedirs(dpath)
         else:
-            from dvc.utils.fs import makedirs
 
             open_func = open
-            makedirs(dpath, exist_ok=True)
+            os.makedirs(dpath, exist_ok=True)
 
         data = {"args": args, "kwargs": kwargs}
         if extra is not None:
@@ -319,34 +306,28 @@ class BaseExecutor(ABC):
     def fetch_exps(
         self,
         dest_scm: "Git",
+        refs: List[str],
         force: bool = False,
-        on_diverged: Callable[[str, bool], None] = None,
+        on_diverged: Optional[Callable[[str, bool], None]] = None,
         **kwargs,
     ) -> Iterable[str]:
         """Fetch reproduced experiment refs into the specified SCM.
 
         Args:
             dest_scm: Destination Git instance.
+            refs: reference names to be fetched from the remotes.
             force: If True, diverged refs will be overwritten
             on_diverged: Callback in the form on_diverged(ref, is_checkpoint)
                 to be called when an experiment ref has diverged.
 
         Extra kwargs will be passed into the remote git client.
         """
-        from ..utils import iter_remote_refs
 
-        refs = []
-        has_checkpoint = False
-        for ref in iter_remote_refs(
-            dest_scm,
-            self.git_url,
-            base=EXPS_NAMESPACE,
-            **kwargs,
-        ):
-            if ref == EXEC_CHECKPOINT:
-                has_checkpoint = True
-            elif not ref.startswith(EXEC_NAMESPACE) and ref != EXPS_STASH:
-                refs.append(ref)
+        if EXEC_CHECKPOINT in refs:
+            refs.remove(EXEC_CHECKPOINT)
+            has_checkpoint = True
+        else:
+            has_checkpoint = False
 
         def on_diverged_ref(orig_ref: str, new_rev: str):
             if force:
@@ -401,7 +382,7 @@ class BaseExecutor(ABC):
         try:
             dvc.scm.validate_git_remote(git_remote)
         except InvalidRemote as exc:
-            raise InvalidRemoteSCMRepo(str(exc))
+            raise InvalidRemoteSCMRepo(str(exc))  # noqa: B904
         dvc.cloud.get_remote_odb()
 
     @classmethod
@@ -514,14 +495,13 @@ class BaseExecutor(ABC):
                     info,
                     stages,
                     exp_hash,
-                    checkpoint_reset,
                     auto_push,
                     git_remote,
                     repro_force,
                 )
-            info.result_hash = exp_hash
-            info.result_ref = ref
-            info.result_force = repro_force
+                info.result_hash = exp_hash
+                info.result_ref = ref
+                info.result_force = repro_force
 
         # ideally we would return stages here like a normal repro() call, but
         # stages is not currently picklable and cannot be returned across
@@ -535,29 +515,20 @@ class BaseExecutor(ABC):
         info,
         stages,
         exp_hash,
-        checkpoint_reset,
         auto_push,
         git_remote,
         repro_force,
     ) -> Tuple[Optional[str], Optional["ExpRefInfo"], bool]:
-        try:
-            is_checkpoint = any(stage.is_checkpoint for stage in stages)
-            if is_checkpoint and checkpoint_reset:
-                # For reset checkpoint stages, we need to force
-                # overwriting existing checkpoint refs even though
-                # repro may not have actually been run with --force
-                repro_force = True
-            cls.commit(
-                dvc.scm,
-                exp_hash,
-                exp_name=info.name,
-                force=repro_force,
-                checkpoint=is_checkpoint,
-            )
-            if auto_push:
-                cls._auto_push(dvc, dvc.scm, git_remote)
-        except UnchangedExperimentError:
-            pass
+        is_checkpoint = any(stage.is_checkpoint for stage in stages)
+        cls.commit(
+            dvc.scm,
+            exp_hash,
+            exp_name=info.name,
+            force=repro_force,
+            checkpoint=is_checkpoint,
+        )
+        if auto_push:
+            cls._auto_push(dvc, dvc.scm, git_remote)
         ref: Optional[str] = dvc.scm.get_ref(EXEC_BRANCH, follow=False)
         exp_ref: Optional["ExpRefInfo"] = (
             ExpRefInfo.from_ref(ref) if ref else None
@@ -576,10 +547,10 @@ class BaseExecutor(ABC):
 
     @classmethod
     @contextmanager
-    def _repro_dvc(
+    def _repro_dvc(  # noqa: C901
         cls,
         info: "ExecutorInfo",
-        infofile: str = None,
+        infofile: Optional[str] = None,
         log_errors: bool = True,
         **kwargs,
     ):
@@ -650,7 +621,7 @@ class BaseExecutor(ABC):
                 push_cache=push_cache,
                 run_cache=run_cache,
             )
-        except BaseException as exc:  # pylint: disable=broad-except
+        except BaseException as exc:  # noqa: BLE001, pylint: disable=W0703
             logger.warning(
                 "Something went wrong while auto pushing experiment "
                 "to the remote '%s': %s",
@@ -668,18 +639,15 @@ class BaseExecutor(ABC):
         unchanged: Iterable["PipelineStage"],
         stages: Iterable["PipelineStage"],
     ):
-        try:
-            exp_hash = cls.hash_exp(list(stages) + list(unchanged))
-            exp_rev = cls.commit(
-                scm, exp_hash, exp_name=name, force=force, checkpoint=True
-            )
+        exp_hash = cls.hash_exp(list(stages) + list(unchanged))
+        exp_rev = cls.commit(
+            scm, exp_hash, exp_name=name, force=force, checkpoint=True
+        )
 
-            if env2bool(DVC_EXP_AUTO_PUSH):
-                git_remote = os.getenv(DVC_EXP_GIT_REMOTE)
-                cls._auto_push(dvc, scm, git_remote)
-            ui.write(f"Checkpoint experiment iteration '{exp_rev[:7]}'.")
-        except UnchangedExperimentError:
-            pass
+        if env2bool(DVC_EXP_AUTO_PUSH):
+            git_remote = os.getenv(DVC_EXP_GIT_REMOTE)
+            cls._auto_push(dvc, scm, git_remote)
+        ui.write(f"Checkpoint experiment iteration '{exp_rev[:7]}'.")
 
     @classmethod
     def commit(
@@ -694,7 +662,6 @@ class BaseExecutor(ABC):
         rev = scm.get_rev()
         if not scm.is_dirty(untracked_files=False):
             logger.debug("No changes to commit")
-            raise UnchangedExperimentError(rev)
 
         check_conflict = False
         branch = scm.get_ref(EXEC_BRANCH, follow=False)
@@ -755,21 +722,12 @@ class BaseExecutor(ABC):
             dvc_logger.setLevel(level)
 
     @contextmanager
-    def set_exec_refs(
-        self, scm: "Git", stash_rev: str, entry: "ExpStashEntry"
-    ):
+    def set_temp_refs(self, scm: "Git", temp_dict: Dict[str, str]):
         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)
+            for ref, rev in temp_dict.items():
+                scm.set_ref(ref, rev)
             yield
         finally:
-            for ref in (EXEC_HEAD, EXEC_MERGE, EXEC_BASELINE):
+            for ref in temp_dict:
                 if scm.get_ref(ref):
                     scm.remove_ref(ref)

dvc/repo/experiments/executor/local.py

@@ -2,15 +2,14 @@ import logging
 import os
 from contextlib import ExitStack
 from tempfile import mkdtemp
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Optional, Union
 
-from funcy import cached_property
+from funcy import retry
 from scmrepo.exceptions import SCMError as _SCMError
+from shortuuid import uuid
 
-from dvc.scm import SCM, GitMergeError
-from dvc.utils.fs import makedirs, remove
-
-from ..refs import (
+from dvc.lock import LockError
+from dvc.repo.experiments.refs import (
     EXEC_APPLY,
     EXEC_BASELINE,
     EXEC_BRANCH,
@@ -18,16 +17,22 @@ from ..refs import (
     EXEC_HEAD,
     EXEC_MERGE,
     EXEC_NAMESPACE,
+    EXPS_TEMP,
+    ExpRefInfo,
 )
-from .base import EXEC_TMP_DIR, BaseExecutor, TaskStatus
+from dvc.repo.experiments.utils import EXEC_TMP_DIR, get_exp_rwlock
+from dvc.scm import SCM, GitMergeError
+from dvc.utils.fs import remove
+from dvc.utils.objects import cached_property
+
+from .base import BaseExecutor, TaskStatus
 
 if TYPE_CHECKING:
     from scmrepo.git import Git
 
     from dvc.repo import Repo
-
-    from ..refs import ExpRefInfo
-    from ..stash import ExpStashEntry
+    from dvc.repo.experiments.stash import ExpStashEntry
+    from dvc.scm import NoSCM
 
 logger = logging.getLogger(__name__)
 
@@ -43,7 +48,7 @@ class BaseLocalExecutor(BaseExecutor):
         return f"file://{root_dir}"
 
     @cached_property
-    def scm(self):
+    def scm(self) -> Union["Git", "NoSCM"]:
         return SCM(self.root_dir)
 
     def cleanup(self, infofile: str):
@@ -67,8 +72,10 @@ class TempDirExecutor(BaseLocalExecutor):
     QUIET = True
     DEFAULT_LOCATION = "tempdir"
 
+    @retry(180, errors=LockError, timeout=1)
     def init_git(
         self,
+        repo: "Repo",
         scm: "Git",
         stash_rev: str,
         entry: "ExpStashEntry",
@@ -77,7 +84,7 @@ class TempDirExecutor(BaseLocalExecutor):
     ):
         from dulwich.repo import Repo as DulwichRepo
 
-        from ..utils import push_refspec
+        from dvc.repo.experiments.utils import push_refspec
 
         DulwichRepo.init(os.fspath(self.root_dir))
 
@@ -85,18 +92,43 @@ class TempDirExecutor(BaseLocalExecutor):
         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)
+        temp_head = f"{EXPS_TEMP}/head-{uuid()}"
+        temp_merge = f"{EXPS_TEMP}/merge-{uuid()}"
+        temp_baseline = f"{EXPS_TEMP}/baseline-{uuid()}"
 
-        if branch:
-            push_refspec(scm, self.git_url, branch, branch)
-            self.scm.set_ref(EXEC_BRANCH, branch, symbolic=True)
-        elif self.scm.get_ref(EXEC_BRANCH):
-            self.scm.remove_ref(EXEC_BRANCH)
+        temp_ref_dict = {
+            temp_head: entry.head_rev,
+            temp_merge: stash_rev,
+            temp_baseline: entry.baseline_rev,
+        }
+        with get_exp_rwlock(
+            repo, writes=[temp_head, temp_merge, temp_baseline]
+        ), self.set_temp_refs(scm, temp_ref_dict):
+            # 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
+            refspec = [
+                (temp_head, EXEC_HEAD),
+                (temp_merge, EXEC_MERGE),
+                (temp_baseline, EXEC_BASELINE),
+            ]
+
+            if branch:
+                refspec.append((branch, branch))
+                with get_exp_rwlock(repo, reads=[branch]):
+                    push_refspec(scm, self.git_url, refspec)
+                self.scm.set_ref(EXEC_BRANCH, branch, symbolic=True)
+            else:
+                push_refspec(scm, self.git_url, refspec)
+                if self.scm.get_ref(EXEC_BRANCH):
+                    self.scm.remove_ref(EXEC_BRANCH)
+
+            if self.scm.get_ref(EXEC_CHECKPOINT):
+                self.scm.remove_ref(EXEC_CHECKPOINT)
 
-        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
@@ -106,7 +138,7 @@ class TempDirExecutor(BaseLocalExecutor):
         try:
             self.scm.merge(merge_rev, squash=True, commit=False)
         except _SCMError as exc:
-            raise GitMergeError(str(exc), scm=self.scm)
+            raise GitMergeError(str(exc), scm=self.scm)  # noqa: B904
 
     def _config(self, cache_dir):
         local_config = os.path.join(
@@ -135,8 +167,9 @@ class TempDirExecutor(BaseLocalExecutor):
         wdir: Optional[str] = None,
         **kwargs,
     ):
+        assert repo.tmp_dir
         parent_dir: str = wdir or os.path.join(repo.tmp_dir, EXEC_TMP_DIR)
-        makedirs(parent_dir, exist_ok=True)
+        os.makedirs(parent_dir, exist_ok=True)
         tmp_dir = mkdtemp(dir=parent_dir)
         try:
             executor = cls._from_stash_entry(repo, entry, tmp_dir, **kwargs)
@@ -167,8 +200,10 @@ class WorkspaceExecutor(BaseLocalExecutor):
         logger.debug("Init workspace executor in '%s'", root_dir)
         return executor
 
+    @retry(180, errors=LockError, timeout=1)
     def init_git(
         self,
+        repo: "Repo",
         scm: "Git",
         stash_rev: str,
         entry: "ExpStashEntry",
@@ -179,25 +214,26 @@ class WorkspaceExecutor(BaseLocalExecutor):
         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),
-                force=True,
-                client="dvc",
+        with get_exp_rwlock(repo, writes=[EXEC_NAMESPACE]):
+            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),
+                    force=True,
+                    client="dvc",
+                )
             )
-        )
-        merge_rev = self.scm.get_ref(EXEC_MERGE)
-        try:
-            self.scm.merge(merge_rev, squash=True, commit=False)
-        except _SCMError as exc:
-            raise GitMergeError(str(exc), scm=self.scm)
-        if branch:
-            self.scm.set_ref(EXEC_BRANCH, branch, symbolic=True)
-        elif scm.get_ref(EXEC_BRANCH):
-            self.scm.remove_ref(EXEC_BRANCH)
+            merge_rev = self.scm.get_ref(EXEC_MERGE)
+            try:
+                self.scm.merge(merge_rev, squash=True, commit=False)
+            except _SCMError as exc:
+                raise GitMergeError(str(exc), scm=self.scm)  # noqa: B904
+            if branch:
+                self.scm.set_ref(EXEC_BRANCH, branch, symbolic=True)
+            elif scm.get_ref(EXEC_BRANCH):
+                self.scm.remove_ref(EXEC_BRANCH)
 
     def init_cache(self, repo: "Repo", rev: str, run_cache: bool = True):
         pass

dvc/repo/experiments/executor/ssh.py

@@ -8,13 +8,15 @@ from typing import TYPE_CHECKING, Callable, Iterable, Optional
 from dvc_ssh import SSHFileSystem
 from funcy import first
 
-from ..refs import (
+from dvc.repo.experiments.refs import (
+    EXEC_BASELINE,
     EXEC_BRANCH,
     EXEC_CHECKPOINT,
     EXEC_HEAD,
     EXEC_MERGE,
     EXEC_NAMESPACE,
 )
+
 from .base import BaseExecutor, ExecutorInfo, ExecutorResult, TaskStatus
 
 if TYPE_CHECKING:
@@ -79,6 +81,7 @@ class SSHExecutor(BaseExecutor):
         **kwargs,
     ):
         machine_name: Optional[str] = kwargs.pop("machine_name", None)
+        assert repo.machine
         executor = cls._from_stash_entry(
             repo,
             entry,
@@ -111,21 +114,21 @@ class SSHExecutor(BaseExecutor):
 
     @staticmethod
     def _git_client_args(fs):
-        kwargs = {
+        return {
             "password": fs.fs_args.get("password"),
             "key_filename": first(fs.fs_args.get("client_keys", [])),
         }
-        return kwargs
 
     def init_git(
         self,
+        repo: "Repo",
         scm: "Git",
         stash_rev: str,
         entry: "ExpStashEntry",
         infofile: Optional[str],
         branch: Optional[str] = None,
     ):
-        from ..utils import push_refspec
+        from dvc.repo.experiments.utils import push_refspec
 
         self.status = TaskStatus.PREPARING
         if infofile:
@@ -147,12 +150,18 @@ class SSHExecutor(BaseExecutor):
             # (see https://github.com/iterative/dvc/issues/6508)
             kwargs = self._git_client_args(fs)
 
-            with self.set_exec_refs(scm, stash_rev, entry):
-                refspec = f"{EXEC_NAMESPACE}/"
-                push_refspec(scm, self.git_url, refspec, refspec, **kwargs)
+            ref_dict = {
+                EXEC_HEAD: entry.head_rev,
+                EXEC_MERGE: stash_rev,
+                EXEC_BASELINE: entry.baseline_rev,
+            }
+            with self.set_temp_refs(scm, ref_dict):
+                exec_namespace = f"{EXEC_NAMESPACE}/"
+                refspec = [(exec_namespace, exec_namespace)]
+                push_refspec(scm, self.git_url, refspec, **kwargs)
 
             if branch:
-                push_refspec(scm, self.git_url, branch, branch, **kwargs)
+                push_refspec(scm, self.git_url, [(branch, branch)], **kwargs)
                 self._ssh_cmd(fs, f"git symbolic-ref {EXEC_BRANCH} {branch}")
             else:
                 self._ssh_cmd(
@@ -235,8 +244,8 @@ class SSHExecutor(BaseExecutor):
     def reproduce(
         cls,
         info: "ExecutorInfo",
-        rev: str,
-        queue: Optional["Queue"] = None,
+        rev: str,  # noqa: ARG003
+        queue: Optional["Queue"] = None,  # noqa: ARG003
         infofile: Optional[str] = None,
         log_errors: bool = True,
         log_level: Optional[int] = None,

dvc/repo/experiments/init.py

@@ -22,8 +22,8 @@ from dvc.stage import PipelineStage
 from dvc.types import OptStr
 
 if TYPE_CHECKING:
+    from dvc.dvcfile import ProjectFile, SingleStageFile
     from dvc.repo import Repo
-    from dvc.dvcfile import DVCFile
     from dvc.dependency import Dependency
 
 from dvc.ui import ui
@@ -41,8 +41,10 @@ PROMPTS = {
 
 def _prompts(
     keys: Iterable[str],
-    defaults: Dict[str, str] = None,
-    validator: Callable[[str, str], Union[str, Tuple[str, str]]] = None,
+    defaults: Optional[Dict[str, str]] = None,
+    validator: Optional[
+        Callable[[str, str], Union[str, Tuple[str, str]]]
+    ] = None,
     allow_omission: bool = True,
     stream: Optional[TextIO] = None,
 ) -> Dict[str, OptStr]:
@@ -77,7 +79,9 @@ def _disable_logging(highest_level=logging.CRITICAL):
 def init_interactive(
     defaults: Dict[str, str],
     provided: Dict[str, str],
-    validator: Callable[[str, str], Union[str, Tuple[str, str]]] = None,
+    validator: Optional[
+        Callable[[str, str], Union[str, Tuple[str, str]]]
+    ] = None,
     stream: Optional[TextIO] = None,
 ) -> Dict[str, str]:
     command_prompts = lremove(provided.keys(), ["cmd"])
@@ -112,7 +116,9 @@ def init_interactive(
 
 
 def _check_stage_exists(
-    dvcfile: "DVCFile", name: str, force: bool = False
+    dvcfile: Union["ProjectFile", "SingleStageFile"],
+    name: str,
+    force: bool = False,
 ) -> None:
     if not force and dvcfile.exists() and name in dvcfile.stages:
         from dvc.stage.exceptions import DuplicateStageName
@@ -142,14 +148,13 @@ def validate_prompts(
         except MissingParamsFile:
             return value, msg_format.format(value, "file")
         except ParamsIsADirectoryError:
-            raise InvalidResponse(
+            raise InvalidResponse(  # noqa: B904
                 f"[prompt.invalid]'{value}' is a directory. "
                 "Please retry with an existing parameters file."
             )
-    elif key in ("code", "data"):
-        if not os.path.exists(value):
-            typ = "file" if is_file(value) else "directory"
-            return value, msg_format.format(value, typ)
+    elif key in ("code", "data") and not os.path.exists(value):
+        typ = "file" if is_file(value) else "directory"
+        return value, msg_format.format(value, typ)
     return value
 
 
@@ -200,16 +205,16 @@ def init_out_dirs(stage: PipelineStage) -> List[str]:
 def init(
     repo: "Repo",
     name: str = "train",
-    type: str = "default",  # pylint: disable=redefined-builtin
-    defaults: Dict[str, str] = None,
-    overrides: Dict[str, str] = None,
+    type: str = "default",  # noqa: A002, pylint: disable=redefined-builtin
+    defaults: Optional[Dict[str, str]] = None,
+    overrides: Optional[Dict[str, str]] = None,
     interactive: bool = False,
     force: bool = False,
     stream: Optional[TextIO] = None,
 ) -> Tuple[PipelineStage, List["Dependency"], List[str]]:
-    from dvc.dvcfile import make_dvcfile
+    from dvc.dvcfile import PROJECT_FILE, load_file
 
-    dvcfile = make_dvcfile(repo, "dvc.yaml")
+    dvcfile = load_file(repo, PROJECT_FILE)
     _check_stage_exists(dvcfile, name, force=force)
 
     defaults = defaults.copy() if defaults else {}
@@ -244,7 +249,9 @@ def init(
         try:
             ParamsDependency(None, params, repo=repo).validate_filepath()
         except ParamsIsADirectoryError as exc:
-            raise DvcException(f"{exc}.")  # swallow cause for display
+            raise DvcException(  # noqa: B904
+                f"{exc}."
+            )  # swallow cause for display
         except MissingParamsFile:
             pass
 
@@ -276,6 +283,7 @@ def init(
             plots_key: compact(plots),
         },
     )
+    assert isinstance(stage, PipelineStage)
 
     with _disable_logging(), repo.scm_context(autostage=True, quiet=True):
         stage.dump(update_lock=False)
@@ -285,5 +293,4 @@ def init(
         if params:
             repo.scm_context.track_file(params)
 
-    assert isinstance(stage, PipelineStage)
     return stage, initialized_deps, initialized_out_dirs

dvc/repo/experiments/ls.py

@@ -6,7 +6,7 @@ from dvc.repo.scm_context import scm_context
 from dvc.scm import iter_revs
 from dvc.types import Optional
 
-from .utils import exp_refs, exp_refs_by_baseline
+from .utils import exp_refs_by_baseline
 
 logger = logging.getLogger(__name__)
 
@@ -20,19 +20,22 @@ def ls(
     num: int = 1,
     git_remote: Optional[str] = None,
 ):
-    results = defaultdict(list)
-    if all_commits:
-        gen = exp_refs(repo.scm, git_remote)
-        for info in gen:
-            results[info.baseline_sha].append(info.name)
-        return results
-
-    rev = rev or "HEAD"
-
-    revs = iter_revs(repo.scm, [rev], num)
-    rev_set = set(revs.keys())
+    rev_set = None
+    if not all_commits:
+        revs = iter_revs(repo.scm, [rev or "HEAD"], num)
+        rev_set = set(revs.keys())
     ref_info_dict = exp_refs_by_baseline(repo.scm, rev_set, git_remote)
-    for rev, ref_info_list in ref_info_dict.items():
-        results[rev] = [ref_info.name for ref_info in ref_info_list]
+
+    tags = repo.scm.describe(ref_info_dict.keys())
+    remained = {baseline for baseline, tag in tags.items() if tag is None}
+    base = "refs/heads"
+    ref_heads = repo.scm.describe(remained, base=base)
+
+    results = defaultdict(list)
+    for baseline in ref_info_dict:
+        name = baseline[:7]
+        if tags[baseline] or ref_heads[baseline]:
+            name = tags[baseline] or ref_heads[baseline][len(base) + 1 :]
+        results[name] = [info.name for info in ref_info_dict[baseline]]
 
     return results

dvc/repo/experiments/pull.py

@@ -18,7 +18,7 @@ logger = logging.getLogger(__name__)
 
 @locked
 @scm_context
-def pull(
+def pull(  # noqa: C901
     repo,
     git_remote: str,
     exp_names: Union[Iterable[str], str],

dvc/repo/experiments/push.py

@@ -18,7 +18,7 @@ logger = logging.getLogger(__name__)
 
 @locked
 @scm_context
-def push(
+def push(  # noqa: C901
     repo,
     git_remote: str,
     exp_names: Union[Iterable[str], str],
@@ -82,7 +82,7 @@ def _push(
 ) -> Mapping[SyncStatus, List["ExpRefInfo"]]:
     from scmrepo.exceptions import AuthError
 
-    from ...scm import GitAuthError
+    from dvc.scm import GitAuthError
 
     refspec_list = [f"{exp_ref}:{exp_ref}" for exp_ref in refs]
     logger.debug("git push experiment '%s' -> '%s'", refspec_list, git_remote)
@@ -96,7 +96,7 @@ def _push(
                 progress=pbar.update_git,
             )
         except AuthError as exc:
-            raise GitAuthError(str(exc))
+            raise GitAuthError(str(exc))  # noqa: B904
 
     def group_result(refspec):
         return results[str(refspec)]

dvc/repo/experiments/queue/base.py

@@ -17,25 +17,31 @@ from typing import (
     Union,
 )
 
-from funcy import cached_property, retry
+from funcy import retry
 
 from dvc.dependency import ParamsDependency
-from dvc.env import DVCLIVE_RESUME
+from dvc.env import DVC_EXP_BASELINE_REV, DVC_EXP_NAME, DVCLIVE_RESUME
 from dvc.exceptions import DvcException
 from dvc.lock import LockError
-from dvc.ui import ui
-
-from ..exceptions import CheckpointExistsError, ExperimentExistsError
-from ..executor.base import (
+from dvc.repo.experiments.exceptions import (
+    CheckpointExistsError,
+    ExperimentExistsError,
+)
+from dvc.repo.experiments.executor.base import BaseExecutor, ExecutorResult
+from dvc.repo.experiments.executor.local import WorkspaceExecutor
+from dvc.repo.experiments.refs import ExpRefInfo
+from dvc.repo.experiments.stash import ExpStash, ExpStashEntry
+from dvc.repo.experiments.utils import (
     EXEC_PID_DIR,
     EXEC_TMP_DIR,
-    BaseExecutor,
-    ExecutorResult,
+    exp_refs_by_rev,
+    get_exp_rwlock,
+    get_random_exp_name,
 )
-from ..executor.local import WorkspaceExecutor
-from ..refs import ExpRefInfo
-from ..stash import ExpStash, ExpStashEntry
-from ..utils import exp_refs_by_rev, scm_locked
+from dvc.ui import ui
+from dvc.utils.objects import cached_property
+
+from .utils import get_remote_executor_refs
 
 if TYPE_CHECKING:
     from scmrepo.git import Git
@@ -84,6 +90,11 @@ class QueueDoneResult(NamedTuple):
     result: Optional[ExecutorResult]
 
 
+class ExpRefAndQueueEntry(NamedTuple):
+    exp_ref_info: Optional["ExpRefInfo"]
+    queue_entry: Optional["QueueEntry"]
+
+
 class BaseStashQueue(ABC):
     """Naive Git-stash based experiment queue.
 
@@ -101,6 +112,7 @@ class BaseStashQueue(ABC):
             failed_ref: Failed run Git stash ref for this queue.
         """
         self.repo = repo
+        assert self.repo.tmp_dir
         self.ref = ref
         self.failed_ref = failed_ref
 
@@ -118,10 +130,12 @@ class BaseStashQueue(ABC):
 
     @cached_property
     def pid_dir(self) -> str:
+        assert self.repo.tmp_dir is not None
         return os.path.join(self.repo.tmp_dir, EXEC_TMP_DIR, EXEC_PID_DIR)
 
     @cached_property
-    def args_file(self):
+    def args_file(self) -> str:
+        assert self.repo.tmp_dir is not None
         return os.path.join(self.repo.tmp_dir, BaseExecutor.PACKED_ARGS_FILE)
 
     @abstractmethod
@@ -285,7 +299,7 @@ class BaseStashQueue(ABC):
                 output.
         """
 
-    def _stash_exp(
+    def _stash_exp(  # noqa: C901
         self,
         *args,
         params: Optional[Dict[str, List[str]]] = None,
@@ -400,7 +414,14 @@ class BaseStashQueue(ABC):
                             )
 
                     # save additional repro command line arguments
-                    run_env = {DVCLIVE_RESUME: "1"} if resume_rev else {}
+                    run_env = {
+                        DVC_EXP_BASELINE_REV: baseline_rev,
+                    }
+                    if not name:
+                        name = get_random_exp_name(self.scm, baseline_rev)
+                    run_env[DVC_EXP_NAME] = name
+                    if resume_rev:
+                        run_env[DVCLIVE_RESUME] = "1"
                     self._pack_args(*args, run_env=run_env, **kwargs)
 
                     # save experiment as a stash commit
@@ -494,7 +515,7 @@ class BaseStashQueue(ABC):
             with open(self.args_file, "rb") as fobj:
                 try:
                     data = pickle.load(fobj)
-                except Exception:  # pylint: disable=broad-except
+                except Exception:  # noqa: BLE001, pylint: disable=broad-except
                     data = {}
             extra = int(data.get("extra", 0)) + 1
         else:
@@ -555,28 +576,39 @@ class BaseStashQueue(ABC):
 
     @staticmethod
     @retry(180, errors=LockError, timeout=1)
-    @scm_locked
+    def get_stash_entry(
+        exp: "Experiments",
+        queue_entry: QueueEntry,
+    ) -> "ExpStashEntry":
+        stash = ExpStash(exp.scm, queue_entry.stash_ref)
+        stash_rev = queue_entry.stash_rev
+        with get_exp_rwlock(exp.repo, writes=[queue_entry.stash_ref]):
+            stash_entry = stash.stash_revs.get(
+                stash_rev,
+                ExpStashEntry(None, stash_rev, stash_rev, None, None),
+            )
+            if stash_entry.stash_index is not None:
+                stash.drop(stash_entry.stash_index)
+        return stash_entry
+
+    @classmethod
     def init_executor(
+        cls,
         exp: "Experiments",
         queue_entry: QueueEntry,
         executor_cls: Type[BaseExecutor] = WorkspaceExecutor,
         **kwargs,
     ) -> BaseExecutor:
-        scm = exp.scm
-        stash = ExpStash(scm, queue_entry.stash_ref)
-        stash_rev = queue_entry.stash_rev
-        stash_entry = stash.stash_revs.get(
-            stash_rev,
-            ExpStashEntry(None, stash_rev, stash_rev, None, None),
-        )
-        if stash_entry.stash_index is not None:
-            stash.drop(stash_entry.stash_index)
+        stash_entry = cls.get_stash_entry(exp, queue_entry)
+
         executor = executor_cls.from_stash_entry(
             exp.repo, stash_entry, **kwargs
         )
 
+        stash_rev = queue_entry.stash_rev
         infofile = exp.celery_queue.get_infofile_path(stash_rev)
         executor.init_git(
+            exp.repo,
             exp.repo.scm,
             stash_rev,
             stash_entry,
@@ -597,7 +629,6 @@ class BaseStashQueue(ABC):
 
     @staticmethod
     @retry(180, errors=LockError, timeout=1)
-    @scm_locked
     def collect_git(
         exp: "Experiments",
         executor: BaseExecutor,
@@ -611,16 +642,20 @@ class BaseStashQueue(ABC):
                 raise CheckpointExistsError(ref_info.name)
             raise ExperimentExistsError(ref_info.name)
 
-        for ref in executor.fetch_exps(
-            exp.scm,
-            force=exec_result.force,
-            on_diverged=on_diverged,
-        ):
-            exp_rev = exp.scm.get_ref(ref)
-            if exp_rev:
-                assert exec_result.exp_hash
-                logger.debug("Collected experiment '%s'.", exp_rev[:7])
-                results[exp_rev] = exec_result.exp_hash
+        refs = get_remote_executor_refs(exp.scm, executor.git_url)
+
+        with get_exp_rwlock(exp.repo, writes=refs):
+            for ref in executor.fetch_exps(
+                exp.scm,
+                refs,
+                force=exec_result.force,
+                on_diverged=on_diverged,
+            ):
+                exp_rev = exp.scm.get_ref(ref)
+                if exp_rev:
+                    assert exec_result.exp_hash
+                    logger.debug("Collected experiment '%s'.", exp_rev[:7])
+                    results[exp_rev] = exec_result.exp_hash
 
         return results
 

dvc/repo/experiments/queue/celery.py

@@ -12,24 +12,35 @@ from typing import (
     NamedTuple,
     Optional,
     Set,
+    Union,
 )
 
 from celery.result import AsyncResult
-from funcy import cached_property, first
+from funcy import first
 from kombu.message import Message
 
 from dvc.daemon import daemonize
 from dvc.exceptions import DvcException
+from dvc.repo.experiments.exceptions import UnresolvedQueueExpNamesError
+from dvc.repo.experiments.executor.base import ExecutorInfo, ExecutorResult
+from dvc.repo.experiments.refs import CELERY_STASH
+from dvc.repo.experiments.utils import EXEC_TMP_DIR, get_exp_rwlock
 from dvc.ui import ui
+from dvc.utils.objects import cached_property
 
-from ..exceptions import UnresolvedQueueExpNamesError
-from ..executor.base import EXEC_TMP_DIR, ExecutorInfo, ExecutorResult
-from .base import BaseStashQueue, QueueDoneResult, QueueEntry, QueueGetResult
+from .base import (
+    BaseStashQueue,
+    ExpRefAndQueueEntry,
+    QueueDoneResult,
+    QueueEntry,
+    QueueGetResult,
+)
 from .exceptions import CannotKillTasksError
 from .tasks import run_exp
 from .utils import fetch_running_exp_from_temp_dir
 
 if TYPE_CHECKING:
+    from dvc.repo.experiments.refs import ExpRefInfo
     from dvc_task.app import FSApp
     from dvc_task.proc.manager import ProcessManager
     from dvc_task.worker import TemporaryWorker
@@ -55,15 +66,14 @@ class LocalCeleryQueue(BaseStashQueue):
 
     CELERY_DIR = "celery"
 
-    _shutdown_task_ids: Set[str] = set()
-
     @cached_property
     def wdir(self) -> str:
+        assert self.repo.tmp_dir is not None
         return os.path.join(self.repo.tmp_dir, EXEC_TMP_DIR, self.CELERY_DIR)
 
     @cached_property
     def celery(self) -> "FSApp":
-        from kombu.transport.filesystem import Channel  # type: ignore
+        from kombu.transport.filesystem import Channel
 
         # related to https://github.com/iterative/dvc-task/issues/61
         Channel.QoS.restore_at_shutdown = False
@@ -114,7 +124,7 @@ class LocalCeleryQueue(BaseStashQueue):
             timeout=10,
         )
 
-    def spawn_worker(self, num: int = 1):
+    def _spawn_worker(self, num: int = 1):
         """spawn one single worker to process to queued tasks.
 
         Argument:
@@ -157,14 +167,15 @@ class LocalCeleryQueue(BaseStashQueue):
             if node_name in active_worker:
                 logger.debug(f"Exp queue worker {node_name} already exist")
                 continue
-            self.spawn_worker(num)
+            self._spawn_worker(num)
             started += 1
 
         return started
 
     def put(self, *args, **kwargs) -> QueueEntry:
         """Stash an experiment and add it to the queue."""
-        entry = self._stash_exp(*args, **kwargs)
+        with get_exp_rwlock(self.repo, writes=["workspace", CELERY_STASH]):
+            entry = self._stash_exp(*args, **kwargs)
         self.celery.signature(run_exp.s(entry.asdict())).delay()
         return entry
 
@@ -222,10 +233,10 @@ class LocalCeleryQueue(BaseStashQueue):
                 exp_result = self.get_result(entry)
             except FileNotFoundError:
                 if result.status == "SUCCESS":
-                    raise DvcException(
+                    raise DvcException(  # noqa: B904
                         f"Invalid experiment '{entry.stash_rev[:7]}'."
                     )
-                elif result.status == "FAILURE":
+                if result.status == "FAILURE":
                     exp_result = None
             yield QueueDoneResult(entry, exp_result)
 
@@ -297,18 +308,23 @@ class LocalCeleryQueue(BaseStashQueue):
         return running_task_ids
 
     def _try_to_kill_tasks(
-        self, to_kill: Dict[QueueEntry, str]
+        self, to_kill: Dict[QueueEntry, str], force: bool
     ) -> Dict[QueueEntry, str]:
         fail_to_kill_entries: Dict[QueueEntry, str] = {}
         for queue_entry, rev in to_kill.items():
             try:
-                self.proc.kill(queue_entry.stash_rev)
+                if force:
+                    self.proc.kill(queue_entry.stash_rev)
+                else:
+                    self.proc.interrupt(queue_entry.stash_rev)
                 logger.debug(f"Task {rev} had been killed.")
             except ProcessLookupError:
                 fail_to_kill_entries[queue_entry] = rev
         return fail_to_kill_entries
 
-    def _mark_inactive_tasks_failure(self, remained_entries):
+    def _mark_inactive_tasks_failure(
+        self, remained_entries: Dict[QueueEntry, str]
+    ) -> None:
         remained_revs: List[str] = []
         running_ids = self._get_running_task_ids()
         logger.debug(f"Current running tasks ids: {running_ids}.")
@@ -330,36 +346,47 @@ class LocalCeleryQueue(BaseStashQueue):
         if remained_revs:
             raise CannotKillTasksError(remained_revs)
 
-    def kill(self, revs: Collection[str]) -> None:
+    def _kill_entries(
+        self, entries: Dict[QueueEntry, str], force: bool
+    ) -> None:
+        logger.debug(
+            "Found active tasks: '%s' to kill",
+            list(entries.values()),
+        )
+        inactive_entries: Dict[QueueEntry, str] = self._try_to_kill_tasks(
+            entries, force
+        )
+
+        if inactive_entries:
+            self._mark_inactive_tasks_failure(inactive_entries)
+
+    def kill(self, revs: Collection[str], force: bool = False) -> None:
         name_dict: Dict[
             str, Optional[QueueEntry]
         ] = self.match_queue_entry_by_name(set(revs), self.iter_active())
 
-        to_kill: Dict[QueueEntry, str] = {}
         missing_revs: List[str] = []
+        to_kill: Dict[QueueEntry, str] = {}
         for rev, queue_entry in name_dict.items():
             if queue_entry is None:
                 missing_revs.append(rev)
             else:
                 to_kill[queue_entry] = rev
+
         if missing_revs:
             raise UnresolvedQueueExpNamesError(missing_revs)
 
-        fail_to_kill_entries: Dict[QueueEntry, str] = self._try_to_kill_tasks(
-            to_kill
-        )
-
-        if fail_to_kill_entries:
-            self._mark_inactive_tasks_failure(fail_to_kill_entries)
+        if to_kill:
+            self._kill_entries(to_kill, force)
 
     def shutdown(self, kill: bool = False):
         self.celery.control.shutdown()
         if kill:
-            for _, task_entry in self._iter_active_tasks():
-                try:
-                    self.proc.kill(task_entry.stash_rev)
-                except ProcessLookupError:
-                    continue
+            to_kill: Dict[QueueEntry, str] = {}
+            for entry in self.iter_active():
+                to_kill[entry] = entry.name or entry.stash_rev
+            if to_kill:
+                self._kill_entries(to_kill, True)
 
     def follow(
         self,
@@ -399,7 +426,9 @@ class LocalCeleryQueue(BaseStashQueue):
         try:
             proc_info = self.proc[queue_entry.stash_rev]
         except KeyError:
-            raise DvcException(f"No output logs found for experiment '{rev}'")
+            raise DvcException(  # noqa: B904
+                f"No output logs found for experiment '{rev}'"
+            )
         with open(
             proc_info.stdout,
             encoding=encoding or locale.getpreferredencoding(),
@@ -437,3 +466,31 @@ class LocalCeleryQueue(BaseStashQueue):
                 )
             )
         return result
+
+    def get_ref_and_entry_by_names(
+        self,
+        exp_names: Union[str, List[str]],
+        git_remote: Optional[str] = None,
+    ) -> Dict[str, ExpRefAndQueueEntry]:
+        """Find finished ExpRefInfo or queued or failed QueueEntry by name"""
+        from dvc.repo.experiments.utils import resolve_name
+
+        if isinstance(exp_names, str):
+            exp_names = [exp_names]
+        results: Dict[str, ExpRefAndQueueEntry] = {}
+
+        exp_ref_match: Dict[str, Optional["ExpRefInfo"]] = resolve_name(
+            self.scm, exp_names, git_remote
+        )
+        if not git_remote:
+            queue_entry_match: Dict[
+                str, Optional["QueueEntry"]
+            ] = self.match_queue_entry_by_name(
+                exp_names, self.iter_queued(), self.iter_done()
+            )
+
+        for exp_name in exp_names:
+            exp_ref = exp_ref_match[exp_name]
+            queue_entry = None if git_remote else queue_entry_match[exp_name]
+            results[exp_name] = ExpRefAndQueueEntry(exp_ref, queue_entry)
+        return results

dvc/repo/experiments/queue/remove.py

@@ -9,7 +9,7 @@ if TYPE_CHECKING:
     from dvc.repo.experiments.stash import ExpStashEntry
 
 
-def remove_tasks(
+def remove_tasks(  # noqa: C901
     celery_queue: "LocalCeleryQueue",
     queue_entries: Iterable["QueueEntry"],
 ):
@@ -24,7 +24,9 @@ def remove_tasks(
     failed_stash_revs: List["ExpStashEntry"] = []
     done_entry_set: Set["QueueEntry"] = set()
     stash_rev_all = celery_queue.stash.stash_revs
-    failed_rev_all = celery_queue.failed_stash.stash_revs
+    failed_rev_all: Dict[str, "ExpStashEntry"] = {}
+    if celery_queue.failed_stash:
+        failed_rev_all = celery_queue.failed_stash.stash_revs
     for entry in queue_entries:
         if entry.stash_rev in stash_rev_all:
             stash_revs[entry.stash_rev] = stash_rev_all[entry.stash_rev]
@@ -58,7 +60,8 @@ def remove_tasks(
                 result.forget()
             celery_queue.celery.purge(msg.delivery_tag)
     finally:
-        celery_queue.failed_stash.remove_revs(failed_stash_revs)
+        if celery_queue.failed_stash:
+            celery_queue.failed_stash.remove_revs(failed_stash_revs)
 
 
 def _get_names(entries: Iterable[Union["QueueEntry", "QueueDoneResult"]]):

dvc/repo/experiments/queue/tasks.py

@@ -1,17 +1,22 @@
-from typing import Any, Dict
+from typing import TYPE_CHECKING, Any, Dict
 
 from celery import shared_task
 from celery.utils.log import get_task_logger
 
-from ..executor.base import ExecutorInfo
-from ..executor.local import TempDirExecutor
+from dvc.repo.experiments.executor.base import ExecutorInfo
+from dvc.repo.experiments.executor.local import TempDirExecutor
+
 from .base import BaseStashQueue, QueueEntry
 
+if TYPE_CHECKING:
+    from dvc.repo.experiments.executor.base import BaseExecutor
+
+
 logger = get_task_logger(__name__)
 
 
 @shared_task
-def setup_exp(entry_dict: Dict[str, Any]) -> TempDirExecutor:
+def setup_exp(entry_dict: Dict[str, Any]) -> "BaseExecutor":
     """Setup an experiment.
 
     Arguments:
@@ -39,7 +44,7 @@ def setup_exp(entry_dict: Dict[str, Any]) -> TempDirExecutor:
 
 @shared_task
 def collect_exp(
-    proc_dict: Dict[str, Any],  # pylint: disable=unused-argument
+    proc_dict: Dict[str, Any],  # noqa: ARG001, pylint: disable=unused-argument
     entry_dict: Dict[str, Any],
 ) -> str:
     """Collect results for an experiment.
@@ -69,7 +74,7 @@ def collect_exp(
         else:
             logger.debug("Experiment failed (Exec result was None)")
             celery_queue.stash_failed(entry)
-    except Exception:  # pylint: disable=broad-except
+    except Exception:  # noqa: BLE001, pylint: disable=broad-except
         # Log exceptions but do not re-raise so that task chain execution
         # continues
         logger.exception("Failed to collect experiment")

dvc/repo/experiments/queue/tempdir.py

@@ -1,19 +1,22 @@
 import logging
 import os
+from collections import defaultdict
 from typing import TYPE_CHECKING, Dict, Generator, Optional
 
-from funcy import cached_property, first
+from funcy import first
 
-from ..exceptions import ExpQueueEmptyError
-from ..executor.base import (
-    EXEC_PID_DIR,
-    EXEC_TMP_DIR,
+from dvc.exceptions import DvcException
+from dvc.repo.experiments.exceptions import ExpQueueEmptyError
+from dvc.repo.experiments.executor.base import (
     BaseExecutor,
     ExecutorInfo,
     ExecutorResult,
     TaskStatus,
 )
-from ..executor.local import TempDirExecutor
+from dvc.repo.experiments.executor.local import TempDirExecutor
+from dvc.repo.experiments.utils import EXEC_PID_DIR, EXEC_TMP_DIR
+from dvc.utils.objects import cached_property
+
 from .base import BaseStashQueue, QueueEntry, QueueGetResult
 from .utils import fetch_running_exp_from_temp_dir
 from .workspace import WorkspaceQueue
@@ -35,6 +38,7 @@ class TempDirQueue(WorkspaceQueue):
 
     @cached_property
     def _standalone_tmp_dir(self) -> str:
+        assert self.repo.tmp_dir is not None
         return os.path.join(self.repo.tmp_dir, _STANDALONE_TMP_DIR)
 
     @cached_property
@@ -91,6 +95,51 @@ class TempDirQueue(WorkspaceQueue):
                     None,
                 )
 
+    def _reproduce_entry(
+        self, entry: QueueEntry, executor: BaseExecutor
+    ) -> Dict[str, Dict[str, str]]:
+        from dvc.stage.monitor import CheckpointKilledError
+
+        results: Dict[str, Dict[str, str]] = defaultdict(dict)
+        exec_name = self._EXEC_NAME or entry.stash_rev
+        infofile = self.get_infofile_path(exec_name)
+        try:
+            rev = entry.stash_rev
+            exec_result = executor.reproduce(
+                info=executor.info,
+                rev=rev,
+                infofile=infofile,
+                log_level=logger.getEffectiveLevel(),
+                log_errors=True,
+            )
+            if not exec_result.exp_hash:
+                raise DvcException(
+                    f"Failed to reproduce experiment '{rev[:7]}'"
+                )
+            if exec_result.ref_info:
+                results[rev].update(
+                    self.collect_executor(
+                        self.repo.experiments, executor, exec_result
+                    )
+                )
+        except CheckpointKilledError:
+            results[rev].update(
+                self.collect_executor(
+                    self.repo.experiments, executor, exec_result
+                )
+            )
+
+            return results
+        except DvcException:
+            raise
+        except Exception as exc:  # noqa: BLE001
+            raise DvcException(
+                f"Failed to reproduce experiment '{rev[:7]}'"
+            ) from exc
+        finally:
+            executor.cleanup(infofile)
+        return results
+
     @staticmethod
     def collect_executor(
         exp: "Experiments",

dvc/repo/experiments/queue/utils.py

@@ -1,17 +1,45 @@
 import logging
-from typing import TYPE_CHECKING, Dict
+from typing import TYPE_CHECKING, Dict, List
 
 from scmrepo.exceptions import SCMError
 
-from ..executor.base import ExecutorInfo, TaskStatus
+from dvc.repo.experiments.executor.base import ExecutorInfo, TaskStatus
+from dvc.repo.experiments.refs import (
+    EXEC_CHECKPOINT,
+    EXEC_NAMESPACE,
+    EXPS_NAMESPACE,
+    EXPS_STASH,
+)
+from dvc.repo.experiments.utils import get_exp_rwlock, iter_remote_refs
 
 logger = logging.getLogger(__name__)
 
 
 if TYPE_CHECKING:
+    from scmrepo.git import Git
+
     from .base import BaseStashQueue
 
 
+def get_remote_executor_refs(scm: "Git", remote_url: str) -> List[str]:
+    """Get result list refs from a remote repository
+
+    Args:
+        remote_url : remote executor's url
+    """
+    refs = []
+    for ref in iter_remote_refs(
+        scm,
+        remote_url,
+        base=EXPS_NAMESPACE,
+    ):
+        if ref == EXEC_CHECKPOINT or (
+            not ref.startswith(EXEC_NAMESPACE) and ref != EXPS_STASH
+        ):
+            refs.append(ref)
+    return refs
+
+
 def fetch_running_exp_from_temp_dir(
     queue: "BaseStashQueue", rev: str, fetch_refs: bool
 ) -> Dict[str, Dict]:
@@ -25,18 +53,17 @@ def fetch_running_exp_from_temp_dir(
     Returns:
         Dict[str, Dict]: _description_
     """
+    from dvc.repo.experiments.executor.local import TempDirExecutor
     from dvc.scm import InvalidRemoteSCMRepo
     from dvc.utils.serialize import load_json
 
-    from ..executor.local import TempDirExecutor
-
     result: Dict[str, Dict] = {}
     infofile = queue.get_infofile_path(rev)
     try:
         info = ExecutorInfo.from_dict(load_json(infofile))
     except OSError:
         return result
-    if info.status < TaskStatus.FAILED:
+    if info.status <= TaskStatus.RUNNING:
         result[rev] = info.asdict()
         if info.git_url and fetch_refs and info.status > TaskStatus.PREPARING:
 
@@ -45,15 +72,18 @@ def fetch_running_exp_from_temp_dir(
 
             executor = TempDirExecutor.from_info(info)
             try:
-                for ref in executor.fetch_exps(
-                    queue.scm,
-                    on_diverged=on_diverged,
-                ):
-                    logger.debug("Updated running experiment '%s'.", ref)
-                    last_rev = queue.scm.get_ref(ref)
-                    result[rev]["last"] = last_rev
-                    if last_rev:
-                        result[last_rev] = info.asdict()
+                refs = get_remote_executor_refs(queue.scm, executor.git_url)
+                with get_exp_rwlock(queue.repo, writes=refs):
+                    for ref in executor.fetch_exps(
+                        queue.scm,
+                        refs,
+                        on_diverged=on_diverged,
+                    ):
+                        logger.debug("Updated running experiment '%s'.", ref)
+                        last_rev = queue.scm.get_ref(ref)
+                        result[rev]["last"] = last_rev
+                        if last_rev:
+                            result[last_rev] = info.asdict()
             except (InvalidRemoteSCMRepo, SCMError):
                 # ignore stale info files
                 del result[rev]

dvc/repo/experiments/queue/workspace.py

@@ -9,19 +9,20 @@ from voluptuous import Invalid
 
 from dvc.exceptions import DvcException
 from dvc.lock import make_lock
-from dvc.rwlock import RWLOCK_FILE, RWLOCK_LOCK, SCHEMA
-from dvc.utils import relpath
-from dvc.utils.fs import remove
-
-from ..exceptions import ExpQueueEmptyError
-from ..executor.base import (
+from dvc.repo.experiments.exceptions import ExpQueueEmptyError
+from dvc.repo.experiments.executor.base import (
     BaseExecutor,
     ExecutorInfo,
     ExecutorResult,
     TaskStatus,
 )
-from ..executor.local import WorkspaceExecutor
-from ..refs import EXEC_BRANCH
+from dvc.repo.experiments.executor.local import WorkspaceExecutor
+from dvc.repo.experiments.refs import EXEC_BRANCH, WORKSPACE_STASH
+from dvc.repo.experiments.utils import get_exp_rwlock, to_studio_params
+from dvc.rwlock import RWLOCK_FILE, RWLOCK_LOCK, SCHEMA
+from dvc.utils import relpath
+from dvc.utils.fs import remove
+
 from .base import BaseStashQueue, QueueDoneResult, QueueEntry, QueueGetResult
 
 if TYPE_CHECKING:
@@ -34,7 +35,8 @@ class WorkspaceQueue(BaseStashQueue):
     _EXEC_NAME: Optional[str] = "workspace"
 
     def put(self, *args, **kwargs) -> QueueEntry:
-        return self._stash_exp(*args, **kwargs)
+        with get_exp_rwlock(self.repo, writes=["workspace", WORKSPACE_STASH]):
+            return self._stash_exp(*args, **kwargs)
 
     def get(self) -> QueueGetResult:
         revs = self.stash.stash_revs
@@ -55,7 +57,7 @@ class WorkspaceQueue(BaseStashQueue):
         return QueueGetResult(entry, executor)
 
     def iter_queued(self) -> Generator[QueueEntry, None, None]:
-        for rev, entry in self.stash.stash_revs:
+        for rev, entry in self.stash.stash_revs.items():
             yield QueueEntry(
                 self.repo.root_dir,
                 self.scm.root_dir,
@@ -94,6 +96,8 @@ class WorkspaceQueue(BaseStashQueue):
     def _reproduce_entry(
         self, entry: QueueEntry, executor: BaseExecutor
     ) -> Dict[str, Dict[str, str]]:
+        from dvc_studio_client.post_live_metrics import post_live_metrics
+
         from dvc.stage.monitor import CheckpointKilledError
 
         results: Dict[str, Dict[str, str]] = defaultdict(dict)
@@ -101,6 +105,13 @@ class WorkspaceQueue(BaseStashQueue):
         infofile = self.get_infofile_path(exec_name)
         try:
             rev = entry.stash_rev
+            post_live_metrics(
+                "start",
+                executor.info.baseline_rev,
+                executor.info.name,
+                "dvc",
+                params=to_studio_params(self.repo.params.show()),
+            )
             exec_result = executor.reproduce(
                 info=executor.info,
                 rev=rev,
@@ -123,18 +134,25 @@ class WorkspaceQueue(BaseStashQueue):
             return {}
         except DvcException:
             raise
-        except Exception as exc:
+        except Exception as exc:  # noqa: BLE001
             raise DvcException(
                 f"Failed to reproduce experiment '{rev[:7]}'"
             ) from exc
         finally:
+            post_live_metrics(
+                "done",
+                executor.info.baseline_rev,
+                executor.info.name,
+                "dvc",
+                experiment_rev=first(results.get(rev, {})),
+            )
             executor.cleanup(infofile)
         return results
 
     @staticmethod
     def collect_executor(  # pylint: disable=unused-argument
         exp: "Experiments",
-        executor: BaseExecutor,
+        executor: BaseExecutor,  # noqa: ARG004
         exec_result: ExecutorResult,
     ) -> Dict[str, str]:
         results: Dict[str, str] = {}

dvc/repo/experiments/refs.py

@@ -16,6 +16,7 @@ EXEC_BRANCH = f"{EXEC_NAMESPACE}/EXEC_BRANCH"
 EXEC_BASELINE = f"{EXEC_NAMESPACE}/EXEC_BASELINE"
 EXEC_HEAD = f"{EXEC_NAMESPACE}/EXEC_HEAD"
 EXEC_MERGE = f"{EXEC_NAMESPACE}/EXEC_MERGE"
+EXPS_TEMP = f"{EXPS_NAMESPACE}/temp"
 STASHES = {WORKSPACE_STASH, CELERY_STASH}
 
 
@@ -23,9 +24,7 @@ class ExpRefInfo:
 
     namespace = EXPS_NAMESPACE
 
-    def __init__(
-        self, baseline_sha: Optional[str] = None, name: Optional[str] = None
-    ):
+    def __init__(self, baseline_sha: str, name: Optional[str] = None):
         self.baseline_sha = baseline_sha
         self.name: str = name if name else ""
 
@@ -33,7 +32,7 @@ class ExpRefInfo:
         return "/".join(self.parts)
 
     def __repr__(self):
-        baseline = f"'{self.baseline_sha}'" if self.baseline_sha else "None"
+        baseline = f"'{self.baseline_sha}'"
         name = f"'{self.name}'" if self.name else "None"
         return f"ExpRefInfo(baseline_sha={baseline}, name={name})"
 
@@ -41,11 +40,7 @@ class ExpRefInfo:
     def parts(self):
         return (
             (self.namespace,)
-            + (
-                (self.baseline_sha[:2], self.baseline_sha[2:])
-                if self.baseline_sha
-                else ()
-            )
+            + ((self.baseline_sha[:2], self.baseline_sha[2:]))
             + ((self.name,) if self.name else ())
         )
 
@@ -54,14 +49,13 @@ class ExpRefInfo:
         try:
             parts = ref.split("/")
             if (
-                len(parts) < 2
-                or len(parts) == 3
+                len(parts) < 4
                 or len(parts) > 5
                 or "/".join(parts[:2]) != EXPS_NAMESPACE
             ):
                 raise InvalidExpRefError(ref)
         except ValueError:
-            raise InvalidExpRefError(ref)
-        baseline_sha = parts[2] + parts[3] if len(parts) >= 4 else None
+            raise InvalidExpRefError(ref)  # noqa: B904
+        baseline_sha = parts[2] + parts[3]
         name = parts[4] if len(parts) == 5 else None
         return cls(baseline_sha, name)

dvc/repo/experiments/remove.py

@@ -1,19 +1,12 @@
 import logging
-from typing import (
-    TYPE_CHECKING,
-    Dict,
-    Iterable,
-    List,
-    NamedTuple,
-    Optional,
-    Union,
-)
+from typing import TYPE_CHECKING, Dict, Iterable, List, Optional, Union
 
 from dvc.repo import locked
 from dvc.repo.scm_context import scm_context
 from dvc.scm import iter_revs
 
 from .exceptions import UnresolvedExpNamesError
+from .queue.base import QueueEntry
 from .refs import ExpRefInfo
 from .utils import exp_refs, exp_refs_by_baseline, push_refspec
 
@@ -22,61 +15,15 @@ if TYPE_CHECKING:
     from dvc.repo.experiments.queue.celery import LocalCeleryQueue
     from dvc.scm import Git
 
-    from .queue.base import QueueEntry
+    from .queue.base import ExpRefAndQueueEntry
 
 
 logger = logging.getLogger(__name__)
 
 
-class ExpRefAndQueueEntry(NamedTuple):
-    exp_ref_list: List["ExpRefInfo"]
-    queue_entry_list: List["QueueEntry"]
-    removed: List["str"]
-
-
-def _get_ref_and_entry_by_names(
-    exp_names: Union[str, List[str]],
-    scm: "Git",
-    celery_queue: "LocalCeleryQueue",
-    git_remote: Optional[str],
-) -> ExpRefAndQueueEntry:
-    from .utils import resolve_name
-
-    exp_ref_list: List["ExpRefInfo"] = []
-    queue_entry_list: List["QueueEntry"] = []
-    removed: List[str] = []
-    if isinstance(exp_names, str):
-        exp_names = [exp_names]
-    exp_ref_match: Dict[str, Optional["ExpRefInfo"]] = resolve_name(
-        scm, exp_names, git_remote
-    )
-    if not git_remote:
-        queue_entry_match: Dict[
-            str, Optional["QueueEntry"]
-        ] = celery_queue.match_queue_entry_by_name(
-            exp_names, celery_queue.iter_queued(), celery_queue.iter_done()
-        )
-
-    remained = []
-    for exp_name in exp_names:
-        exp_ref = exp_ref_match[exp_name]
-        queue_entry = None if git_remote else queue_entry_match[exp_name]
-        if exp_ref or queue_entry:
-            if exp_ref:
-                exp_ref_list.append(exp_ref)
-            if queue_entry:
-                queue_entry_list.append(queue_entry)
-            removed.append(exp_name)
-        else:
-            remained.append(exp_name)
-    if remained:
-        raise UnresolvedExpNamesError(remained)
-    return ExpRefAndQueueEntry(exp_ref_list, queue_entry_list, removed)
-
-
 @locked
 @scm_context
-def remove(
+def remove(  # noqa: C901
     repo: "Repo",
     exp_names: Union[None, str, List[str]] = None,
     rev: Optional[str] = None,
@@ -103,12 +50,22 @@ def remove(
     exp_ref_list: List["ExpRefInfo"] = []
     queue_entry_list: List["QueueEntry"] = []
     if exp_names:
-        result = _get_ref_and_entry_by_names(
-            exp_names, repo.scm, celery_queue, git_remote
-        )
-        removed.extend(result.removed)
-        exp_ref_list.extend(result.exp_ref_list)
-        queue_entry_list.extend(result.queue_entry_list)
+        results: Dict[
+            str, "ExpRefAndQueueEntry"
+        ] = celery_queue.get_ref_and_entry_by_names(exp_names, git_remote)
+        remained: List[str] = []
+        for name, result in results.items():
+            if not result.exp_ref_info and not result.queue_entry:
+                remained.append(name)
+                continue
+            removed.append(name)
+            if result.exp_ref_info:
+                exp_ref_list.append(result.exp_ref_info)
+            if result.queue_entry:
+                queue_entry_list.append(result.queue_entry)
+
+        if remained:
+            raise UnresolvedExpNamesError(remained)
     elif rev:
         exp_ref_dict = _resolve_exp_by_baseline(repo, rev, num, git_remote)
         removed.extend(exp_ref_dict.keys())
@@ -152,8 +109,7 @@ def _remove_commited_exps(
                 push_refspec(
                     scm,
                     remote,
-                    None,
-                    str(ref_info),
+                    [(None, str(ref_info))],
                     progress=pbar.update_git,
                 )
     else:

dvc/repo/experiments/run.py

@@ -11,7 +11,7 @@ logger = logging.getLogger(__name__)
 
 
 @locked
-def run(
+def run(  # noqa: C901
     repo,
     targets: Optional[Iterable[str]] = None,
     params: Optional[Iterable[str]] = None,