From 1d3524f14403363046048ea4aa348042fa4006f4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Peter=20Rowlands=20=28=EB=B3=80=EA=B8=B0=ED=98=B8=29?= Date: Thu, 17 Jun 2021 20:41:33 +0900 Subject: [PATCH] odb: treat external repo dependencies as objects (#6109) * objects: add ExternalRepoFile object type * RepoDependency: implement get_obj (use ExternalRepoFile) * drop distinction between used external and used objects * update tests * remove unneeded fetch from ExternalRepoFile * wrap erepo errors as an object error * update cloud/remote functions for new used_objs behavior * fix update() rev/locked behavior * update import tests * remove ExternalRepoFile * move get_obj and get_used_objs implementations back into repo dependency * return objects plus associated ODB (remote) in get_used_objs * odb: add tmp_dir field and support creating Remote from an odb * return dict mapping odb -> objects in get_used_objs * use dummy git odb for git imports * fetch: support pulling from specific odb * update cloud funcs for get_used_objs usage * update tests for used_objs behavior * use fetch in dep.download * catch circular imports from local fs repos * add tests for chained and circular imports * move tmp_dir into odb config * update return docstring for repo.used_objs --- dvc/dependency/repo.py | 159 +++++++++++++++++++++---------- dvc/exceptions.py | 8 ++ dvc/fs/repo.py | 6 ++ dvc/objects/__init__.py | 14 ++- dvc/objects/db/base.py | 17 ++++ dvc/objects/db/git.py | 24 +++++ dvc/output.py | 32 ++++--- dvc/remote/base.py | 17 ++++ dvc/repo/__init__.py | 39 ++++---- dvc/repo/fetch.py | 122 ++++++++++++++---------- dvc/repo/gc.py | 6 +- dvc/repo/push.py | 10 +- dvc/repo/status.py | 32 ++++--- dvc/stage/__init__.py | 23 ++--- dvc/stage/cache.py | 19 ++-- tests/func/test_data_cloud.py | 2 +- tests/func/test_import.py | 54 ++++++++++- tests/unit/output/test_output.py | 5 +- tests/unit/repo/test_repo.py | 18 ++-- 19 files changed, 410 insertions(+), 197 deletions(-) create mode 100644 dvc/objects/db/git.py diff --git a/dvc/dependency/repo.py b/dvc/dependency/repo.py index df65a4dbd..3e5d81134 100644 --- a/dvc/dependency/repo.py +++ b/dvc/dependency/repo.py @@ -1,4 +1,6 @@ -from typing import NamedTuple, Optional +import os +from collections import defaultdict +from typing import TYPE_CHECKING, Dict, Optional, Set from voluptuous import Required @@ -6,10 +8,9 @@ from dvc.path_info import PathInfo from .base import Dependency - -class RepoPair(NamedTuple): - url: str - rev: Optional[str] = None +if TYPE_CHECKING: + from dvc.objects.db.base import ObjectDB + from dvc.objects.file import HashFile class RepoDependency(Dependency): @@ -28,6 +29,7 @@ class RepoDependency(Dependency): def __init__(self, def_repo, stage, *args, **kwargs): self.def_repo = def_repo + self._staged_objs: Dict[str, "HashFile"] = {} super().__init__(stage, *args, **kwargs) def _parse_path(self, fs, path_info): @@ -37,37 +39,12 @@ class RepoDependency(Dependency): def is_in_repo(self): return False - @property - def repo_pair(self) -> RepoPair: - d = self.def_repo - rev = d.get(self.PARAM_REV_LOCK) or d.get(self.PARAM_REV) - return RepoPair(d[self.PARAM_URL], rev) - def __str__(self): return "{} ({})".format(self.def_path, self.def_repo[self.PARAM_URL]) - def _make_repo(self, *, locked=True, **kwargs): - from dvc.external_repo import external_repo - - d = self.def_repo - rev = (d.get("rev_lock") if locked else None) or d.get("rev") - return external_repo(d["url"], rev=rev, **kwargs) - - def _get_hash(self, locked=True): - from dvc.objects.stage import stage - - with self._make_repo(locked=locked) as repo: - path_info = PathInfo(repo.root_dir) / self.def_path - return stage( - self.repo.odb.local, - path_info, - repo.repo_fs, - self.repo.odb.local.fs.PARAM_CHECKSUM, - ).hash_info - def workspace_status(self): - current = self._get_hash(locked=True) - updated = self._get_hash(locked=False) + current = self.get_obj(locked=True).hash_info + updated = self.get_obj(locked=False).hash_info if current != updated: return {str(self): "update available"} @@ -85,31 +62,21 @@ class RepoDependency(Dependency): def download(self, to, jobs=None): from dvc.checkout import checkout - from dvc.config import NoRemoteError - from dvc.exceptions import NoOutputOrStageError from dvc.objects import save - from dvc.objects.stage import stage + from dvc.objects.db.git import GitObjectDB + from dvc.repo.fetch import fetch_from_odb - odb = self.repo.odb.local - - with self._make_repo(cache_dir=odb.cache_dir) as repo: - if self.def_repo.get(self.PARAM_REV_LOCK) is None: - self.def_repo[self.PARAM_REV_LOCK] = repo.get_rev() - path_info = PathInfo(repo.root_dir) / self.def_path - try: - repo.fetch([path_info.fspath], jobs=jobs, recursive=True) - except (NoOutputOrStageError, NoRemoteError): - pass - obj = stage( - odb, path_info, repo.repo_fs, odb.fs.PARAM_CHECKSUM, jobs=jobs - ) - save(odb, obj, jobs=jobs) + for odb, objs in self.get_used_objs().items(): + if not isinstance(odb, GitObjectDB): + fetch_from_odb(self.repo, odb, objs, jobs=jobs) + obj = self.get_obj() + save(self.repo.odb.local, obj, jobs=jobs) checkout( to.path_info, to.fs, obj, - odb, + self.repo.odb.local, dvcignore=None, state=self.repo.state, ) @@ -117,7 +84,6 @@ class RepoDependency(Dependency): def update(self, rev=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() @@ -126,3 +92,94 @@ class RepoDependency(Dependency): # origin project url and rev_lock, and it makes RepoDependency # immutable, hence its impossible for checksum to change. return False + + def get_used_objs( + self, **kwargs + ) -> Dict[Optional["ObjectDB"], Set["HashFile"]]: + from dvc.config import NoRemoteError + from dvc.exceptions import NoOutputOrStageError + from dvc.objects.db.git import GitObjectDB + from dvc.objects.stage import stage + + local_odb = self.repo.odb.local + locked = kwargs.pop("locked", True) + with self._make_repo( + locked=locked, cache_dir=local_odb.cache_dir + ) as repo: + used_objs = defaultdict(set) + rev = repo.get_rev() + if locked and self.def_repo.get(self.PARAM_REV_LOCK) is None: + self.def_repo[self.PARAM_REV_LOCK] = rev + + path_info = PathInfo(repo.root_dir) / str(self.def_path) + try: + for odb, objs in repo.used_objs( + [os.fspath(path_info)], + force=True, + jobs=kwargs.get("jobs"), + recursive=True, + ).items(): + if odb is None: + odb = repo.cloud.get_remote().odb + self._check_circular_import(odb) + used_objs[odb].update(objs) + except (NoRemoteError, NoOutputOrStageError): + pass + + staged_obj = stage( + local_odb, + path_info, + repo.repo_fs, + local_odb.fs.PARAM_CHECKSUM, + ) + self._staged_objs[rev] = staged_obj + git_odb = GitObjectDB(repo.repo_fs, repo.root_dir) + used_objs[git_odb].add(staged_obj) + return used_objs + + def _check_circular_import(self, odb): + from dvc.exceptions import CircularImportError + from dvc.fs.repo import RepoFileSystem + + if not odb or not isinstance(odb.fs, RepoFileSystem): + return + + self_url = self.repo.url or self.repo.root_dir + if odb.fs.repo_url is not None and odb.fs.repo_url == self_url: + raise CircularImportError(self, odb.fs.repo_url, self_url) + + def get_obj(self, filter_info=None, **kwargs): + from dvc.objects.stage import stage + + odb = self.repo.odb.local + locked = kwargs.pop("locked", True) + with self._make_repo(locked=locked, cache_dir=odb.cache_dir) as repo: + rev = repo.get_rev() + if locked and self.def_repo.get(self.PARAM_REV_LOCK) is None: + self.def_repo[self.PARAM_REV_LOCK] = rev + obj = self._staged_objs.get(rev) + if obj is not None: + return obj + + path_info = PathInfo(repo.root_dir) / str(self.def_path) + obj = stage( + odb, + path_info, + repo.repo_fs, + odb.fs.PARAM_CHECKSUM, + ) + self._staged_objs[rev] = obj + return obj + + def _make_repo(self, locked=True, **kwargs): + 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): + d = self.def_repo + return (d.get(self.PARAM_REV_LOCK) if locked else None) or d.get( + self.PARAM_REV + ) diff --git a/dvc/exceptions.py b/dvc/exceptions.py index c64847d93..2c08974f3 100644 --- a/dvc/exceptions.py +++ b/dvc/exceptions.py @@ -384,3 +384,11 @@ class CacheLinkError(DvcException): ) super().__init__(msg) self.path_infos = path_infos + + +class CircularImportError(DvcException): + def __init__(self, dep, a, b): + super().__init__( + f"'{dep}' contains invalid circular import. " + f"DVC repo '{a}' already imports from '{b}'." + ) diff --git a/dvc/fs/repo.py b/dvc/fs/repo.py index 634900591..69e3f456d 100644 --- a/dvc/fs/repo.py +++ b/dvc/fs/repo.py @@ -64,6 +64,12 @@ class RepoFileSystem(BaseFileSystem): # pylint:disable=abstract-method if hasattr(repo, "dvc_dir"): self._dvcfss[repo.root_dir] = DvcFileSystem(repo=repo) + @property + def repo_url(self): + if self._main_repo is None: + return None + return self._main_repo.url + def _get_repo(self, path: str) -> Optional["Repo"]: """Returns repo that the path falls in, using prefix. diff --git a/dvc/objects/__init__.py b/dvc/objects/__init__.py index fafacf339..f464ccb4b 100644 --- a/dvc/objects/__init__.py +++ b/dvc/objects/__init__.py @@ -1,14 +1,24 @@ import logging from concurrent.futures import ThreadPoolExecutor, as_completed +from typing import TYPE_CHECKING, Optional from dvc.progress import Tqdm from .tree import Tree +if TYPE_CHECKING: + from .db.base import ObjectDB + from .file import HashFile + logger = logging.getLogger(__name__) -def save(odb, obj, jobs=None, **kwargs): +def save( + odb: "ObjectDB", + obj: "HashFile", + jobs: Optional[int] = None, + **kwargs, +): if isinstance(obj, Tree): with ThreadPoolExecutor(max_workers=jobs) as executor: for future in Tqdm( @@ -18,7 +28,7 @@ def save(odb, obj, jobs=None, **kwargs): entry.path_info, entry.fs, entry.hash_info, - **kwargs + **kwargs, ) for _, entry in obj ), diff --git a/dvc/objects/db/base.py b/dvc/objects/db/base.py index 08db92cff..c0419bc54 100644 --- a/dvc/objects/db/base.py +++ b/dvc/objects/db/base.py @@ -29,6 +29,23 @@ class ObjectDB: self.cache_types = config.get("type") or copy(self.DEFAULT_CACHE_TYPES) self.cache_type_confirmed = False self.slow_link_warning = config.get("slow_link_warning", True) + self.tmp_dir = config.get("tmp_dir") + + @property + def config(self): + return { + "state": self.state, + "verify": self.verify, + "type": self.cache_types, + "slow_link_warning": self.slow_link_warning, + "tmp_dir": self.tmp_dir, + } + + def __eq__(self, other): + return self.fs == other.fs and self.path_info == other.path_info + + def __hash__(self): + return hash((self.fs.scheme, self.path_info)) def move(self, from_info, to_info): self.fs.move(from_info, to_info) diff --git a/dvc/objects/db/git.py b/dvc/objects/db/git.py new file mode 100644 index 000000000..f5869d5dd --- /dev/null +++ b/dvc/objects/db/git.py @@ -0,0 +1,24 @@ +import logging + +from .base import ObjectDB + +logger = logging.getLogger(__name__) + + +class GitObjectDB(ObjectDB): + """Dummy read-only ODB for uncached objects in external Git repos.""" + + def __init__(self, fs, path_info, **config): + from dvc.fs.repo import RepoFileSystem + + assert isinstance(fs, RepoFileSystem) + super().__init__(fs, path_info) + + def get(self, hash_info): + raise NotImplementedError + + def add(self, path_info, fs, hash_info, move=True, **kwargs): + raise NotImplementedError + + def gc(self, used, jobs=None): + raise NotImplementedError diff --git a/dvc/output.py b/dvc/output.py index 4f1fbc40b..07be26acc 100644 --- a/dvc/output.py +++ b/dvc/output.py @@ -2,7 +2,7 @@ import logging import os from collections import defaultdict from copy import copy -from typing import TYPE_CHECKING, Dict, Set, Type +from typing import TYPE_CHECKING, Dict, Optional, Set, Type from urllib.parse import urlparse from funcy import collecting, project @@ -33,8 +33,7 @@ from .utils import relpath from .utils.fs import path_isin if TYPE_CHECKING: - from dvc.dependency.repo import RepoPair - + from .objects.db.base import ObjectDB from .objects.file import HashFile logger = logging.getLogger(__name__) @@ -324,6 +323,8 @@ class Output: self.obj = None self.isexec = False if self.IS_DEPENDENCY else isexec + self.def_remote = None + def _parse_path(self, fs, path_info): if fs.scheme != "local": return path_info @@ -809,7 +810,7 @@ class Output: def collect_used_dir_cache( self, remote=None, force=False, jobs=None, filter_info=None - ) -> Set["HashFile"]: + ) -> Dict[Optional["ObjectDB"], Set["HashFile"]]: """Fetch dir cache and return used objects for this out.""" try: @@ -830,17 +831,22 @@ class Output: "unable to fully collect used cache" " without cache for directory '{}'".format(self) ) - return set() + return {} obj = self.get_obj(filter_info=filter_info, copy=True) self._set_obj_names(obj) - return {obj} + return {None: {obj}} - def get_used_objs(self, **kwargs) -> Set["HashFile"]: + def get_used_objs( + self, **kwargs + ) -> Dict[Optional["ObjectDB"], Set["HashFile"]]: """Return filtered set of used objects for this out.""" if not self.use_cache: - return set() + return {} + + if self.stage.is_repo_import: + return self.get_used_external(**kwargs) if not self.hash_info: msg = ( @@ -859,7 +865,7 @@ class Output: ) ) logger.warning(msg) - return set() + return {} if self.is_dir_checksum: return self.collect_used_dir_cache(**kwargs) @@ -869,7 +875,7 @@ class Output: obj = self.odb.get(self.hash_info) self._set_obj_names(obj) - return {obj} + return {None: {obj}} def _set_obj_names(self, obj): obj.name = str(self) @@ -877,12 +883,14 @@ class Output: for key, entry_obj in obj: entry_obj.name = os.path.join(str(self), *key) - def get_used_external(self, **kwargs) -> Dict["RepoPair", str]: + def get_used_external( + self, **kwargs + ) -> Dict[Optional["ObjectDB"], Set["HashFile"]]: if not self.use_cache or not self.stage.is_repo_import: return {} (dep,) = self.stage.deps - return {dep.repo_pair: dep.def_path} + return dep.get_used_objs() def _validate_output_path(self, path, stage=None): from dvc.dvcfile import is_valid_filename diff --git a/dvc/remote/base.py b/dvc/remote/base.py index ba5e39ca6..1d211484c 100644 --- a/dvc/remote/base.py +++ b/dvc/remote/base.py @@ -72,6 +72,7 @@ class Remote: from dvc.objects.db import get_odb self.fs = fs + config["tmp_dir"] = tmp_dir self.odb = get_odb(self.fs, path_info, **config) url = getattr(path_info, "fspath", path_info.url) @@ -86,6 +87,22 @@ class Remote: path_info=self.odb.path_info or "No path", ) + def __eq__(self, other): + return self.odb == other.odb + + def __hash__(self): + return hash(self.odb) + + @classmethod + def from_odb(cls, odb): + config = odb.config + return cls( + odb.fs, + odb.path_info, + config.pop("tmp_dir"), + **config, + ) + @index_locked def gc(self, *args, **kwargs): removed = self.odb.gc(*args, **kwargs) diff --git a/dvc/repo/__init__.py b/dvc/repo/__init__.py index e44337861..00405c5e2 100644 --- a/dvc/repo/__init__.py +++ b/dvc/repo/__init__.py @@ -341,7 +341,7 @@ class Repo: existing_stages = self.stages if old_stages is None else old_stages build_graph(existing_stages + new_stages) - def used_cache( + def used_objs( self, targets=None, all_branches=False, @@ -367,21 +367,22 @@ class Repo: the scope. Returns: - A dictionary with Schemes (representing output's location) mapped - to items containing the output's `dumpd` names and the output's - children (if the given output is a directory). + A dict mapping (remote) ODB instances to sets of objects that + belong to each ODB. If the ODB instance is None, the objects + are naive and do not belong to a specific remote ODB. """ - used_objs = set() - used_external = defaultdict(set) + used = defaultdict(set) def _add_suffix(objs, suffix): from dvc.objects.tree import Tree for obj in objs: - obj.name += suffix + if obj.name is not None: + obj.name += suffix if isinstance(obj, Tree): for _, entry_obj in obj: - entry_obj.name += suffix + if entry_obj.name is not None: + entry_obj.name += suffix for branch in self.brancher( revs=revs, @@ -400,27 +401,23 @@ class Repo: ) for stage, filter_info in pairs: - objs, external = stage.get_used_cache( + for odb, objs in stage.get_used_objs( remote=remote, force=force, jobs=jobs, filter_info=filter_info, - ) - if branch: - _add_suffix(objs, f" ({branch})") - used_objs.update(objs) - for repo_pair, paths in external.items(): - used_external[repo_pair].update(paths) + ).items(): + if branch: + _add_suffix(objs, f" ({branch})") + used[odb].update(objs) if used_run_cache: - objs, external = self.stage_cache.get_used_cache( + for odb, objs in self.stage_cache.get_used_objs( used_run_cache, remote=remote, force=force, jobs=jobs - ) - used_objs.update(objs) - for repo_pair, paths in external.items(): - used_external[repo_pair].update(paths) + ).items(): + used[odb].update(objs) - return used_objs, used_external + return used @cached_property def outs_trie(self): diff --git a/dvc/repo/fetch.py b/dvc/repo/fetch.py index 5c1c9e277..688a95345 100644 --- a/dvc/repo/fetch.py +++ b/dvc/repo/fetch.py @@ -1,8 +1,7 @@ import logging -import os from dvc.config import NoRemoteError -from dvc.exceptions import DownloadError, NoOutputOrStageError +from dvc.exceptions import DownloadError from ..scheme import Schemes from . import locked @@ -37,11 +36,12 @@ def fetch( config.NoRemoteError: thrown when downloading only local files and no remote is configured """ + from dvc.objects.db.git import GitObjectDB if isinstance(targets, str): targets = [targets] - used_objs, used_external = self.used_cache( + used = self.used_objs( targets, all_branches=all_branches, all_tags=all_tags, @@ -60,22 +60,39 @@ def fetch( try: if run_cache: self.stage_cache.pull(remote) - downloaded += self.cloud.pull( - used_objs, - jobs, - remote=remote, - show_checksums=show_checksums, - ) - except NoRemoteError: - if not used_external and any( - obj.fs.scheme == Schemes.LOCAL for obj in used_objs - ): - raise except DownloadError as exc: failed += exc.amount - for repo_pair, files in used_external.items(): - d, f = _fetch_external(self, repo_pair.url, repo_pair.rev, files, jobs) + external = set() + for odb, objs in used.items(): + if odb is None: + # objs contains naive objects to be pulled from specified remote + d, f = _fetch_naive_objs( + self, + objs, + jobs=jobs, + remote=remote, + show_checksums=show_checksums, + ) + downloaded += d + failed += f + elif isinstance(odb, GitObjectDB): + # objs contains staged import objects which should be saved + # last (after all other objects have been pulled) + external.update(objs) + else: + d, f = fetch_from_odb( + self, + odb, + objs, + jobs=jobs, + show_checksums=show_checksums, + ) + downloaded += d + failed += f + + if external: + d, f = _fetch_external(self, external, jobs=jobs) downloaded += d failed += f @@ -85,44 +102,51 @@ def fetch( return downloaded -def _fetch_external(self, repo_url, repo_rev, files, jobs): - from dvc.external_repo import external_repo - from dvc.objects import save - from dvc.objects.stage import stage - from dvc.path_info import PathInfo - from dvc.scm.base import CloneError - +def _fetch_naive_objs(repo, objs, **kwargs): + # objs contains naive objects to be pulled from specified remote + downloaded = 0 failed = 0 + try: + downloaded += repo.cloud.pull(objs, **kwargs) + except NoRemoteError: + if any(obj.fs.scheme == Schemes.LOCAL for obj in objs): + raise + except DownloadError as exc: + failed += exc.amount + return downloaded, failed + + +def fetch_from_odb(repo, odb, objs, **kwargs): + from dvc.remote.base import Remote + + downloaded = 0 + failed = 0 + remote = Remote.from_odb(odb) + try: + downloaded += remote.pull( + repo.odb.local, + objs, + **kwargs, + ) + except DownloadError as exc: + failed += exc.amount + return downloaded, failed + + +def _fetch_external(repo, objs, **kwargs): + from dvc.objects import save + from dvc.objects.errors import ObjectError results = [] + failed = 0 - def cb(result): + def callback(result): results.append(result) - odb = self.odb.local - try: - with external_repo( - repo_url, repo_rev, cache_dir=odb.cache_dir - ) as repo: - root = PathInfo(repo.root_dir) - for path in files: - path_info = root / path - try: - used_objs, _ = repo.used_cache( - [os.fspath(path_info)], - force=True, - jobs=jobs, - recursive=True, - ) - cb(repo.cloud.pull(used_objs, jobs)) - except (NoOutputOrStageError, NoRemoteError): - pass - obj = stage(odb, path_info, repo.repo_fs, "md5", jobs=jobs) - save(odb, obj, jobs=jobs, download_callback=cb) - except CloneError: - failed += 1 - logger.exception( - "failed to fetch data for '{}'".format(", ".join(files)) - ) + for obj in objs: + try: + save(repo.odb.local, obj, download_callback=callback, **kwargs) + except ObjectError: + failed += 1 return sum(results), failed diff --git a/dvc/repo/gc.py b/dvc/repo/gc.py index 7cce7e8cc..46eb9e3fa 100644 --- a/dvc/repo/gc.py +++ b/dvc/repo/gc.py @@ -61,7 +61,7 @@ def gc( stack.enter_context(repo.lock) for repo in all_repos + [self]: - objs, _ = repo.used_cache( + for objs in repo.used_objs( all_branches=all_branches, with_deps=with_deps, all_tags=all_tags, @@ -70,8 +70,8 @@ def gc( remote=remote, force=force, jobs=jobs, - ) - used_objs.update(objs) + ).values(): + used_objs.update(objs) for scheme, odb in self.odb.by_scheme(): if not odb: diff --git a/dvc/repo/push.py b/dvc/repo/push.py index 532eede43..63afbcfa9 100644 --- a/dvc/repo/push.py +++ b/dvc/repo/push.py @@ -24,7 +24,7 @@ def push( expanded_targets = glob_targets(targets, glob=glob) - used_objs, _ = self.used_cache( + used = self.used_objs( expanded_targets, all_branches=all_branches, all_tags=all_tags, @@ -38,6 +38,8 @@ def push( revs=revs, ) - return len(used_run_cache) + self.cloud.push( - used_objs, jobs, remote=remote - ) + pushed = len(used_run_cache) + for odb, objs in used.items(): + if odb is None: + pushed += self.cloud.push(objs, jobs, remote=remote) + return pushed diff --git a/dvc/repo/status.py b/dvc/repo/status.py index 60f6dd111..3a304cc90 100644 --- a/dvc/repo/status.py +++ b/dvc/repo/status.py @@ -78,7 +78,7 @@ def _cloud_status( """ import dvc.remote.base as cloud - used_objs, _ = self.used_cache( + used = self.used_objs( targets, all_branches=all_branches, all_tags=all_tags, @@ -91,22 +91,26 @@ def _cloud_status( ) ret = {} - status_info = self.cloud.status( - used_objs, jobs, remote=remote, log_missing=False - ) - for info in status_info.values(): - name = info["name"] - status_ = info["status"] - if status_ == cloud.STATUS_OK: + for odb, objs in used.items(): + if odb is not None: + # ignore imported objects continue + status_info = self.cloud.status( + objs, jobs, remote=remote, log_missing=False + ) + for info in status_info.values(): + name = info["name"] + status_ = info["status"] + if status_ == cloud.STATUS_OK: + continue - prefix_map = { - cloud.STATUS_DELETED: "deleted", - cloud.STATUS_NEW: "new", - cloud.STATUS_MISSING: "missing", - } + prefix_map = { + cloud.STATUS_DELETED: "deleted", + cloud.STATUS_NEW: "new", + cloud.STATUS_MISSING: "missing", + } - ret[name] = prefix_map[status_] + ret[name] = prefix_map[status_] return ret diff --git a/dvc/stage/__init__.py b/dvc/stage/__init__.py index 0ba452793..8559b615a 100644 --- a/dvc/stage/__init__.py +++ b/dvc/stage/__init__.py @@ -4,7 +4,7 @@ import string from collections import defaultdict from contextlib import suppress from dataclasses import dataclass -from typing import TYPE_CHECKING, Dict, Optional, Set, Tuple +from typing import TYPE_CHECKING, Dict, Optional, Set from funcy import cached_property, project @@ -35,8 +35,8 @@ from .utils import ( ) if TYPE_CHECKING: - from dvc.dependency.repo import RepoPair from dvc.dvcfile import DVCFile + from dvc.objects.db.base import ObjectDB from dvc.objects.file import HashFile logger = logging.getLogger(__name__) @@ -624,20 +624,15 @@ class Stage(params.StageParams): for out in self.filter_outs(filter_info) ) - def get_used_cache( + def get_used_objs( self, *args, **kwargs - ) -> Tuple[Set["HashFile"], Dict["RepoPair", Set[str]]]: - """Return tuple of (used_objects, used_external) for this stage.""" - used_objs = set() - used_external = defaultdict(set) + ) -> Dict[Optional["ObjectDB"], Set["HashFile"]]: + """Return set of objects used by this stage.""" + used_objs = defaultdict(set) for out in self.filter_outs(kwargs.get("filter_info")): - used_objs.update(out.get_used_objs(*args, **kwargs)) - for repo_pair, path in out.get_used_external( - *args, **kwargs - ).items(): - used_external[repo_pair].add(path) - - return used_objs, used_external + for odb, objs in out.get_used_objs(*args, **kwargs).items(): + used_objs[odb].update(objs) + return used_objs @staticmethod def _check_can_merge(stage, ancestor_out=None): diff --git a/dvc/stage/cache.py b/dvc/stage/cache.py index c06e71e68..3d3b1d153 100644 --- a/dvc/stage/cache.py +++ b/dvc/stage/cache.py @@ -1,7 +1,6 @@ import logging import os import tempfile -from collections import defaultdict from contextlib import contextmanager from funcy import cached_property, first @@ -197,8 +196,8 @@ class StageCache: cached_stage = self._create_stage(cache, wdir=stage.wdir) if pull: - objs, _ = cached_stage.get_used_cache() - self.repo.cloud.pull(objs) + for objs in cached_stage.get_used_objs().values(): + self.repo.cloud.pull(objs) if not cached_stage.outs_cached(): raise RunCacheNotFoundError(stage) @@ -251,18 +250,16 @@ class StageCache: self.repo.odb.local, ) - def get_used_cache(self, used_run_cache, *args, **kwargs): + def get_used_objs(self, used_run_cache, *args, **kwargs): """Return used cache for the specified run-cached stages.""" - used_objs = set() - used_external = defaultdict(set) + from collections import defaultdict + used_objs = defaultdict(set) for key, value in used_run_cache: entry = self._load_cache(key, value) if not entry: continue stage = self._create_stage(entry) - objs, external = stage.get_used_cache(*args, **kwargs) - used_objs.update(objs) - for repo_pair, paths in external: - used_external[repo_pair].update(paths) - return used_objs, used_external + for odb, objs in stage.get_used_objs(*args, **kwargs).items(): + used_objs[odb].update(objs) + return used_objs diff --git a/tests/func/test_data_cloud.py b/tests/func/test_data_cloud.py index 9fa0c623c..cec4d3fcd 100644 --- a/tests/func/test_data_cloud.py +++ b/tests/func/test_data_cloud.py @@ -56,7 +56,7 @@ def test_cloud(tmp_dir, dvc, remote): # pylint:disable=unused-argument out = stage.outs[0] cache = out.cache_path md5 = out.hash_info.value - foo_objs = out.get_used_objs() + foo_objs = out.get_used_objs().get(None, set()) (stage_dir,) = tmp_dir.dvc_gen( { diff --git a/tests/func/test_import.py b/tests/func/test_import.py index a65c3aac1..1e68fe005 100644 --- a/tests/func/test_import.py +++ b/tests/func/test_import.py @@ -530,7 +530,7 @@ def test_import_with_no_exec(tmp_dir, dvc, erepo_dir): def test_import_with_jobs(mocker, dvc, erepo_dir): - from dvc.data_cloud import DataCloud + import dvc as dvc_module with erepo_dir.chdir(): erepo_dir.dvc_gen( @@ -545,7 +545,57 @@ def test_import_with_jobs(mocker, dvc, erepo_dir): commit="init", ) - spy = mocker.spy(DataCloud, "pull") + spy = mocker.spy(dvc_module.objects, "save") dvc.imp(os.fspath(erepo_dir), "dir1", jobs=3) run_jobs = tuple(spy.call_args_list[0])[1].get("jobs") assert run_jobs == 3 + + +def test_chained_import(tmp_dir, dvc, make_tmp_dir, erepo_dir, local_cloud): + erepo_dir.add_remote(config=local_cloud.config) + with erepo_dir.chdir(): + erepo_dir.dvc_gen({"dir": {"foo": "foo", "bar": "bar"}}, commit="init") + erepo_dir.dvc.push() + remove(erepo_dir.dvc.odb.local.cache_dir) + remove(os.fspath(erepo_dir / "dir")) + + erepo2 = make_tmp_dir("erepo2", scm=True, dvc=True) + with erepo2.chdir(): + erepo2.dvc.imp(os.fspath(erepo_dir), "dir") + erepo2.scm.add("dir.dvc") + erepo2.scm.commit("import") + remove(erepo2.dvc.odb.local.cache_dir) + remove(os.fspath(erepo2 / "dir")) + + dvc.imp(os.fspath(erepo2), "dir", "dir_imported") + dst = tmp_dir / "dir_imported" + assert (dst / "foo").read_text() == "foo" + assert (dst / "bar").read_text() == "bar" + + remove(dvc.odb.local.cache_dir) + remove("dir_imported") + + # pulled objects should come from the original upstream repo's remote, + # no cache or remote should be needed from the intermediate repo + dvc.pull(["dir_imported.dvc"]) + assert not os.path.exists(erepo_dir.dvc.odb.local.cache_dir) + assert not os.path.exists(erepo2.dvc.odb.local.cache_dir) + assert (dst / "foo").read_text() == "foo" + assert (dst / "bar").read_text() == "bar" + + +def test_circular_import(tmp_dir, dvc, scm, erepo_dir): + from dvc.exceptions import CircularImportError + + with erepo_dir.chdir(): + erepo_dir.dvc_gen({"dir": {"foo": "foo", "bar": "bar"}}, commit="init") + + dvc.imp(os.fspath(erepo_dir), "dir", "dir_imported") + scm.add("dir_imported.dvc") + scm.commit("import") + + with erepo_dir.chdir(): + with pytest.raises(CircularImportError): + erepo_dir.dvc.imp( + os.fspath(tmp_dir), "dir_imported", "circular_import" + ) diff --git a/tests/unit/output/test_output.py b/tests/unit/output/test_output.py index 197ac57ce..effe0c3fc 100644 --- a/tests/unit/output/test_output.py +++ b/tests/unit/output/test_output.py @@ -79,7 +79,7 @@ def test_checksum_schema_fail(value): ), ], ) -def test_get_used_cache(exists, expected_message, mocker, caplog): +def test_get_used_objs(exists, expected_message, mocker, caplog): stage = mocker.MagicMock() mocker.patch.object(stage, "__str__", return_value="stage: 'stage.dvc'") mocker.patch.object(stage, "addressing", "stage.dvc") @@ -99,6 +99,5 @@ def test_get_used_cache(exists, expected_message, mocker, caplog): ).return_value = exists with caplog.at_level(logging.WARNING, logger="dvc"): - assert set() == output.get_used_objs() - assert {} == output.get_used_external() + assert {} == output.get_used_objs() assert first(caplog.messages) == expected_message diff --git a/tests/unit/repo/test_repo.py b/tests/unit/repo/test_repo.py index 8a7401c43..316a5411a 100644 --- a/tests/unit/repo/test_repo.py +++ b/tests/unit/repo/test_repo.py @@ -44,26 +44,24 @@ def test_find_outs_by_path_does_graph_checks(tmp_dir, dvc): "path", [os.path.join("dir", "subdir", "file"), os.path.join("dir", "subdir")], ) -def test_used_cache(tmp_dir, dvc, path): +def test_used_objs(tmp_dir, dvc, path): from dvc.objects.tree import Tree tmp_dir.dvc_gen({"dir": {"subdir": {"file": "file"}, "other": "other"}}) - expected_objs = { + expected = { HashInfo("md5", "70922d6bf66eb073053a82f77d58c536.dir"), HashInfo("md5", "8c7dd922ad47494fc02c388e12c00eac"), } - expected_external = {} - objs, external = dvc.used_cache([path]) used = set() - for obj in objs: - used.add(obj.hash_info) - if isinstance(obj, Tree): - used.update(entry_obj.hash_info for _, entry_obj in obj) + for _, objs in dvc.used_objs([path]).items(): + for obj in objs: + used.add(obj.hash_info) + if isinstance(obj, Tree): + used.update(entry_obj.hash_info for _, entry_obj in obj) - assert used == expected_objs - assert external == expected_external + assert used == expected def test_locked(mocker):