fetch: introduce --type metrics/plots

Needed for https://github.com/iterative/studio/pull/6541 but is also generally useful.

.github/workflows/build.yaml

@@ -1,4 +1,4 @@
-name: Build packages
+name: Build package
 on:
   push:
     branches:
@@ -7,39 +7,35 @@ on:
     types: [released, prereleased]
   workflow_dispatch:
 
-permissions:  # added using https://github.com/step-security/secure-workflows
+permissions:
   contents: read
+  id-token: write
 
 jobs:
   pip:
-    runs-on: ubuntu-20.04
-    permissions:
-      id-token: write
+    runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v3
       with:
         fetch-depth: 0
 
-    - name: Set up Python 3.8
-      uses: actions/setup-python@v4
+    - uses: actions/setup-python@v4
       with:
         python-version: 3.8
 
-    - name: Install
-      run: pip install --upgrade pip wheel
-
-    - name: Force version for Test PyPI uploads
+    - run: python -m pip install --upgrade pip setuptools_scm build twine
+    - name: Force version for Test PyPI upload
       if: ${{ !startsWith(github.ref, 'refs/tags') }}
-      run: |
-        pip install setuptools_scm
-        echo version=$(python -m setuptools_scm | awk -F+ '{print $1}' | tail -1) >> $GITHUB_ENV
+      run: echo version=$(python -m setuptools_scm | awk -F+ '{print $1}' | tail -1) >> $GITHUB_ENV
 
-    - name: Build packages
-      run: ./scripts/build_package.sh
+    - run: |
+        echo 'PKG = "pip"'>dvc/_build.py
+        python -m build
       env:
         SETUPTOOLS_SCM_PRETEND_VERSION: ${{ env.version }}
