From eb46c03c6c02e837c09d76cc3c96c889a84bafe8 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, 3 Jun 2021 04:19:58 +0900 Subject: [PATCH] output/remote: get rid of NamedCache (#6008) * objects: make HashInfo & HashFile hashable * stage/output: collect objects (HashFiles) instead of NamedCache * remote/data_cloud: use objs instead of NamedCache * repo: use objects instead of NamedCache * rerun black after rebase * remove NamedCache usage from run-cache * drop NamedCache * update HashInfo.__hash__ * update unit tests * revert remote_name change * start cleaning up test failures * fix tree obj expansion * update missing cache logging * fix external dict handling * update naive odb.gc * update state save behavior for pull * update test_data_cloud * update test_remote * fix gc with multiple repos * objects.tree: support filtering w/preserved hash_info * fix stage cache restore * drop fs.path_info from HashFile.__hash__ * objects: add optional name attribute * revert missing cache logging name/suffix changes * revert name changes in remote tests * review fixes * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update dvc/remote/base.py * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Co-authored-by: Ruslan Kuprieiev Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- dvc/data_cloud.py | 85 ++++++++++++++-------- dvc/dependency/repo.py | 11 ++- dvc/hash_info.py | 3 + dvc/objects/db/__init__.py | 98 ------------------------- dvc/objects/db/base.py | 11 ++- dvc/objects/file.py | 24 +++++- dvc/objects/tree.py | 35 +++++++-- dvc/output.py | 109 +++++++++++----------------- dvc/remote/base.py | 97 ++++++++++++++++--------- dvc/repo/__init__.py | 30 ++++++-- dvc/repo/fetch.py | 20 +++-- dvc/repo/gc.py | 38 ++++++---- dvc/repo/push.py | 6 +- dvc/repo/status.py | 4 +- dvc/stage/__init__.py | 22 ++++-- dvc/stage/cache.py | 17 +++-- tests/func/test_data_cloud.py | 60 +++++++++------ tests/func/test_remote.py | 20 +++-- tests/unit/objects/db/test_local.py | 18 +++-- tests/unit/output/test_output.py | 4 +- tests/unit/repo/test_repo.py | 34 ++++----- 21 files changed, 402 insertions(+), 344 deletions(-) diff --git a/dvc/data_cloud.py b/dvc/data_cloud.py index 8825d4510..b4a3d5a87 100644 --- a/dvc/data_cloud.py +++ b/dvc/data_cloud.py @@ -1,6 +1,11 @@ """Manages dvc remotes that user can use with push/pull/status commands.""" import logging +from typing import TYPE_CHECKING, Iterable, Optional + +if TYPE_CHECKING: + from dvc.objects.file import HashFile + from dvc.remote.base import Remote logger = logging.getLogger(__name__) @@ -19,7 +24,11 @@ class DataCloud: def __init__(self, repo): self.repo = repo - def get_remote(self, name=None, command=""): + def get_remote( + self, + name: Optional[str] = None, + command: str = "", + ) -> "Remote": from dvc.config import NoRemoteError if not name: @@ -48,76 +57,88 @@ class DataCloud: return get_remote(self.repo, name=name) - def push(self, cache, jobs=None, remote=None, show_checksums=False): + def push( + self, + objs: Iterable["HashFile"], + jobs: Optional[int] = None, + remote: Optional[str] = None, + show_checksums: bool = False, + ): """Push data items in a cloud-agnostic way. Args: - cache (NamedCache): named checksums to push to the cloud. - jobs (int): number of jobs that can be running simultaneously. - remote (dvc.remote.base.BaseRemote): optional remote to push to. + objs: objects to push to the cloud. + jobs: number of jobs that can be running simultaneously. + remote: optional remote to push to. By default remote from core.remote config option is used. - show_checksums (bool): show checksums instead of file names in + show_checksums: show checksums instead of file names in information messages. """ - remote = self.get_remote(remote, "push") + remote_obj = self.get_remote(remote, "push") - return remote.push( + return remote_obj.push( self.repo.odb.local, - cache, + objs, jobs=jobs, show_checksums=show_checksums, ) - def pull(self, cache, jobs=None, remote=None, show_checksums=False): + def pull( + self, + objs: Iterable["HashFile"], + jobs: Optional[int] = None, + remote: Optional[str] = None, + show_checksums: bool = False, + ): """Pull data items in a cloud-agnostic way. Args: - cache (NamedCache): named checksums to pull from the cloud. - jobs (int): number of jobs that can be running simultaneously. - remote (dvc.remote.base.BaseRemote): optional remote to pull from. + objs: objects to pull from the cloud. + jobs: number of jobs that can be running simultaneously. + remote: optional remote to pull from. By default remote from core.remote config option is used. - show_checksums (bool): show checksums instead of file names in + show_checksums: show checksums instead of file names in information messages. """ - remote = self.get_remote(remote, "pull") + remote_obj = self.get_remote(remote, "pull") - return remote.pull( + return remote_obj.pull( self.repo.odb.local, - cache, + objs, jobs=jobs, show_checksums=show_checksums, ) def status( self, - cache, - jobs=None, - remote=None, - show_checksums=False, - log_missing=True, + objs: Iterable["HashFile"], + jobs: Optional[int] = None, + remote: Optional[str] = None, + show_checksums: bool = False, + log_missing: bool = True, ): """Check status of data items in a cloud-agnostic way. Args: - cache (NamedCache): named checksums to check status for. - jobs (int): number of jobs that can be running simultaneously. - remote (dvc.remote.base.BaseRemote): optional remote to compare + objs: objects to check status for. + jobs: number of jobs that can be running simultaneously. + remote: optional remote to compare cache to. By default remote from core.remote config option is used. - show_checksums (bool): show checksums instead of file names in + show_checksums: show checksums instead of file names in information messages. - log_missing (bool): log warning messages if file doesn't exist + log_missing: log warning messages if file doesn't exist neither in cache, neither in cloud. """ - remote = self.get_remote(remote, "status") - return remote.status( + remote_obj = self.get_remote(remote, "status") + return remote_obj.status( self.repo.odb.local, - cache, + objs, jobs=jobs, show_checksums=show_checksums, log_missing=log_missing, ) def get_url_for(self, remote, checksum): - remote = self.get_remote(remote) - return str(remote.odb.hash_to_path_info(checksum)) + remote_obj = self.get_remote(remote) + return str(remote_obj.odb.hash_to_path_info(checksum)) diff --git a/dvc/dependency/repo.py b/dvc/dependency/repo.py index fd5f6f708..df65a4dbd 100644 --- a/dvc/dependency/repo.py +++ b/dvc/dependency/repo.py @@ -1,3 +1,5 @@ +from typing import NamedTuple, Optional + from voluptuous import Required from dvc.path_info import PathInfo @@ -5,6 +7,11 @@ from dvc.path_info import PathInfo from .base import Dependency +class RepoPair(NamedTuple): + url: str + rev: Optional[str] = None + + class RepoDependency(Dependency): PARAM_REPO = "repo" PARAM_URL = "url" @@ -31,10 +38,10 @@ class RepoDependency(Dependency): return False @property - def repo_pair(self): + def repo_pair(self) -> RepoPair: d = self.def_repo rev = d.get(self.PARAM_REV_LOCK) or d.get(self.PARAM_REV) - return d[self.PARAM_URL], rev + return RepoPair(d[self.PARAM_URL], rev) def __str__(self): return "{} ({})".format(self.def_path, self.def_repo[self.PARAM_URL]) diff --git a/dvc/hash_info.py b/dvc/hash_info.py index 6fba0c21f..72fdbe6f0 100644 --- a/dvc/hash_info.py +++ b/dvc/hash_info.py @@ -21,6 +21,9 @@ class HashInfo: def __str__(self): return f"{self.name}: {self.value}" + def __hash__(self): + return hash((self.name, self.value)) + @classmethod def from_dict(cls, d): _d = d.copy() if d else {} diff --git a/dvc/objects/db/__init__.py b/dvc/objects/db/__init__.py index 507184a8e..8d4620514 100644 --- a/dvc/objects/db/__init__.py +++ b/dvc/objects/db/__init__.py @@ -1,5 +1,3 @@ -from collections import defaultdict - from dvc.scheme import Schemes @@ -80,99 +78,3 @@ class ODBManager: def by_scheme(self): self._init_odb(self.CLOUD_SCHEMES) yield from self._odb.items() - - -class NamedCacheItem: - def __init__(self): - self.names = set() - self.children = defaultdict(NamedCacheItem) - - def __eq__(self, other): - return self.names == other.names and self.children == other.children - - def child_keys(self): - for key, child in self.children.items(): - yield key - yield from child.child_keys() - - def child_names(self): - for key, child in self.children.items(): - yield key, child.names - yield from child.child_names() - - def add(self, checksum, item): - self.children[checksum].update(item) - - def update(self, item, suffix=""): - if suffix: - self.names.update(n + suffix for n in item.names) - else: - self.names.update(item.names) - for checksum, child_item in item.children.items(): - self.children[checksum].update(child_item) - - -class NamedCache: - # pylint: disable=protected-access - def __init__(self): - self._items = defaultdict(lambda: defaultdict(NamedCacheItem)) - self.external = defaultdict(set) - - @classmethod - def make(cls, scheme, checksum, name): - cache = cls() - cache.add(scheme, checksum, name) - return cache - - def __getitem__(self, key): - return self._items[key] - - def add(self, scheme, checksum, name): - """Add a mapped name for the specified checksum.""" - self._items[scheme][checksum].names.add(name) - - def add_child_cache(self, checksum, cache, suffix=""): - """Add/update child cache for the specified checksum.""" - for scheme, src in cache._items.items(): - dst = self._items[scheme][checksum].children - for child_checksum, item in src.items(): - dst[child_checksum].update(item, suffix=suffix) - - for repo_pair, files in cache.external.items(): - self.external[repo_pair].update(files) - - def add_external(self, url, rev, path): - self.external[url, rev].add(path) - - def update(self, cache, suffix=""): - for scheme, src in cache._items.items(): - dst = self._items[scheme] - for checksum, item in src.items(): - dst[checksum].update(item, suffix=suffix) - - for repo_pair, files in cache.external.items(): - self.external[repo_pair].update(files) - - def scheme_keys(self, scheme): - """Iterate over a flat list of all keys for the specified scheme, - including children. - """ - for key, item in self._items[scheme].items(): - yield key - yield from item.child_keys() - - def scheme_names(self, scheme): - """Iterate over a flat list of checksum, names items for the specified - scheme, including children. - """ - for key, item in self._items[scheme].items(): - yield key, item.names - yield from item.child_names() - - def dir_keys(self, scheme): - return ( - key for key, item in self._items[scheme].items() if item.children - ) - - def child_keys(self, scheme, checksum): - return self._items[scheme][checksum].child_keys() diff --git a/dvc/objects/db/base.py b/dvc/objects/db/base.py index edb3bb970..08db92cff 100644 --- a/dvc/objects/db/base.py +++ b/dvc/objects/db/base.py @@ -7,6 +7,7 @@ from typing import Optional from dvc.objects.errors import ObjectFormatError from dvc.objects.file import HashFile +from dvc.objects.tree import Tree from dvc.progress import Tqdm logger = logging.getLogger(__name__) @@ -317,6 +318,14 @@ class ObjectDB: pass def gc(self, used, jobs=None): + used_hashes = set() + for obj in used: + used_hashes.add(obj.hash_info.value) + if isinstance(obj, Tree): + used_hashes.update( + entry_obj.hash_info.value for _, entry_obj in obj + ) + removed = False # hashes must be sorted to ensure we always remove .dir files first for hash_ in sorted( @@ -324,7 +333,7 @@ class ObjectDB: key=self.fs.is_dir_hash, reverse=True, ): - if hash_ in used: + if hash_ in used_hashes: continue path_info = self.hash_to_path_info(hash_) if self.fs.is_dir_hash(hash_): diff --git a/dvc/objects/file.py b/dvc/objects/file.py index a9545d290..f39ceed8c 100644 --- a/dvc/objects/file.py +++ b/dvc/objects/file.py @@ -1,17 +1,30 @@ import errno import logging import os +from typing import TYPE_CHECKING, Optional from .errors import ObjectFormatError +if TYPE_CHECKING: + from dvc.fs.base import BaseFileSystem + from dvc.hash_info import HashInfo + from dvc.types import DvcPath + logger = logging.getLogger(__name__) class HashFile: - def __init__(self, path_info, fs, hash_info): + def __init__( + self, + path_info: Optional["DvcPath"], + fs: Optional["BaseFileSystem"], + hash_info: "HashInfo", + name: Optional[str] = None, + ): self.path_info = path_info self.fs = fs self.hash_info = hash_info + self.name = name @property def size(self): @@ -37,6 +50,15 @@ class HashFile: and self.hash_info == other.hash_info ) + def __hash__(self): + return hash( + ( + self.hash_info, + self.path_info, + self.fs.scheme if self.fs else None, + ) + ) + def check(self, odb, check_hash=True): from .stage import get_file_hash diff --git a/dvc/objects/tree.py b/dvc/objects/tree.py index 79cd181da..783a41366 100644 --- a/dvc/objects/tree.py +++ b/dvc/objects/tree.py @@ -1,6 +1,7 @@ import json import logging import posixpath +from typing import TYPE_CHECKING, Optional, Tuple from funcy import cached_property @@ -8,6 +9,9 @@ from .errors import ObjectFormatError from .file import HashFile from .stage import get_file_hash +if TYPE_CHECKING: + from .db.base import ObjectDB + logger = logging.getLogger(__name__) @@ -115,24 +119,45 @@ class Tree(HashFile): tree = cls.from_list(raw) tree.path_info = obj.path_info tree.fs = obj.fs + for _, entry_obj in tree: + entry_obj.fs = obj.fs tree.hash_info = hash_info return tree - def filter(self, odb, prefix): + def filter( + self, odb: "ObjectDB", prefix: Tuple[str], copy: bool = False + ) -> Optional[HashFile]: + """Return filter object(s) for this tree. + + If copy is True, returned object will be a Tree containing + filtered entries, but with hash_info copied from the original tree. + + If copy is False, returned object will be a raw HashFile or Tree with + newly computed hash_info for the filtered object. + """ obj = self._dict.get(prefix) if obj: + if copy: + tree = Tree(self.path_info, self.fs, self.hash_info) + tree.add(prefix, obj) + return tree return obj - depth = len(prefix) - tree = Tree(None, None, None) + if copy: + tree = Tree(self.path_info, self.fs, self.hash_info) + depth = 0 + else: + tree = Tree(None, None, None) + depth = len(prefix) try: for key, obj in self.trie.items(prefix): tree.add(key[depth:], obj) except KeyError: return None - tree.digest() - odb.add(tree.path_info, tree.fs, tree.hash_info) + if not copy: + tree.digest() + odb.add(tree.path_info, tree.fs, tree.hash_info) return tree diff --git a/dvc/output.py b/dvc/output.py index 0c21bd071..4f1fbc40b 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 +from typing import TYPE_CHECKING, Dict, Set, Type from urllib.parse import urlparse from funcy import collecting, project @@ -24,14 +24,19 @@ from .fs.local import LocalFileSystem from .fs.s3 import S3FileSystem from .hash_info import HashInfo from .istextfile import istextfile +from .objects import Tree from .objects import save as osave -from .objects.db import NamedCache from .objects.errors import ObjectFormatError from .objects.stage import stage as ostage from .scheme import Schemes from .utils import relpath from .utils.fs import path_isin +if TYPE_CHECKING: + from dvc.dependency.repo import RepoPair + + from .objects.file import HashFile + logger = logging.getLogger(__name__) @@ -647,7 +652,7 @@ class Output: def download(self, to, jobs=None): self.fs.download(self.path_info, to.path_info, jobs=jobs) - def get_obj(self, filter_info=None): + def get_obj(self, filter_info=None, **kwargs): if self.obj: obj = self.obj elif self.hash_info: @@ -660,7 +665,7 @@ class Output: if filter_info and filter_info != self.path_info: prefix = filter_info.relative_to(self.path_info).parts - obj = obj.filter(self.odb, prefix) + obj = obj.filter(self.odb, prefix, **kwargs) return obj @@ -786,18 +791,14 @@ class Output: self.odb.unprotect(self.path_info) def get_dir_cache(self, **kwargs): - if not self.is_dir_checksum: raise DvcException("cannot get dir cache for file checksum") + obj = self.odb.get(self.hash_info) try: - objects.check(self.odb, self.odb.get(self.hash_info)) + objects.check(self.odb, obj) except (FileNotFoundError, ObjectFormatError): - self.repo.cloud.pull( - NamedCache.make("local", self.hash_info.value, str(self)), - show_checksums=False, - **kwargs, - ) + self.repo.cloud.pull([obj], show_checksums=False, **kwargs) try: self.obj = objects.load(self.odb, self.hash_info) @@ -808,27 +809,8 @@ class Output: def collect_used_dir_cache( self, remote=None, force=False, jobs=None, filter_info=None - ): - """Get a list of `info`s related to the given directory. - - - Pull the directory entry from the remote cache if it was changed. - - Example: - - Given the following commands: - - $ echo "foo" > directory/foo - $ echo "bar" > directory/bar - $ dvc add directory - - It will return a NamedCache like: - - nc = NamedCache() - nc.add(self.scheme, 'c157a79031e1', 'directory/foo') - nc.add(self.scheme, 'd3b07384d113', 'directory/bar') - """ - - cache = NamedCache() + ) -> Set["HashFile"]: + """Fetch dir cache and return used objects for this out.""" try: self.get_dir_cache(jobs=jobs, remote=remote) @@ -848,38 +830,17 @@ class Output: "unable to fully collect used cache" " without cache for directory '{}'".format(self) ) - return cache + return set() - path = str(self.path_info) - filter_path = str(filter_info) if filter_info else None - for entry_key, entry_obj in self.obj: - entry_path = os.path.join(path, *entry_key) - if ( - not filter_path - or entry_path == filter_path - or entry_path.startswith(filter_path + os.sep) - ): - cache.add(self.scheme, entry_obj.hash_info.value, entry_path) + obj = self.get_obj(filter_info=filter_info, copy=True) + self._set_obj_names(obj) + return {obj} - return cache - - def get_used_cache(self, **kwargs): - """Get a dumpd of the given `out`, with an entry including the branch. - - The `used_cache` of an output is no more than its `info`. - - In case that the given output is a directory, it will also - include the `info` of its files. - """ + def get_used_objs(self, **kwargs) -> Set["HashFile"]: + """Return filtered set of used objects for this out.""" if not self.use_cache: - return NamedCache() - - if self.stage.is_repo_import: - cache = NamedCache() - (dep,) = self.stage.deps - cache.external[dep.repo_pair].add(dep.def_path) - return cache + return set() if not self.hash_info: msg = ( @@ -898,18 +859,30 @@ class Output: ) ) logger.warning(msg) - return NamedCache() + return set() - ret = NamedCache.make(self.scheme, self.hash_info.value, str(self)) + if self.is_dir_checksum: + return self.collect_used_dir_cache(**kwargs) - if not self.is_dir_checksum: - return ret + obj = self.get_obj(filter_info=kwargs.get("filter_info")) + if not obj: + obj = self.odb.get(self.hash_info) + self._set_obj_names(obj) - ret.add_child_cache( - self.hash_info.value, self.collect_used_dir_cache(**kwargs) - ) + return {obj} - return ret + def _set_obj_names(self, obj): + obj.name = str(self) + if isinstance(obj, Tree): + for key, entry_obj in obj: + entry_obj.name = os.path.join(str(self), *key) + + def get_used_external(self, **kwargs) -> Dict["RepoPair", str]: + if not self.use_cache or not self.stage.is_repo_import: + return {} + + (dep,) = self.stage.deps + return {dep.repo_pair: dep.def_path} 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 e878196c3..ba5e39ca6 100644 --- a/dvc/remote/base.py +++ b/dvc/remote/base.py @@ -8,6 +8,7 @@ from functools import partial, wraps from dvc.exceptions import DownloadError, UploadError from dvc.hash_info import HashInfo +from dvc.objects.tree import Tree from ..progress import Tqdm from .index import RemoteIndex @@ -98,7 +99,7 @@ class Remote: def status( self, cache, - named_cache, + objs, jobs=None, show_checksums=False, download=False, @@ -107,7 +108,7 @@ class Remote: # Return flattened dict containing all status info dir_status, file_status, _ = self._status( cache, - named_cache, + objs, jobs=jobs, show_checksums=show_checksums, download=download, @@ -126,10 +127,10 @@ class Remote: return indexed_hashes + self.odb.hashes_exist(list(hashes), **kwargs) - def _status( + def _status( # pylint: disable=unused-argument self, cache, - named_cache, + objs, jobs=None, show_checksums=False, download=False, @@ -142,8 +143,18 @@ class Remote: {dir_hash: set(file_hash, ...)} which can be used to map a .dir file to its file contents. """ + logger.debug(f"Preparing to collect status from {self.odb.path_info}") - md5s = set(named_cache.scheme_keys(cache.fs.scheme)) + md5s = set() + dir_objs = {} + for obj in objs: + if obj.fs.scheme != cache.fs.scheme: + continue + md5s.add(obj.hash_info.value) + if isinstance(obj, Tree): + dir_objs[obj.hash_info.value] = obj + for _, entry_obj in obj: + md5s.add(entry_obj.hash_info.value) logger.debug("Collecting information from local cache...") local_exists = frozenset( @@ -159,11 +170,8 @@ class Remote: else: logger.debug("Collecting information from remote cache...") remote_exists = set() - dir_md5s = set(named_cache.dir_keys(cache.fs.scheme)) - if dir_md5s: - remote_exists.update( - self._indexed_dir_hashes(cache, named_cache, dir_md5s) - ) + if dir_objs: + remote_exists.update(self._indexed_dir_hashes(dir_objs)) md5s.difference_update(remote_exists) if md5s: remote_exists.update( @@ -173,7 +181,7 @@ class Remote: ) return self._make_status( cache, - named_cache, + objs, show_checksums, local_exists, remote_exists, @@ -183,29 +191,36 @@ class Remote: def _make_status( self, cache, - named_cache, + objs, show_checksums, local_exists, remote_exists, log_missing, ): - def make_names(hash_, names): - return {"name": hash_ if show_checksums else " ".join(names)} + def make_names(obj): + if show_checksums or obj.name is None: + name = obj.hash_info.value + else: + name = obj.name + return {"name": name} dir_status = {} file_status = {} dir_contents = {} - for hash_, item in named_cache[cache.fs.scheme].items(): - if item.children: - dir_status[hash_] = make_names(hash_, item.names) + for obj in objs: + if obj.fs.scheme != cache.fs.scheme: + continue + + hash_ = obj.hash_info.value + if isinstance(obj, Tree): + dir_status[hash_] = make_names(obj) dir_contents[hash_] = set() - for child_hash, child in item.children.items(): - file_status[child_hash] = make_names( - child_hash, child.names - ) + for _, child_obj in obj: + child_hash = child_obj.hash_info.value + file_status[child_hash] = make_names(child_obj) dir_contents[hash_].add(child_hash) else: - file_status[hash_] = make_names(hash_, item.names) + file_status[hash_] = make_names(obj) self._fill_statuses(dir_status, local_exists, remote_exists) self._fill_statuses(file_status, local_exists, remote_exists) @@ -215,9 +230,10 @@ class Remote: return dir_status, file_status, dir_contents - def _indexed_dir_hashes(self, cache, named_cache, dir_md5s): + def _indexed_dir_hashes(self, dir_objs): # Validate our index by verifying all indexed .dir hashes # still exist on the remote + dir_md5s = set(dir_objs.keys()) indexed_dirs = set(self.index.dir_hashes()) indexed_dir_exists = set() if indexed_dirs: @@ -239,9 +255,9 @@ class Remote: # If .dir hash exists on the remote, assume directory contents # still exists on the remote for dir_hash in dir_exists: - file_hashes = list( - named_cache.child_keys(cache.fs.scheme, dir_hash) - ) + file_hashes = [ + obj.hash_info.value for _, obj in dir_objs[dir_hash] + ] if dir_hash not in self.index: logger.debug( "Indexing new .dir '{}' with '{}' nested files".format( @@ -291,7 +307,7 @@ class Remote: def _process( self, cache, - named_cache, + objs, jobs=None, show_checksums=False, download=False, @@ -317,7 +333,7 @@ class Remote: dir_status, file_status, dir_contents = self._status( cache, - named_cache, + objs, jobs=jobs, show_checksums=show_checksums, download=download, @@ -457,17 +473,20 @@ class Remote: self.index.update([dir_hash], file_hashes) @index_locked - def push(self, cache, named_cache, jobs=None, show_checksums=False): + def push(self, cache, objs, jobs=None, show_checksums=False): ret = self._process( cache, - named_cache, + objs, jobs=jobs, show_checksums=show_checksums, download=False, ) if self.fs.scheme == "local": - for checksum in named_cache.scheme_keys("local"): + for obj in objs: + if obj.fs.scheme != "local": + continue + checksum = obj.hash_info.value cache_file = self.odb.hash_to_path_info(checksum) if self.fs.exists(cache_file): hash_info = HashInfo(self.fs.PARAM_CHECKSUM, checksum) @@ -477,17 +496,19 @@ class Remote: return ret @index_locked - def pull(self, cache, named_cache, jobs=None, show_checksums=False): + def pull(self, cache, objs, jobs=None, show_checksums=False): ret = self._process( cache, - named_cache, + objs, jobs=jobs, show_checksums=show_checksums, download=True, ) if not self.odb.verify: - for checksum in named_cache.scheme_keys("local"): + + def save_state(obj): + checksum = obj.hash_info.value cache_file = cache.hash_to_path_info(checksum) if cache.fs.exists(cache_file): # We can safely save here, as existing corrupted files @@ -498,6 +519,14 @@ class Remote: cache.state.save(cache_file, cache.fs, hash_info) cache.protect(cache_file) + for obj in objs: + if obj.fs.scheme != "local": + continue + save_state(obj) + if isinstance(obj, Tree): + for _, entry_obj in obj: + save_state(entry_obj) + return ret @staticmethod diff --git a/dvc/repo/__init__.py b/dvc/repo/__init__.py index b844cf6f1..e44337861 100644 --- a/dvc/repo/__init__.py +++ b/dvc/repo/__init__.py @@ -1,5 +1,6 @@ import logging import os +from collections import defaultdict from contextlib import contextmanager from functools import wraps from typing import TYPE_CHECKING, Callable, Optional @@ -370,9 +371,17 @@ class Repo: to items containing the output's `dumpd` names and the output's children (if the given output is a directory). """ - from dvc.objects.db import NamedCache + used_objs = set() + used_external = defaultdict(set) - cache = NamedCache() + def _add_suffix(objs, suffix): + from dvc.objects.tree import Tree + + for obj in objs: + obj.name += suffix + if isinstance(obj, Tree): + for _, entry_obj in obj: + entry_obj.name += suffix for branch in self.brancher( revs=revs, @@ -390,23 +399,28 @@ class Repo: for target in targets ) - suffix = f"({branch})" if branch else "" for stage, filter_info in pairs: - used_cache = stage.get_used_cache( + objs, external = stage.get_used_cache( remote=remote, force=force, jobs=jobs, filter_info=filter_info, ) - cache.update(used_cache, suffix=suffix) + if branch: + _add_suffix(objs, f" ({branch})") + used_objs.update(objs) + for repo_pair, paths in external.items(): + used_external[repo_pair].update(paths) if used_run_cache: - used_cache = self.stage_cache.get_used_cache( + objs, external = self.stage_cache.get_used_cache( used_run_cache, remote=remote, force=force, jobs=jobs ) - cache.update(used_cache) + used_objs.update(objs) + for repo_pair, paths in external.items(): + used_external[repo_pair].update(paths) - return cache + return used_objs, used_external @cached_property def outs_trie(self): diff --git a/dvc/repo/fetch.py b/dvc/repo/fetch.py index b4891499c..5c1c9e277 100644 --- a/dvc/repo/fetch.py +++ b/dvc/repo/fetch.py @@ -4,6 +4,7 @@ import os from dvc.config import NoRemoteError from dvc.exceptions import DownloadError, NoOutputOrStageError +from ..scheme import Schemes from . import locked logger = logging.getLogger(__name__) @@ -40,7 +41,7 @@ def fetch( if isinstance(targets, str): targets = [targets] - used = self.used_cache( + used_objs, used_external = self.used_cache( targets, all_branches=all_branches, all_tags=all_tags, @@ -60,16 +61,21 @@ def fetch( if run_cache: self.stage_cache.pull(remote) downloaded += self.cloud.pull( - used, jobs, remote=remote, show_checksums=show_checksums + used_objs, + jobs, + remote=remote, + show_checksums=show_checksums, ) except NoRemoteError: - if not used.external and used["local"]: + 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_url, repo_rev), files in used.external.items(): - d, f = _fetch_external(self, repo_url, repo_rev, files, jobs) + for repo_pair, files in used_external.items(): + d, f = _fetch_external(self, repo_pair.url, repo_pair.rev, files, jobs) downloaded += d failed += f @@ -102,13 +108,13 @@ def _fetch_external(self, repo_url, repo_rev, files, jobs): for path in files: path_info = root / path try: - used = repo.used_cache( + used_objs, _ = repo.used_cache( [os.fspath(path_info)], force=True, jobs=jobs, recursive=True, ) - cb(repo.cloud.pull(used, jobs)) + cb(repo.cloud.pull(used_objs, jobs)) except (NoOutputOrStageError, NoRemoteError): pass obj = stage(odb, path_info, repo.repo_fs, "md5", jobs=jobs) diff --git a/dvc/repo/gc.py b/dvc/repo/gc.py index 896638fa7..7cce7e8cc 100644 --- a/dvc/repo/gc.py +++ b/dvc/repo/gc.py @@ -1,11 +1,14 @@ import logging +from typing import TYPE_CHECKING, Set from dvc.exceptions import InvalidArgumentError -from dvc.objects.db import NamedCache from ..scheme import Schemes from . import locked +if TYPE_CHECKING: + from dvc.objects.file import HashFile + logger = logging.getLogger(__name__) @@ -52,30 +55,32 @@ def gc( repos = [] all_repos = [Repo(path) for path in repos] + used_objs: Set["HashFile"] = set() with ExitStack() as stack: for repo in all_repos: stack.enter_context(repo.lock) - used = NamedCache() for repo in all_repos + [self]: - used.update( - repo.used_cache( - all_branches=all_branches, - with_deps=with_deps, - all_tags=all_tags, - all_commits=all_commits, - all_experiments=all_experiments, - remote=remote, - force=force, - jobs=jobs, - ) + objs, _ = repo.used_cache( + all_branches=all_branches, + with_deps=with_deps, + all_tags=all_tags, + all_commits=all_commits, + all_experiments=all_experiments, + remote=remote, + force=force, + jobs=jobs, ) + used_objs.update(objs) for scheme, odb in self.odb.by_scheme(): if not odb: continue - removed = odb.gc(set(used.scheme_keys(scheme)), jobs=jobs) + removed = odb.gc( + {obj for obj in used_objs if obj.fs.scheme == scheme}, + jobs=jobs, + ) if not removed: logger.info(f"No unused '{scheme}' cache to remove.") @@ -83,6 +88,9 @@ def gc( return remote = self.cloud.get_remote(remote, "gc -c") - removed = remote.gc(set(used.scheme_keys(Schemes.LOCAL)), jobs=jobs) + removed = remote.gc( + {obj for obj in used_objs if obj.fs.scheme == Schemes.LOCAL}, + jobs=jobs, + ) if not removed: logger.info("No unused cache to remove from remote.") diff --git a/dvc/repo/push.py b/dvc/repo/push.py index 23bb5244d..532eede43 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 = self.used_cache( + used_objs, _ = self.used_cache( expanded_targets, all_branches=all_branches, all_tags=all_tags, @@ -38,4 +38,6 @@ def push( revs=revs, ) - return len(used_run_cache) + self.cloud.push(used, jobs, remote=remote) + return len(used_run_cache) + self.cloud.push( + used_objs, jobs, remote=remote + ) diff --git a/dvc/repo/status.py b/dvc/repo/status.py index fb4f86b01..60f6dd111 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 = self.used_cache( + used_objs, _ = self.used_cache( targets, all_branches=all_branches, all_tags=all_tags, @@ -92,7 +92,7 @@ def _cloud_status( ret = {} status_info = self.cloud.status( - used, jobs, remote=remote, log_missing=False + used_objs, jobs, remote=remote, log_missing=False ) for info in status_info.values(): name = info["name"] diff --git a/dvc/stage/__init__.py b/dvc/stage/__init__.py index 176a56a7b..0ba452793 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 +from typing import TYPE_CHECKING, Dict, Optional, Set, Tuple from funcy import cached_property, project @@ -35,7 +35,9 @@ from .utils import ( ) if TYPE_CHECKING: + from dvc.dependency.repo import RepoPair from dvc.dvcfile import DVCFile + from dvc.objects.file import HashFile logger = logging.getLogger(__name__) # Disallow all punctuation characters except hyphen and underscore @@ -622,14 +624,20 @@ class Stage(params.StageParams): for out in self.filter_outs(filter_info) ) - def get_used_cache(self, *args, **kwargs): - from dvc.objects.db import NamedCache - - cache = NamedCache() + def get_used_cache( + 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) for out in self.filter_outs(kwargs.get("filter_info")): - cache.update(out.get_used_cache(*args, **kwargs)) + 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 cache + return used_objs, used_external @staticmethod def _check_can_merge(stage, ancestor_out=None): diff --git a/dvc/stage/cache.py b/dvc/stage/cache.py index 7b9df6f77..c06e71e68 100644 --- a/dvc/stage/cache.py +++ b/dvc/stage/cache.py @@ -1,6 +1,7 @@ import logging import os import tempfile +from collections import defaultdict from contextlib import contextmanager from funcy import cached_property, first @@ -196,7 +197,8 @@ class StageCache: cached_stage = self._create_stage(cache, wdir=stage.wdir) if pull: - self.repo.cloud.pull(cached_stage.get_used_cache()) + objs, _ = cached_stage.get_used_cache() + self.repo.cloud.pull(objs) if not cached_stage.outs_cached(): raise RunCacheNotFoundError(stage) @@ -250,14 +252,17 @@ class StageCache: ) def get_used_cache(self, used_run_cache, *args, **kwargs): - from dvc.objects.db import NamedCache - - cache = NamedCache() + """Return used cache for the specified run-cached stages.""" + used_objs = set() + used_external = defaultdict(set) for key, value in used_run_cache: entry = self._load_cache(key, value) if not entry: continue stage = self._create_stage(entry) - cache.update(stage.get_used_cache(*args, **kwargs)) - return cache + 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 diff --git a/tests/func/test_data_cloud.py b/tests/func/test_data_cloud.py index 474cb0705..9fa0c623c 100644 --- a/tests/func/test_data_cloud.py +++ b/tests/func/test_data_cloud.py @@ -8,7 +8,6 @@ from flaky.flaky_decorator import flaky import dvc as dvc_module from dvc.external_repo import clean_repos from dvc.main import main -from dvc.objects.db import NamedCache from dvc.remote.base import ( STATUS_DELETED, STATUS_MISSING, @@ -57,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 - info = out.get_used_cache() + foo_objs = out.get_used_objs() (stage_dir,) = tmp_dir.dvc_gen( { @@ -70,29 +69,39 @@ def test_cloud(tmp_dir, dvc, remote): # pylint:disable=unused-argument ) out_dir = stage_dir.outs[0] cache_dir = out_dir.cache_path - name_dir = str(out_dir) md5_dir = out_dir.hash_info.value - info_dir = NamedCache.make(out_dir.scheme, md5_dir, name_dir) + dir_objs = {out_dir.obj} + entry_md5s = [entry_obj.hash_info.value for _, entry_obj in out_dir.obj] + + def _make_dir_status(status): + expected = {md5_dir: {"name": md5_dir, "status": status}} + expected.update( + { + entry_md5: {"name": entry_md5, "status": status} + for entry_md5 in entry_md5s + } + ) + return expected # Check status - status = dvc.cloud.status(info, show_checksums=True) + status = dvc.cloud.status(foo_objs, show_checksums=True) expected = {md5: {"name": md5, "status": STATUS_NEW}} assert status == expected - status_dir = dvc.cloud.status(info_dir, show_checksums=True) - expected = {md5_dir: {"name": md5_dir, "status": STATUS_NEW}} + status_dir = dvc.cloud.status(dir_objs, show_checksums=True) + expected = _make_dir_status(STATUS_NEW) assert status_dir == expected # Move cache and check status # See issue https://github.com/iterative/dvc/issues/4383 for details backup_dir = dvc.odb.local.cache_dir + ".backup" move(dvc.odb.local.cache_dir, backup_dir) - status = dvc.cloud.status(info, show_checksums=True) + status = dvc.cloud.status(foo_objs, show_checksums=True) expected = {md5: {"name": md5, "status": STATUS_MISSING}} assert status == expected - status_dir = dvc.cloud.status(info_dir, show_checksums=True) - expected = {md5_dir: {"name": md5_dir, "status": STATUS_MISSING}} + status_dir = dvc.cloud.status(dir_objs, show_checksums=True) + expected = _make_dir_status(STATUS_MISSING) assert status_dir == expected # Restore original cache: @@ -100,48 +109,48 @@ def test_cloud(tmp_dir, dvc, remote): # pylint:disable=unused-argument move(backup_dir, dvc.odb.local.cache_dir) # Push and check status - dvc.cloud.push(info) + dvc.cloud.push(foo_objs) assert os.path.exists(cache) assert os.path.isfile(cache) - dvc.cloud.push(info_dir) + dvc.cloud.push(dir_objs) assert os.path.isfile(cache_dir) - status = dvc.cloud.status(info, show_checksums=True) + status = dvc.cloud.status(foo_objs, show_checksums=True) expected = {md5: {"name": md5, "status": STATUS_OK}} assert status == expected - status_dir = dvc.cloud.status(info_dir, show_checksums=True) - expected = {md5_dir: {"name": md5_dir, "status": STATUS_OK}} + status_dir = dvc.cloud.status(dir_objs, show_checksums=True) + expected = _make_dir_status(STATUS_OK) assert status_dir == expected # Remove and check status remove(dvc.odb.local.cache_dir) - status = dvc.cloud.status(info, show_checksums=True) + status = dvc.cloud.status(foo_objs, show_checksums=True) expected = {md5: {"name": md5, "status": STATUS_DELETED}} assert status == expected - status_dir = dvc.cloud.status(info_dir, show_checksums=True) - expected = {md5_dir: {"name": md5_dir, "status": STATUS_DELETED}} + status_dir = dvc.cloud.status(dir_objs, show_checksums=True) + expected = _make_dir_status(STATUS_DELETED) assert status_dir == expected # Pull and check status - dvc.cloud.pull(info) + dvc.cloud.pull(foo_objs) assert os.path.exists(cache) assert os.path.isfile(cache) with open(cache) as fd: assert fd.read() == "foo" - dvc.cloud.pull(info_dir) + dvc.cloud.pull(dir_objs) assert os.path.isfile(cache_dir) - status = dvc.cloud.status(info, show_checksums=True) + status = dvc.cloud.status(foo_objs, show_checksums=True) expected = {md5: {"name": md5, "status": STATUS_OK}} assert status == expected - status_dir = dvc.cloud.status(info_dir, show_checksums=True) - expected = {md5_dir: {"name": md5_dir, "status": STATUS_OK}} + status_dir = dvc.cloud.status(dir_objs, show_checksums=True) + expected = _make_dir_status(STATUS_OK) assert status_dir == expected @@ -277,7 +286,10 @@ def test_missing_cache(tmp_dir, dvc, local_remote, caplog): assert bar in caplog.text caplog.clear() - assert dvc.status(cloud=True) == {"bar": "missing", "foo": "missing"} + assert dvc.status(cloud=True) == { + "bar": "missing", + "foo": "missing", + } assert header not in caplog.text assert foo not in caplog.text assert bar not in caplog.text diff --git a/tests/func/test_remote.py b/tests/func/test_remote.py index e7cd7eb4f..14f3f2b32 100644 --- a/tests/func/test_remote.py +++ b/tests/func/test_remote.py @@ -204,6 +204,7 @@ def test_partial_push_n_pull(tmp_dir, dvc, tmp_path_factory, local_remote): dvc.push() remove(dvc.odb.local.cache_dir) + baz.collect_used_dir_cache() with patch.object(LocalFileSystem, "_download", side_effect=Exception): with pytest.raises(DownloadError) as download_error_info: dvc.pull() @@ -430,17 +431,22 @@ def test_push_incomplete_dir(tmp_dir, dvc, mocker, local_remote): remote = dvc.cloud.get_remote("upstream") odb = dvc.odb.local - dir_hash = stage.outs[0].hash_info.value - used = stage.get_used_cache(remote=remote) + out = stage.outs[0] + file_objs = [entry_obj for _, entry_obj in out.obj] # remove one of the cache files for directory - file_hashes = list(used.child_keys(odb.fs.scheme, dir_hash)) - remove(odb.hash_to_path_info(file_hashes[0])) + remove(odb.hash_to_path_info(file_objs[0].hash_info.value)) dvc.push() - assert not remote.fs.exists(remote.odb.hash_to_path_info(dir_hash)) - assert not remote.fs.exists(remote.odb.hash_to_path_info(file_hashes[0])) - assert remote.fs.exists(remote.odb.hash_to_path_info(file_hashes[1])) + assert not remote.fs.exists( + remote.odb.hash_to_path_info(out.hash_info.value) + ) + assert not remote.fs.exists( + remote.odb.hash_to_path_info(file_objs[0].hash_info.value) + ) + assert remote.fs.exists( + remote.odb.hash_to_path_info(file_objs[1].hash_info.value) + ) def test_upload_exists(tmp_dir, dvc, local_remote): diff --git a/tests/unit/objects/db/test_local.py b/tests/unit/objects/db/test_local.py index 6448b60cc..dc8f89e3e 100644 --- a/tests/unit/objects/db/test_local.py +++ b/tests/unit/objects/db/test_local.py @@ -4,8 +4,9 @@ import os import pytest from dvc.fs.local import LocalFileSystem -from dvc.objects.db import NamedCache +from dvc.hash_info import HashInfo from dvc.objects.db.local import LocalObjectDB +from dvc.objects.file import HashFile from dvc.path_info import PathInfo from dvc.remote.index import RemoteIndexNoop @@ -17,11 +18,16 @@ def test_status_download_optimization(mocker, dvc): """ odb = LocalObjectDB(LocalFileSystem(), PathInfo(".")) - infos = NamedCache() - infos.add("local", "acbd18db4cc2f85cedef654fccc4a4d8", "foo") - infos.add("local", "37b51d194a7513e45b56f6524f2d51f2", "bar") + objs = { + HashFile( + None, odb.fs, HashInfo("md5", "acbd18db4cc2f85cedef654fccc4a4d8") + ), + HashFile( + None, odb.fs, HashInfo("md5", "37b51d194a7513e45b56f6524f2d51f2") + ), + } - local_exists = list(infos["local"]) + local_exists = [obj.hash_info.value for obj in objs] mocker.patch.object(odb, "hashes_exist", return_value=local_exists) other_remote = mocker.Mock() @@ -29,7 +35,7 @@ def test_status_download_optimization(mocker, dvc): other_remote.hashes_exist.return_value = [] other_remote.index = RemoteIndexNoop() - other_remote.status(odb, infos, download=True) + other_remote.status(odb, objs, download=True) assert other_remote.hashes_exist.call_count == 0 diff --git a/tests/unit/output/test_output.py b/tests/unit/output/test_output.py index 667453912..197ac57ce 100644 --- a/tests/unit/output/test_output.py +++ b/tests/unit/output/test_output.py @@ -6,7 +6,6 @@ from funcy import first from voluptuous import MultipleInvalid, Schema from dvc.ignore import _no_match -from dvc.objects.db import NamedCache from dvc.output import CHECKSUM_SCHEMA, Output from dvc.path_info import PathInfo from dvc.stage import Stage @@ -100,5 +99,6 @@ def test_get_used_cache(exists, expected_message, mocker, caplog): ).return_value = exists with caplog.at_level(logging.WARNING, logger="dvc"): - assert isinstance(output.get_used_cache(), NamedCache) + assert set() == output.get_used_objs() + assert {} == output.get_used_external() assert first(caplog.messages) == expected_message diff --git a/tests/unit/repo/test_repo.py b/tests/unit/repo/test_repo.py index 4e797bc5f..8a7401c43 100644 --- a/tests/unit/repo/test_repo.py +++ b/tests/unit/repo/test_repo.py @@ -4,6 +4,7 @@ import shutil import pytest from dvc.exceptions import OutputDuplicationError +from dvc.hash_info import HashInfo from dvc.repo import NotDvcRepoError, Repo, locked @@ -44,26 +45,25 @@ def test_find_outs_by_path_does_graph_checks(tmp_dir, dvc): [os.path.join("dir", "subdir", "file"), os.path.join("dir", "subdir")], ) def test_used_cache(tmp_dir, dvc, path): - from dvc.objects.db import NamedCache + from dvc.objects.tree import Tree tmp_dir.dvc_gen({"dir": {"subdir": {"file": "file"}, "other": "other"}}) - expected = NamedCache.make( - "local", "70922d6bf66eb073053a82f77d58c536.dir", "dir" - ) - expected.add_child_cache( - "70922d6bf66eb073053a82f77d58c536.dir", - NamedCache.make( - "local", - "8c7dd922ad47494fc02c388e12c00eac", - os.path.join("dir", "subdir", "file"), - ), - ) - used_cache = dvc.used_cache([path]) - assert ( - used_cache._items == expected._items - and used_cache.external == expected.external - ) + expected_objs = { + 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) + + assert used == expected_objs + assert external == expected_external def test_locked(mocker):