+    - run: twine check --strict dist/*
 
-    - name: Publish packages to PyPI
+    - name: Publish to PyPI
       if: github.event_name == 'release'
       uses: pypa/gh-action-pypi-publish@release/v1
 
@@ -47,5 +43,5 @@ jobs:
       if: ${{ github.event_name == 'release' || (github.event_name == 'push' && github.ref == 'refs/heads/main') }}
       uses: pypa/gh-action-pypi-publish@release/v1
       with:
-        repository_url: https://test.pypi.org/legacy/
-        skip_existing: true
+        repository-url: https://test.pypi.org/legacy/
+        skip-existing: true

README.rst

@@ -280,8 +280,8 @@ Barrak, A., Eghan, E.E. and Adams, B. `On the Co-evolution of ML Pipelines and S
    :target: https://pypi.org/project/dvc
    :alt: PyPI Downloads
 
-.. |Packages| image:: https://img.shields.io/github/v/release/iterative/dvc?label=deb|pkg|rpm|exe&logo=GitHub
-   :target: https://github.com/iterative/dvc/releases/latest
+.. |Packages| image:: https://img.shields.io/badge/deb|pkg|rpm|exe-blue
+   :target: https://dvc.org/doc/install
    :alt: deb|pkg|rpm|exe
 
 .. |DOI| image:: https://img.shields.io/badge/DOI-10.5281/zenodo.3677553-blue.svg

dvc/__pyinstaller/hook-dvc.system.py

@@ -1,5 +0,0 @@
-# ruff: noqa: N999
-import os
-
-if os.name == "nt":
-    hiddenimports = ["win32file"]

dvc/__pyinstaller/hook-dvc.tree.gs.py

@@ -1,10 +0,0 @@
-# ruff: noqa: N999
-from PyInstaller.utils.hooks import copy_metadata  # pylint:disable=import-error
-
-datas = copy_metadata("google-cloud-storage")
-
-# NOTE: including "requests" metadata because of google-resumable-media
-# bug. See [1] for more info.
-#
-# [1] https://github.com/googleapis/google-resumable-media-python/issues/59
-datas += copy_metadata("requests")

dvc/__pyinstaller/hook-google_compute_engine.logger.py

@@ -1,5 +0,0 @@
-# ruff: noqa: N999
-from PyInstaller.utils.hooks import copy_metadata  # pylint:disable=import-error
-
-datas = copy_metadata("google-compute-engine")
-hiddenimports = ["google_cloud_engine"]

dvc/_debug.py

@@ -198,33 +198,84 @@ def debugtools(args: Optional["Namespace"] = None, **kwargs):
 def add_debugging_flags(parser):
     from argparse import SUPPRESS
 
-    parser.add_argument("--cprofile", action="store_true", default=False, help=SUPPRESS)
-    parser.add_argument("--yappi", action="store_true", default=False, help=SUPPRESS)
+    # For detailed info see:
+    # https://github.com/iterative/dvc/wiki/Debugging,-Profiling-and-Benchmarking-DVC
+    args, _ = parser.parse_known_args()
+    verbose = args.verbose
+
+    def debug_help(msg):
+        if verbose:
+            return msg
+        return SUPPRESS
+
+    parser = parser.add_argument_group("debug options")
+
+    parser.add_argument(
+        "--cprofile",
+        action="store_true",
+        default=False,
+        help=debug_help("Generate cprofile data for tools like snakeviz / tuna"),
+    )
+    parser.add_argument(
+        "--cprofile-dump", help=debug_help("Location to dump cprofile file")
+    )
+    parser.add_argument(
+        "--yappi",
+        action="store_true",
+        default=False,
+        help=debug_help(
+            "Generate a callgrind file for use with tools like "
+            "kcachegrind / qcachegrind"
+        ),
+    )
     parser.add_argument(
         "--yappi-separate-threads",
         action="store_true",
         default=False,
-        help=SUPPRESS,
+        help=debug_help("Generate one callgrind file per thread"),
     )
     parser.add_argument(
-        "--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)
-    parser.add_argument(
-        "--instrument", action="store_true", default=False, help=SUPPRESS
+        "--viztracer",
+        action="store_true",
+        default=False,
+        help=debug_help("Generate a viztracer file for use with vizviewer"),
     )
     parser.add_argument(
-        "--instrument-open", action="store_true", default=False, help=SUPPRESS
+        "--viztracer-depth",
+        type=int,
+        help=debug_help("Set viztracer maximum stack depth"),
+    )
+    parser.add_argument(
+        "--viztracer-async",
+        action="store_true",
+        default=False,
+        help=debug_help("Treat async tasks as threads"),
+    )
+    parser.add_argument(
+        "--pdb",
+        action="store_true",
+        default=False,
+        help=debug_help("Drop into the pdb/ipdb debugger on any exception"),
+    )
+    parser.add_argument(
+        "--instrument",
+        action="store_true",
+        default=False,
+        help=debug_help("Use pyinstrument CLI profiler"),
+    )
+    parser.add_argument(
+        "--instrument-open",
+        action="store_true",
+        default=False,
+        help=debug_help("Use pyinstrument web profiler"),
     )
     parser.add_argument(
         "--show-stack",
         "--ss",
         action="store_true",
         default=False,
-        help=SUPPRESS,
+        help=debug_help(
+            r"Use Ctrl+T on macOS or Ctrl+\ on Linux to print the stack "
+            "frame currently executing. Unavailable on Windows."
+        ),
     )

dvc/api/data.py

@@ -36,21 +36,31 @@ def get_url(path, repo=None, rev=None, remote=None):
     NOTE: This function does not check for the actual existence of the file or
     directory in the remote storage.
     """
-    with Repo.open(repo, rev=rev, subrepos=True, uninitialized=True) as _repo:
+    from dvc.config import NoRemoteError
+    from dvc_data.index import StorageKeyError
+
+    repo_kwargs: Dict[str, Any] = {}
+    if remote:
+        repo_kwargs["config"] = {"core": {"remote": remote}}
+    with Repo.open(
+        repo, rev=rev, subrepos=True, uninitialized=True, **repo_kwargs
+    ) as _repo:
         with _wrap_exceptions(_repo, path):
             fs_path = _repo.dvcfs.from_os_path(path)
-
-            with reraise(FileNotFoundError, PathMissingError(path, repo)):
-                info = _repo.dvcfs.info(fs_path)
-
-            dvc_info = info.get("dvc_info")
-            if not dvc_info:
-                raise OutputNotFoundError(path, repo)
-
-            dvc_repo = info["repo"]  # pylint: disable=unsubscriptable-object
-            md5 = dvc_info["md5"]
-
-            return dvc_repo.cloud.get_url_for(remote, checksum=md5)
+            fs = _repo.dvcfs.fs
+            # pylint: disable-next=protected-access
+            key = fs._get_key_from_relative(fs_path)
+            # pylint: disable-next=protected-access
+            subrepo, _, subkey = fs._get_subrepo_info(key)
+            index = subrepo.index.data["repo"]
+            with reraise(KeyError, OutputNotFoundError(path, repo)):
+                entry = index[subkey]
+            with reraise(
+                (StorageKeyError, ValueError),
+                NoRemoteError(f"no remote specified in {_repo}"),
+            ):
+                remote_fs, remote_path = index.storage_map.get_remote(entry)
+                return remote_fs.unstrip_protocol(remote_path)
 
 
 class _OpenContextManager(GCM):

dvc/cli/parser.py

@@ -137,7 +137,6 @@ def get_parent_parser():
     from dvc._debug import add_debugging_flags
 
     parent_parser = argparse.ArgumentParser(add_help=False)
-    add_debugging_flags(parent_parser)
     log_level_group = parent_parser.add_mutually_exclusive_group()
     log_level_group.add_argument(
         "-q", "--quiet", action="count", default=0, help="Be quiet."
@@ -145,6 +144,7 @@ def get_parent_parser():
     log_level_group.add_argument(
         "-v", "--verbose", action="count", default=0, help="Be verbose."
     )
+    add_debugging_flags(parent_parser)
 
     return parent_parser
 

dvc/commands/data_sync.py

@@ -91,6 +91,8 @@ class CmdDataFetch(CmdDataBase):
                 with_deps=self.args.with_deps,
                 recursive=self.args.recursive,
                 run_cache=self.args.run_cache,
+                max_size=self.args.max_size,
+                types=self.args.types,
             )
             self.log_summary({"fetched": processed_files_count})
         except DvcException:
@@ -322,6 +324,22 @@ def add_parser(subparsers, _parent_parser):
         default=False,
         help="Fetch run history for all stages.",
     )
+    fetch_parser.add_argument(
+        "--max-size",
+        type=int,
+        help="Fetch data files/directories that are each below specified size (bytes).",
+    )
+    fetch_parser.add_argument(
+        "--type",
+        dest="types",
+        action="append",
+        default=[],
+        help=(
+            "Only fetch data files/directories that are of a particular "
+            "type (metrics, plots)."
+        ),
+        choices=["metrics", "plots"],
+    )
     fetch_parser.set_defaults(func=CmdDataFetch)
 
     # Status

dvc/fs/dvc.py

@@ -5,7 +5,7 @@ import ntpath
 import os
 import posixpath
 import threading
-from contextlib import suppress
+from contextlib import ExitStack, suppress
 from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple, Type, Union
 
 from fsspec.spec import AbstractFileSystem
@@ -54,6 +54,8 @@ def _merge_info(repo, key, fs_info, dvc_info):
         ret["size"] = dvc_info["size"]
         if not fs_info and "md5" in dvc_info:
             ret["md5"] = dvc_info["md5"]
+        if not fs_info and "md5-dos2unix" in dvc_info:
+            ret["md5-dos2unix"] = dvc_info["md5-dos2unix"]
 
     if fs_info:
         ret["type"] = fs_info["type"]
@@ -121,11 +123,13 @@ class _DVCFileSystem(AbstractFileSystem):  # pylint:disable=abstract-method
         from pygtrie import Trie
 
         super().__init__()
+        self._repo_stack = ExitStack()
         if repo is None:
             repo = self._make_repo(url=url, rev=rev, subrepos=subrepos, **repo_kwargs)
             assert repo is not None
             # pylint: disable=protected-access
             repo_factory = repo._fs_conf["repo_factory"]
+            self._repo_stack.enter_context(repo)
 
         if not repo_factory:
             from dvc.repo import Repo
@@ -214,6 +218,7 @@ class _DVCFileSystem(AbstractFileSystem):  # pylint:disable=abstract-method
                     scm=self.repo.scm,
                     repo_factory=self.repo_factory,
                 )
+                self._repo_stack.enter_context(repo)
                 self._datafss[key] = DataFileSystem(index=repo.index.data["repo"])
             self._subrepos_trie[key] = repo
 
@@ -394,6 +399,9 @@ class _DVCFileSystem(AbstractFileSystem):  # pylint:disable=abstract-method
         dvc_path = _get_dvc_path(dvc_fs, subkey)
         return dvc_fs.get_file(dvc_path, lpath, **kwargs)
 
+    def close(self):
+        self._repo_stack.close()
+
 
 class DVCFileSystem(FileSystem):
     protocol = "local"
@@ -427,3 +435,7 @@ class DVCFileSystem(FileSystem):
             path = os.path.relpath(path, self.repo.root_dir)
 
         return as_posix(path)
+
+    def close(self):
+        if "fs" in self.__dict__:
+            self.fs.close()

dvc/output.py

@@ -548,12 +548,24 @@ class Output:
         _, hash_info = self._get_hash_meta()
         return hash_info
 
+    def _build(
+        self, *args, no_progress_bar=False, **kwargs
+    ) -> Tuple["HashFileDB", "Meta", "HashFile"]:
+        from dvc.ui import ui
+
+        with ui.progress(
+            unit="file",
+            desc=f"Collecting files and computing hashes in {self}",
+            disable=no_progress_bar,
+        ) as pb:
+            return build(*args, callback=pb.as_callback(), **kwargs)
+
     def _get_hash_meta(self):
         if self.use_cache:
             odb = self.cache
         else:
             odb = self.local_cache
-        _, meta, obj = build(
+        _, meta, obj = self._build(
             odb,
             self.fs_path,
             self.fs,
@@ -703,7 +715,7 @@ class Output:
 
         self._update_legacy_hash_name()
         if self.use_cache:
-            _, self.meta, self.obj = build(
+            _, self.meta, self.obj = self._build(
                 self.cache,
                 self.fs_path,
                 self.fs,
@@ -711,7 +723,7 @@ class Output:
                 ignore=self.dvcignore,
             )
         else:
-            _, self.meta, self.obj = build(
+            _, self.meta, self.obj = self._build(
                 self.local_cache,
                 self.fs_path,
                 self.fs,
@@ -763,7 +775,7 @@ class Output:
             if granular:
                 obj = self._commit_granular_dir(filter_info, hardlink)
             else:
-                staging, _, obj = build(
+                staging, _, obj = self._build(
                     self.cache,
                     filter_info or self.fs_path,
                     self.fs,
@@ -794,7 +806,7 @@ class Output:
 
     def _commit_granular_dir(self, filter_info, hardlink) -> Optional["HashFile"]:
         prefix = self.fs.path.parts(self.fs.path.relpath(filter_info, self.fs_path))
-        staging, _, obj = build(
+        staging, _, obj = self._build(
             self.cache,
             self.fs_path,
             self.fs,
@@ -1021,7 +1033,7 @@ class Output:
 
         upload = not (update and from_fs.isdir(from_info))
         jobs = jobs or min((from_fs.jobs, odb.fs.jobs))
-        staging, self.meta, obj = build(
+        staging, self.meta, obj = self._build(
             odb,
             from_info,
             from_fs,
@@ -1357,7 +1369,7 @@ class Output:
         new: "HashFile"
         try:
             assert self.hash_name
-            staging, meta, obj = build(
+            staging, meta, obj = self._build(
                 cache,
                 path,
                 self.fs,
@@ -1402,7 +1414,9 @@ class Output:
         otransfer(staging, self.cache, {obj.hash_info}, hardlink=relink, shallow=False)
 
         if relink:
-            with CheckoutCallback(desc="Checking out {path}", unit="files") as callback:
+            with CheckoutCallback(
+                desc=f"Checking out {path}", unit="files"
+            ) as callback:
                 self._checkout(
                     path,
                     self.fs,

dvc/repo/__init__.py

@@ -557,6 +557,7 @@ class Repo:
 
         from dvc.dirs import site_cache_dir
         from dvc.fs import GitFileSystem
+        from dvc.version import version_tuple
 
         cache_dir = self.config["core"].get("site_cache_dir") or site_cache_dir()
 
@@ -578,7 +579,13 @@ class Repo:
             os.umask(umask)
 
         md5 = hashlib.md5(  # noqa: S324  # nosec B324, B303
-            str((root_dir, getpass.getuser())).encode()
+            str(
+                (
+                    root_dir,
+                    getpass.getuser(),
+                    version_tuple[0],
+                )
+            ).encode()
         )
         repo_token = md5.hexdigest()
         return os.path.join(repos_dir, repo_token)
@@ -586,12 +593,16 @@ class Repo:
     def close(self):
         self.scm.close()
         self.state.close()
+        if "dvcfs" in self.__dict__:
+            self.dvcfs.close()
         if self._data_index is not None:
             self._data_index.close()
 
     def _reset(self):
         self.scm._reset()  # pylint: disable=protected-access
         self.state.close()
+        if "dvcfs" in self.__dict__:
+            self.dvcfs.close()
         self.__dict__.pop("index", None)
         self.__dict__.pop("dvcignore", None)
         self.__dict__.pop("dvcfs", None)

dvc/repo/experiments/executor/base.py

@@ -44,7 +44,7 @@ if TYPE_CHECKING:
     from dvc.repo import Repo
     from dvc.repo.experiments.stash import ExpStashEntry
     from dvc.scm import Git
-    from dvc.stage import PipelineStage
+    from dvc.stage import PipelineStage, Stage
 
 logger = logging.getLogger(__name__)
 
@@ -255,6 +255,15 @@ class BaseExecutor(ABC):
             **kwargs,
         )
 
+    @classmethod
+    def _get_stage_files(cls, stages: List["Stage"]) -> List[str]:
+        from dvc.stage.utils import _get_stage_files
+
+        ret: List[str] = []
+        for stage in stages:
+            ret.extend(_get_stage_files(stage))
+        return ret
+
     @classmethod
     def _get_top_level_paths(cls, repo: "Repo") -> List["str"]:
         return list(
@@ -511,6 +520,9 @@ class BaseExecutor(ABC):
                 )
 
             stages = dvc.reproduce(*args, **kwargs)
+            if paths := cls._get_stage_files(stages):
+                logger.debug("Staging stage-related files: %s", paths)
+                dvc.scm_context.add(paths)
             if paths := cls._get_top_level_paths(dvc):
                 logger.debug("Staging top-level files: %s", paths)
                 dvc.scm_context.add(paths)

dvc/repo/fetch.py

@@ -1,6 +1,8 @@
 import logging
+from typing import List, Tuple
 
 from dvc.exceptions import DownloadError
+from dvc_data.index import DataIndex, FileStorage
 
 from . import locked
 
@@ -20,6 +22,8 @@ def fetch(  # noqa: C901, PLR0913
     all_commits=False,
     run_cache=False,
     revs=None,
+    max_size=None,
+    types=None,
 ) -> int:
     """Download data items from a cloud and imported repositories
 
@@ -68,6 +72,8 @@ def fetch(  # noqa: C901, PLR0913
                 targets,
                 with_deps=with_deps,
                 recursive=recursive,
+                max_size=max_size,
+                types=types,
             )
             index_keys.add(idx.data_tree.hash_info.value)
             indexes.append(idx.data["repo"])
@@ -82,8 +88,14 @@ def fetch(  # noqa: C901, PLR0913
         unit="entry",
     ) as cb:
         data = collect(
-            indexes, cache_index=self.data_index, cache_key=cache_key, callback=cb
+            indexes,
+            "remote",
+            cache_index=self.data_index,
+            cache_key=cache_key,
+            callback=cb,
         )
+    data, unversioned_count = _log_unversioned(data)
+    failed_count += unversioned_count
 
     with Callback.as_tqdm_callback(
         desc="Fetching",
@@ -96,7 +108,7 @@ def fetch(  # noqa: C901, PLR0913
                 callback=cb,
             )  # pylint: disable=assignment-from-no-return
         finally:
-            for fs_index in data.values():
+            for fs_index in data:
                 fs_index.close()
 
     if fetch_transferred:
@@ -109,3 +121,34 @@ def fetch(  # noqa: C901, PLR0913
         raise DownloadError(failed_count)
 
     return transferred_count
+
+
+def _log_unversioned(data: List["DataIndex"]) -> Tuple[List["DataIndex"], int]:
+    ret: List["DataIndex"] = []
+    unversioned: List[str] = []
+    for fs_index in data:
+        remote = fs_index.storage_map[()].remote
+        if not isinstance(remote, FileStorage) or not remote.fs.version_aware:
+            ret.append(fs_index)
+            continue
+
+        fs = remote.fs
+        index = DataIndex()
+        index.storage_map = fs_index.storage_map
+        for key, entry in fs_index.iteritems():
+            if entry.meta and not entry.meta.isdir and entry.meta.version_id is None:
+                unversioned.append(fs.unstrip_protocol(fs.path.join(remote.path, *key)))
+            else:
+                index[key] = entry
+        fs_index.close()
+        ret.append(index)
+
+    if unversioned:
+        logger.warning(
+            (
+                "Some files are missing cloud version information and will not be "
+                "fetched from the remote:\n%s"
+            ),
+            "\n".join(unversioned),
+        )
+    return ret, len(unversioned)

dvc/repo/graph.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Iterator, List, Set
+from typing import TYPE_CHECKING, Any, Iterator, List, Optional, Set, TypeVar
 
 from dvc.fs import localfs
 from dvc.utils.fs import path_isin
@@ -8,6 +8,8 @@ if TYPE_CHECKING:
 
     from dvc.stage import Stage
 
+T = TypeVar("T")
+
 
 def check_acyclic(graph: "DiGraph") -> None:
     import networkx as nx
@@ -39,6 +41,21 @@ def get_pipelines(graph: "DiGraph"):
     return [graph.subgraph(c).copy() for c in nx.weakly_connected_components(graph)]
 
 
+def get_subgraph_of_nodes(
+    graph: "DiGraph", sources: Optional[List[Any]] = None, downstream: bool = False
+) -> "DiGraph":
+    from networkx import dfs_postorder_nodes, reverse_view
+
+    if not sources:
+        return graph
+
+    g = reverse_view(graph) if downstream else graph
+    nodes = []
+    for source in sources:
+        nodes.extend(dfs_postorder_nodes(g, source))
+    return graph.subgraph(nodes)
+
+
 def collect_pipeline(stage: "Stage", graph: "DiGraph") -> Iterator["Stage"]:
     import networkx as nx
 

dvc/repo/index.py

@@ -374,6 +374,50 @@ class Index:
 
         return dict(by_workspace)
 
+    @cached_property
+    def metric_keys(self) -> Dict[str, Set["DataIndexKey"]]:
+        from collections import defaultdict
+
+        from .metrics.show import _collect_top_level_metrics
+
+        by_workspace: Dict[str, Set["DataIndexKey"]] = defaultdict(set)
+
+        by_workspace["repo"] = set()
+
+        for out in self.outs:
+            if not out.metric:
+                continue
+
+            workspace, key = out.index_key
+            by_workspace[workspace].add(key)
+
+        for path in _collect_top_level_metrics(self.repo):
+            key = self.repo.fs.path.relparts(path, self.repo.root_dir)
+            by_workspace["repo"].add(key)
+
+        return dict(by_workspace)
+
+    @cached_property
+    def plot_keys(self) -> Dict[str, Set["DataIndexKey"]]:
+        from collections import defaultdict
+
+        by_workspace: Dict[str, Set["DataIndexKey"]] = defaultdict(set)
+
+        by_workspace["repo"] = set()
+
+        for out in self.outs:
+            if not out.plot:
+                continue
+
+            workspace, key = out.index_key
+            by_workspace[workspace].add(key)
+
+        for path in self._plot_sources:
+            key = self.repo.fs.path.relparts(path, self.repo.root_dir)
+            by_workspace["repo"].add(key)
+
+        return dict(by_workspace)
+
     @cached_property
     def data_tree(self):
         from dvc_data.hashfile.tree import Tree
@@ -487,11 +531,31 @@ class Index:
                 used[odb].update(objs)
         return used
 
+    def _types_filter(self, types, out):
+        ws, okey = out.index_key
+        for typ in types:
+            if typ == "plots":
+                keys = self.plot_keys
+            elif typ == "metrics":
+                keys = self.metric_keys
+            else:
+                raise ValueError(f"unsupported type {typ}")
+
+            for key in keys.get(ws, []):
+                if (len(key) >= len(okey) and key[: len(okey)] == okey) or (
+                    len(key) < len(okey) and okey[: len(key)] == key
+                ):
+                    return True
+
+        return False
+
     def targets_view(
         self,
         targets: Optional["TargetType"],
         stage_filter: Optional[Callable[["Stage"], bool]] = None,
         outs_filter: Optional[Callable[["Output"], bool]] = None,
+        max_size: Optional[int] = None,
+        types: Optional[List[str]] = None,
         **kwargs: Any,
     ) -> "IndexView":
         """Return read-only view of index for the specified targets.
@@ -514,7 +578,20 @@ class Index:
             for stage_info in self.collect_targets(targets, **kwargs)
             if not stage_filter or stage_filter(stage_info.stage)
         ]
-        return IndexView(self, stage_infos, outs_filter=outs_filter)
+
+        def _outs_filter(out):
+            if max_size and out.meta and out.meta.size and out.meta.size >= max_size:
+                return False
+
+            if types and not self._types_filter(types, out):
+                return False
+
+            if outs_filter:
+                return outs_filter(out)
+
+            return True
+
+        return IndexView(self, stage_infos, outs_filter=_outs_filter)
 
 
 class _DataPrefixes(NamedTuple):

dvc/repo/open_repo.py

@@ -67,6 +67,7 @@ def open_repo(url, *args, **kwargs):
         url = os.getcwd()
 
     if os.path.exists(url):
+        url = os.path.abspath(url)
         try:
             config = _get_remote_config(url)
             config.update(kwargs.get("config") or {})

dvc/repo/pull.py

@@ -1,5 +1,6 @@
 import logging
 
+from dvc.exceptions import CheckoutError
 from dvc.repo import locked
 from dvc.utils import glob_targets
 
@@ -38,13 +39,17 @@ def pull(  # noqa: PLR0913
         recursive=recursive,
         run_cache=run_cache,
     )
-    stats = self.checkout(
-        targets=expanded_targets,
-        with_deps=with_deps,
-        force=force,
-        recursive=recursive,
-        allow_missing=allow_missing,
-    )
+    try:
+        stats = self.checkout(
+            targets=expanded_targets,
+            with_deps=with_deps,
+            force=force,
+            recursive=recursive,
+            allow_missing=allow_missing,
+        )
+    except CheckoutError as exc:
+        exc.stats["fetched"] = processed_files_count
+        raise
 
     stats["fetched"] = processed_files_count
     return stats

dvc/repo/reproduce.py

@@ -1,13 +1,18 @@
 import logging
-from typing import TYPE_CHECKING, Iterator, List
+from typing import TYPE_CHECKING, Iterable, List, Optional, Union, cast
+
+from funcy import ldistinct
 
 from dvc.exceptions import ReproductionError
 from dvc.repo.scm_context import scm_context
 from dvc.stage.cache import RunCacheNotSupported
+from dvc.utils.collections import ensure_list
 
 from . import locked
 
 if TYPE_CHECKING:
+    from networkx import DiGraph
+
     from dvc.stage import Stage
 
     from . import Repo
@@ -15,115 +20,62 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
-def _reproduce_stage(stage: "Stage", **kwargs) -> List["Stage"]:
-    if stage.frozen and not stage.is_import:
-        logger.warning(
-            "%s is frozen. Its dependencies are not going to be reproduced.",
-            stage,
-        )
-
-    stage = stage.reproduce(**kwargs)
-    if not stage:
-        return []
-
-    if not kwargs.get("dry", False):
-        stage.dump(update_pipeline=False)
-        _track_stage(stage)
-
-    return [stage]
+def collect_stages(
+    repo: "Repo",
+    targets: Iterable[str],
+    recursive: bool = False,
+    glob: bool = False,
+) -> List["Stage"]:
+    stages: List["Stage"] = []
+    for target in targets:
+        stages.extend(repo.stage.collect(target, recursive=recursive, glob=glob))
+    return ldistinct(stages)
 
 
-def _get_stage_files(stage: "Stage") -> Iterator[str]:
-    yield stage.dvcfile.relpath
-    for dep in stage.deps:
-        if (
-            not dep.use_scm_ignore
-            and dep.is_in_repo
-            and not stage.repo.dvcfs.isdvc(stage.repo.dvcfs.from_os_path(str(dep)))
-        ):
-            yield dep.fs_path
-    for out in stage.outs:
-        if not out.use_scm_ignore and out.is_in_repo:
-            yield out.fs_path
+def get_subgraph(
+    graph: "DiGraph",
+    nodes: Optional[List] = None,
+    pipeline: bool = False,
+    downstream: bool = False,
+) -> "DiGraph":
+    import networkx as nx
+
+    from .graph import get_pipeline, get_pipelines, get_subgraph_of_nodes
+
+    if not pipeline or not nodes:
+        return get_subgraph_of_nodes(graph, nodes, downstream=downstream)
+
+    pipelines = get_pipelines(graph)
+    used_pipelines = [get_pipeline(pipelines, node) for node in nodes]
+    return nx.compose_all(used_pipelines)
 
 
-def _track_stage(stage: "Stage") -> None:
-    from dvc.utils import relpath
-
-    context = stage.repo.scm_context
-    for path in _get_stage_files(stage):
-        context.track_file(relpath(path))
-    return context.track_changed_files()
+def _remove_frozen_stages(graph: "DiGraph") -> "DiGraph":
+    g = cast("DiGraph", graph.copy())
+    for stage in graph:
+        if stage.frozen:
+            # NOTE: disconnect frozen stage from its dependencies
+            g.remove_edges_from(graph.out_edges(stage))
+    return g
 
 
-@locked
-@scm_context
-def reproduce(  # noqa: C901, PLR0912
-    self: "Repo",
-    targets=None,
-    recursive=False,
-    pipeline=False,
-    all_pipelines=False,
-    **kwargs,
-):
-    from .graph import get_pipeline, get_pipelines
-
-    glob = kwargs.pop("glob", False)
-
-    if isinstance(targets, str):
-        targets = [targets]
-
-    if not all_pipelines and not targets:
-        from dvc.dvcfile import PROJECT_FILE
-
-        targets = [PROJECT_FILE]
-
-    interactive = kwargs.get("interactive", False)
-    if not interactive:
-        kwargs["interactive"] = self.config["core"].get("interactive", False)
-
-    stages = set()
-    if pipeline or all_pipelines:
-        pipelines = get_pipelines(self.index.graph)
-        if all_pipelines:
-            used_pipelines = pipelines
-        else:
-            used_pipelines = []
-            for target in targets:
-                stage = self.stage.get_target(target)
-                used_pipelines.append(get_pipeline(pipelines, stage))
-
-        for pline in used_pipelines:
-            for stage in pline:
-                if pline.in_degree(stage) == 0:
-                    stages.add(stage)
-    else:
-        for target in targets:
-            stages.update(
-                self.stage.collect(
-                    target,
-                    recursive=recursive,
-                    glob=glob,
-                )
-            )
-
-    if kwargs.get("pull", False) and kwargs.get("run_cache", True):
-        logger.debug("Pulling run cache")
-        try:
-            self.stage_cache.pull(None)
-        except RunCacheNotSupported as e:
-            logger.warning("Failed to pull run cache: %s", e)
-
-    return _reproduce_stages(self.index.graph, list(stages), **kwargs)
+def get_active_graph(
+    graph: "DiGraph",
+    stages: Optional[List["Stage"]] = None,
+    pipeline: bool = False,
+    downstream: bool = False,
+) -> "DiGraph":
+    """Return the graph to operate."""
+    processed = _remove_frozen_stages(graph)
+    return get_subgraph(processed, stages, pipeline=pipeline, downstream=downstream)
 
 
-def _reproduce_stages(  # noqa: C901
-    graph,
-    stages,
-    downstream=False,
-    single_item=False,
-    **kwargs,
-):
+def plan_repro(
+    graph: "DiGraph",
+    stages: Optional[List["Stage"]] = None,
+    pipeline: bool = False,
+    downstream: bool = False,
+) -> List["Stage"]:
     r"""Derive the evaluation of the given node for the given graph.
 
     When you _reproduce a stage_, you want to _evaluate the descendants_
@@ -159,72 +111,89 @@ def _reproduce_stages(  # noqa: C901
 
     The derived evaluation of _downstream_ B would be: [B, D, E]
     """
-    steps = _get_steps(graph, stages, downstream, single_item)
+    import networkx as nx
 
-    force_downstream = kwargs.pop("force_downstream", False)
+    active = get_active_graph(graph, stages, pipeline=pipeline, downstream=downstream)
+    return list(nx.dfs_postorder_nodes(active))
+
+
+def _reproduce_stage(stage: "Stage", **kwargs) -> Optional["Stage"]:
+    if stage.frozen and not stage.is_import:
+        logger.warning(
+            "%s is frozen. Its dependencies are not going to be reproduced.",
+            stage,
+        )
+
+    ret = stage.reproduce(**kwargs)
+    if ret and not kwargs.get("dry", False):
+        stage.dump(update_pipeline=False)
+    return ret
+
+
+def _reproduce_stages(
+    stages: List["Stage"],
+    force_downstream: bool = False,
+    **kwargs,
+) -> List["Stage"]:
     result: List["Stage"] = []
-    unchanged: List["Stage"] = []
-    # `ret` is used to add a cosmetic newline.
-    ret: List["Stage"] = []
-
-    for stage in steps:
-        if ret:
-            logger.info("")
-
+    for i, stage in enumerate(stages):
         try:
-            kwargs["upstream"] = result + unchanged
-            ret = _reproduce_stage(stage, **kwargs)
-            if not ret:
-                unchanged.extend([stage])
-                continue
-
-            result.extend(ret)
-            if force_downstream:
-                # NOTE: we are walking our pipeline from the top to the
-                # bottom. If one stage is changed, it will be reproduced,
-                # which tells us that we should force reproducing all of
-                # the other stages down below, even if their direct
-                # dependencies didn't change.
-                kwargs["force"] = True
+            ret = _reproduce_stage(stage, upstream=stages[:i], **kwargs)
         except Exception as exc:  # noqa: BLE001
             raise ReproductionError(stage.addressing) from exc
+
+        if not ret:
+            continue
+
+        result.append(ret)
+        if force_downstream:
+            # NOTE: we are walking our pipeline from the top to the
+            # bottom. If one stage is changed, it will be reproduced,
+            # which tells us that we should force reproducing all of
+            # the other stages down below, even if their direct
+            # dependencies didn't change.
+            kwargs["force"] = True
+        if i < len(stages) - 1:
+            logger.info("")  # add a newline
     return result
 
 
-def _get_steps(graph, stages, downstream, single_item):
-    import networkx as nx
+@locked
+@scm_context
+def reproduce(
+    self: "Repo",
+    targets: Union[Iterable[str], str, None] = None,
+    recursive: bool = False,
+    pipeline: bool = False,
+    all_pipelines: bool = False,
+    downstream: bool = False,
+    single_item: bool = False,
+    glob: bool = False,
+    **kwargs,
+):
+    from dvc.dvcfile import PROJECT_FILE
 
-    active = graph.copy()
+    if all_pipelines or pipeline:
+        single_item = False
+        downstream = False
+
+    if not kwargs.get("interactive", False):
+        kwargs["interactive"] = self.config["core"].get("interactive", False)
+
+    stages: List["Stage"] = []
+    if not all_pipelines:
+        targets_list = ensure_list(targets or PROJECT_FILE)
+        stages = collect_stages(self, targets_list, recursive=recursive, glob=glob)
+
+    if kwargs.get("pull", False) and kwargs.get("run_cache", True):
+        logger.debug("Pulling run cache")
+        try:
+            self.stage_cache.pull(None)
+        except RunCacheNotSupported as e:
+            logger.warning("Failed to pull run cache: %s", e)
+
+    steps = stages
     if not single_item:
-        # NOTE: frozen stages don't matter for single_item
-        for stage in graph:
-            if stage.frozen:
-                # NOTE: disconnect frozen stage from its dependencies
-                active.remove_edges_from(graph.out_edges(stage))
-
-    all_pipelines: List["Stage"] = []
-    for stage in stages:
-        if downstream:
-            # NOTE (py3 only):
-            # Python's `deepcopy` defaults to pickle/unpickle the object.
-            # Stages are complex objects (with references to `repo`,
-            # `outs`, and `deps`) that cause struggles when you try
-            # to serialize them. We need to create a copy of the graph
-            # itself, and then reverse it, instead of using
-            # graph.reverse() directly because it calls `deepcopy`
-            # underneath -- unless copy=False is specified.
-            nodes = nx.dfs_postorder_nodes(active.reverse(copy=False), stage)
-            all_pipelines += reversed(list(nodes))
-        else:
-            all_pipelines += nx.dfs_postorder_nodes(active, stage)
-
-    steps = []
-    for stage in all_pipelines:
-        if stage not in steps:
-            # NOTE: order of steps still matters for single_item
-            if single_item and stage not in stages:
-                continue
-
-            steps.append(stage)
-
-    return steps
+        graph = self.index.graph
+        steps = plan_repro(graph, stages, pipeline=pipeline, downstream=downstream)
+    return _reproduce_stages(steps, **kwargs)

dvc/stage/utils.py

@@ -102,6 +102,10 @@ def check_no_externals(stage):
     link = format_link(
         "https://dvc.org/doc/user-guide/pipelines/external-dependencies-and-outputs"
     )
+    if stage.is_data_source:
+        link = format_link(
+            "https://dvc.org/doc/user-guide/data-management/importing-external-data"
+        )
     raise StageExternalOutputsError(
         f"Cached output(s) outside of DVC project: {str_outs}. "
         f"See {link} for more info."
@@ -280,3 +284,27 @@ def validate_kwargs(
         kwargs.pop("name", None)
 
     return kwargs
+
+
+def _get_stage_files(stage: "Stage") -> List[str]:
+    from dvc.dvcfile import ProjectFile
+    from dvc.utils import relpath
+
+    ret: List[str] = []
+    file = stage.dvcfile
+    ret.append(file.relpath)
+    if isinstance(file, ProjectFile):
+        ret.append(file._lockfile.relpath)  # pylint: disable=protected-access
+
+    for dep in stage.deps:
+        if (
+            not dep.use_scm_ignore
+            and dep.is_in_repo
+            and not stage.repo.dvcfs.isdvc(stage.repo.dvcfs.from_os_path(str(dep)))
+        ):
+            ret.append(relpath(dep.fs_path))
+
+    for out in stage.outs:
+        if not out.use_scm_ignore and out.is_in_repo:
+            ret.append(relpath(out.fs_path))
+    return ret

pyproject.toml

@@ -34,13 +34,14 @@ dependencies = [
     "configobj>=5.0.6",
     "distro>=1.3",
     "dpath<3,>=2.1.0",
-    "dvc-data>=2.3.2,<2.4.0",
+    "dvc-data>=2.5.0,<2.6.0",
     "dvc-http>=2.29.0",
     "dvc-render>=0.3.1,<1",
     "dvc-studio-client>=0.9.2,<1",
     "dvc-task>=0.3.0,<1",
     "flatten_dict<1,>=0.4.1",
-    "flufl.lock>=5",
+    # https://github.com/iterative/dvc/issues/9654
+    "flufl.lock>=5,<8",
     "funcy>=1.14",
     "grandalf<1,>=0.7",
     "hydra-core>=1.1",
@@ -297,11 +298,10 @@ ignore-names = ["M", "SCM"]
 "dvc/commands/**" = ["N806"]
 "dvc/testing/**" = ["ARG002"]
 "dvc/testing/benchmarks/**" = ["ARG001"]
-"scripts/**" = ["T201", "INP001", "TRY002"]
 
 [tool.ruff.pylint]
 max-args = 10
 
 [tool.bandit]
-exclude_dirs = ["tests", "scripts", "dvc/repo/experiments/utils.py"]
+exclude_dirs = ["tests", "dvc/repo/experiments/utils.py"]
 skips = ["B101"]

scripts/build_package.sh

@@ -1,15 +0,0 @@
-#!/bin/bash
-
-set -e
-set -x
-
-if [ ! -d "dvc" ]; then
-  echo "Please run this script from repository root" >&2
-  exit 1
-fi
-
-echo 'PKG = "pip"' >dvc/_build.py
-
-python -m pip install build twine
-python -m build
-twine check dist/*

tests/func/api/test_data.py

@@ -198,11 +198,13 @@ def test_get_url_subrepos(tmp_dir, scm, local_cloud):
         local_cloud / "files" / "md5" / "ac" / "bd18db4cc2f85cedef654fccc4a4d8"
     )
     assert api.get_url(os.path.join("subrepo", "dir", "foo")) == expected_url
+    assert api.get_url(os.path.join("subrepo", "dir", "foo"), repo=".") == expected_url
 
     expected_url = os.fspath(
         local_cloud / "files" / "md5" / "37" / "b51d194a7513e45b56f6524f2d51f2"
     )
     assert api.get_url("subrepo/bar") == expected_url
+    assert api.get_url("subrepo/bar", repo=".") == expected_url
 
 
 def test_open_from_remote(tmp_dir, erepo_dir, cloud, local_cloud):

tests/func/repro/test_repro.py

@@ -1282,3 +1282,17 @@ def test_repro_rm_recursive(tmp_dir, dvc):
     dvc.reproduce()
     assert (tmp_dir / "dir").exists()
     assert not (tmp_dir / "dir" / "foo").exists()
+
+
+def test_repro_single_item_with_multiple_targets(tmp_dir, dvc, copy_script):
+    stage1 = dvc.stage.add(cmd="echo foo > foo", outs=["foo"], name="gen-foo")
+    with dvc.lock:
+        stage1.run()
+
+    stage2 = dvc.stage.add(
+        cmd="python copy.py foo bar", deps=["foo"], outs=["bar"], name="copy-foo-bar"
+    )
+    assert dvc.reproduce(["copy-foo-bar", "gen-foo"], single_item=True) == [
+        stage2,
+        stage1,
+    ]

tests/unit/command/test_data_sync.py

@@ -18,6 +18,12 @@ def test_fetch(mocker, dvc):
             "--with-deps",
             "--recursive",
             "--run-cache",
+            "--max-size",
+            "10",
+            "--type",
+            "plots",
+            "--type",
+            "metrics",
         ]
     )
     assert cli_args.func == CmdDataFetch
@@ -37,6 +43,8 @@ def test_fetch(mocker, dvc):
         with_deps=True,
         recursive=True,
         run_cache=True,
+        max_size=10,
+        types=["plots", "metrics"],
     )
 
 

tests/unit/repo/test_graph.py

@@ -0,0 +1,38 @@
+import pytest
+from networkx import DiGraph
+from networkx.utils import graphs_equal
+
+from dvc.repo.graph import get_subgraph_of_nodes
+
+
+@pytest.mark.parametrize(
+    "nodes,downstream,expected_edges",
+    [
+        ([], False, {1: [2, 3], 2: [4, 5], 3: [6, 7], 8: [9]}),
+        ([1], False, {1: [2, 3], 2: [4, 5], 3: [6, 7]}),
+        ([2], False, {2: [4, 5]}),
+        ([3], False, {3: [6, 7]}),
+        ([8], False, [(8, 9)]),
+        ([2, 3, 8], False, {2: [4, 5], 3: [6, 7], 8: [9]}),
+        ([4], False, {4: []}),
+        ([], True, {1: [2, 3], 2: [4, 5], 3: [6, 7], 8: [9]}),
+        ([1], True, {1: []}),
+        ([9], True, [(8, 9)]),
+        ([2], True, [(1, 2)]),
+        ([6], True, [(1, 3), (3, 6)]),
+        ([2, 3, 8], True, {1: [2, 3], 8: []}),
+        ([4, 7], True, {1: [2, 3], 2: [4], 3: [7]}),
+    ],
+)
+def test_subgraph_of_nodes(nodes, downstream, expected_edges):
+    r"""
+             1
+           /   \
+          2     3      8
+         / \   / \     |
+        4   5 6   7    9
+    """
+    graph = DiGraph({1: [2, 3], 2: [4, 5], 3: [6, 7], 8: [9]})
+    subgraph = get_subgraph_of_nodes(graph, nodes, downstream=downstream)
+    expected = DiGraph(expected_edges)
+    assert graphs_equal(expected, subgraph)

tests/unit/repo/test_reproduce.py

@@ -1,6 +1,8 @@
-import os
+from itertools import chain
 
-from dvc.repo.reproduce import _get_stage_files
+from networkx import DiGraph
+
+from dvc.repo.reproduce import plan_repro
 
 
 def test_number_reproduces(tmp_dir, dvc, mocker):
@@ -21,40 +23,48 @@ def test_number_reproduces(tmp_dir, dvc, mocker):
     assert reproduce_stage_mock.call_count == 5
 
 
-def test_get_stage_files(tmp_dir, dvc):
-    tmp_dir.dvc_gen("dvc-dep", "dvc-dep")
-    tmp_dir.gen("other-dep", "other-dep")
+def test_repro_plan(mocker):
+    r"""
+             n1
+           /   \
+          n2    m3    n8
+         / \   / \     |
+        n4 n5 n6 n7    n9
+    """
 
-    stage = dvc.stage.add(
-        name="stage",
-        cmd="foo",
-        deps=["dvc-dep", "other-dep"],
-        outs=["dvc-out"],
-        outs_no_cache=["other-out"],
-    )
-    result = set(_get_stage_files(stage))
-    assert result == {
-        stage.dvcfile.relpath,
-        str(tmp_dir / "other-dep"),
-        str(tmp_dir / "other-out"),
-    }
+    # note: downstream steps may not be stable, use AnyOf in such cases
+    class AnyOf:
+        def __init__(self, *items):
+            self.items = items
 
+        def __eq__(self, other: object) -> bool:
+            return any(item == other for item in self.items)
 
-def test_get_stage_files_wdir(tmp_dir, dvc):
-    tmp_dir.gen({"dir": {"dvc-dep": "dvc-dep", "other-dep": "other-dep"}})
-    dvc.add(os.path.join("dir", "dvc-dep"))
+    n = mocker.sentinel
+    n1, n2, n3, n4, n5, n6, n7, n8, n9 = (getattr(n, f"n{i}") for i in range(1, 10))
+    edges = {n1: [n2, n3], n2: [n4, n5], n3: [n6, n7], n8: [n9]}
+    for node in chain.from_iterable([n, *v] for n, v in edges.items()):
+        node.frozen = False
 
-    stage = dvc.stage.add(
-        name="stage",
-        cmd="foo",
-        wdir="dir",
-        deps=["dvc-dep", "other-dep"],
-        outs=["dvc-out"],
-        outs_no_cache=["other-out"],
-    )
-    result = set(_get_stage_files(stage))
-    assert result == {
-        stage.dvcfile.relpath,
-        str(tmp_dir / "dir" / "other-dep"),
-        str(tmp_dir / "dir" / "other-out"),
-    }
+    g = DiGraph(edges)
+    assert plan_repro(g) == [n4, n5, n2, n6, n7, n3, n1, n9, n8]
+    assert plan_repro(g, [n1]) == [n4, n5, n2, n6, n7, n3, n1]
+    assert plan_repro(g, [n4], downstream=True) == [n4, n2, n1]
+    assert plan_repro(g, [n8], True) == plan_repro(g, [n9], True) == [n9, n8]
+    assert plan_repro(g, [n2, n8], True) == [n4, n5, n2, n6, n7, n3, n1, n9, n8]
+    assert plan_repro(g, [n2, n3], downstream=True) == [
+        AnyOf(n2, n3),
+        AnyOf(n2, n3),
+        n1,
+    ]
+
+    n2.frozen = True
+    assert plan_repro(g) == [n2, n6, n7, n3, n1, n9, n8, n4, n5]
+    assert plan_repro(g, [n1]) == [n2, n6, n7, n3, n1]
+    assert plan_repro(g, [n4], downstream=True) == [n4]
+    assert plan_repro(g, [n2, n8], pipeline=True) == [n2, n6, n7, n3, n1, n9, n8]
+    assert plan_repro(g, [n2, n3], downstream=True) == [
+        AnyOf(n2, n3),
+        AnyOf(n2, n3),
+        n1,
+    ]

tests/unit/stage/test_utils.py

@@ -1,7 +1,7 @@
 import os
 
 from dvc.fs import localfs
-from dvc.stage.utils import resolve_paths
+from dvc.stage.utils import _get_stage_files, resolve_paths
 
 
 def test_resolve_paths():
@@ -19,3 +19,40 @@ def test_resolve_paths():
     path, wdir = resolve_paths(fs=localfs, path=file_path, wdir="../../some-dir")
     assert path == os.path.abspath(file_path)
     assert wdir == os.path.abspath("some-dir")
+
+
+def test_get_stage_files(tmp_dir, dvc):
+    tmp_dir.dvc_gen("dvc-dep", "dvc-dep")
+    tmp_dir.gen("other-dep", "other-dep")
+    stage = dvc.stage.create(
+        name="stage",
+        cmd="foo",
+        deps=["dvc-dep", "other-dep"],
+        outs=["dvc-out"],
+        outs_no_cache=["other-out"],
+    )
+    assert _get_stage_files(stage) == [
+        "dvc.yaml",
+        "dvc.lock",
+        "other-dep",
+        "other-out",
+    ]
+
+
+def test_get_stage_files_wdir(tmp_dir, dvc):
+    tmp_dir.gen({"dir": {"dvc-dep": "dvc-dep", "other-dep": "other-dep"}})
+    dvc.add(os.path.join("dir", "dvc-dep"))
+    stage = dvc.stage.create(
+        name="stage",
+        cmd="foo",
+        wdir="dir",
+        deps=["dvc-dep", "other-dep"],
+        outs=["dvc-out"],
+        outs_no_cache=["other-out"],
+    )
+    assert _get_stage_files(stage) == [
+        "dvc.yaml",
+        "dvc.lock",
+        os.path.join("dir", "other-dep"),
+        os.path.join("dir", "other-out"),
+    ]