From 71bae7e4affbc5f37aceb9412dd60822201a942d 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: Sat, 10 Jul 2021 00:48:28 +0900 Subject: [PATCH] objects: use separate staging ODB for staging trees (#6195) * fs: use fsspecwrapper in memfs * odb: add memfs staging ODB in ODBManager * odb: add read_only attribute for odb instances * objects: handle memfs paths when staging trees * odb: remove git ODB * output: use staging ODB * RepoDependency: use staging odb for imports * dvcfs: handle case where dir cache is unavailable * fix circular import check * odb: use temp local ODB for staging * objects: use staging ODB in stage() and tree load * objects: add typing for add/check/load/stage * output/dep: use new staging * add unit tests for local odb staging * localfs: account for string paths in move() * include local in odb.by_scheme * tests: count dir cache like any other object in test_pull_git_imports * odb: clear staging on gc() * objects: add tree.get/filter/insert * output: use obj.filter/get for granular dir commit * squash filter logic output * dependency: don't filter imports at the dep level * don't init staging ODB for external cache * use staging ODB when computing tree hash for external outs * remove unneeded force in test granular commit * use unique-per-manager memfs URLs for staging * odb: make staging per-ODBManager * objects: move state dependent staging into ODBManager * objects: support checking multiple odbs - staging dependent checks can be done via ODBManager.check * checkout: do state lookup before explicit stage() in _changed() * output: use odbmanager based staging * RepoDependency: use odbmanager based staging * update obj staging unit tests * checkout: remove check(), handle individual file exceptions on link() * diff: use odbmanager based staging * objects: add ObjectPermissionError for read/write perm errors * odb: move staging out of odbmanager * objects: handle staging ODB in objects.stage() * objects.tree: don't stage anything in tree methods * odb: clear staged objects on gc() * remote: don't filter objs by scheme, objs are already grouped by ODB * diff: update objects.stage usage * output/dep: update stage() usage * update ODB tests * update func tests * remove memfs from FS_MAP * revert unneeded objects.check() changes * remove unused tree.insert() * revert unused odbmanager changes * add separate odb error classes * stage external outs directly in ODB * check for obj existence in main ODB before staging * use .dvc/tmp/staging instead of per-odb staging * gc: only gc loal staging once --- dvc/checkout.py | 67 +++++---- dvc/dependency/repo.py | 98 ++++++------- dvc/fs/dvc.py | 2 + dvc/fs/local.py | 4 +- dvc/fs/memory.py | 55 +++---- dvc/objects/__init__.py | 20 ++- dvc/objects/db/__init__.py | 1 + dvc/objects/db/base.py | 48 +++++-- dvc/objects/db/git.py | 24 ---- dvc/objects/db/ssh.py | 2 +- dvc/objects/errors.py | 8 ++ dvc/objects/stage.py | 214 +++++++++++++++++++--------- dvc/objects/tree.py | 53 ++++--- dvc/output.py | 49 +++++-- dvc/remote/base.py | 10 -- dvc/repo/diff.py | 2 +- dvc/repo/fetch.py | 4 +- dvc/repo/gc.py | 2 + dvc/scheme.py | 1 + tests/func/test_commit.py | 22 ++- tests/func/test_data_cloud.py | 2 +- tests/func/test_external_repo.py | 2 +- tests/func/test_remote.py | 2 + tests/unit/objects/db/test_local.py | 66 +++++++++ 24 files changed, 481 insertions(+), 277 deletions(-) delete mode 100644 dvc/objects/db/git.py diff --git a/dvc/checkout.py b/dvc/checkout.py index 89d7c422a..ae29a5db9 100644 --- a/dvc/checkout.py +++ b/dvc/checkout.py @@ -21,7 +21,7 @@ from dvc.types import Optional logger = logging.getLogger(__name__) -def _changed(path_info, fs, obj, cache): +def _changed(path_info, fs, obj, cache, state=None): logger.trace("checking if '%s'('%s') has changed.", path_info, obj) try: @@ -32,11 +32,13 @@ def _changed(path_info, fs, obj, cache): ) return True - try: - actual = stage(cache, path_info, fs, obj.hash_info.name).hash_info - except FileNotFoundError: - logger.debug("'%s' doesn't exist.", path_info) - return True + actual = state.get(path_info, fs) if state else None + if actual is None: + try: + actual = stage(cache, path_info, fs, obj.hash_info.name).hash_info + except FileNotFoundError: + logger.debug("'%s' doesn't exist.", path_info) + return True if obj.hash_info != actual: logger.debug( @@ -120,9 +122,11 @@ def _try_links(cache, from_info, to_info, link_types): def _link(cache, from_info, to_info): - assert cache.fs.isfile(from_info) cache.makedirs(to_info.parent) - _try_links(cache, from_info, to_info, cache.cache_types) + try: + _try_links(cache, from_info, to_info, cache.cache_types) + except FileNotFoundError as exc: + raise CheckoutError([to_info]) from exc def _cache_is_copy(cache, path_info): @@ -163,7 +167,7 @@ def _checkout_file( modified = False cache_info = cache.hash_to_path_info(obj.hash_info.value) if fs.exists(path_info): - if not relink and _changed(path_info, fs, obj, cache): + if not relink and _changed(path_info, fs, obj, cache, state=state): modified = True _remove(path_info, fs, cache, force=force) _link(cache, cache_info, path_info) @@ -222,19 +226,23 @@ def _checkout_dir( logger.debug("Linking directory '%s'.", path_info) + failed = [] for entry_key, entry_obj in obj: - entry_modified = _checkout_file( - path_info.joinpath(*entry_key), - fs, - entry_obj, - cache, - force, - progress_callback, - relink, - state=None, - ) - if entry_modified: - modified = True + try: + entry_modified = _checkout_file( + path_info.joinpath(*entry_key), + fs, + entry_obj, + cache, + force, + progress_callback, + relink, + state=None, + ) + if entry_modified: + modified = True + except CheckoutError as exc: + failed.extend(exc.target_infos) modified = ( _remove_redundant_files( @@ -243,6 +251,9 @@ def _checkout_dir( or modified ) + if failed: + raise CheckoutError(failed) + if state: state.save(path_info, fs, obj.hash_info) @@ -310,21 +321,9 @@ def checkout( _remove(path_info, fs, cache, force=force) failed = path_info - elif not relink and not _changed(path_info, fs, obj, cache): + elif not relink and not _changed(path_info, fs, obj, cache, state=state): logger.trace("Data '%s' didn't change.", path_info) # type: ignore skip = True - else: - try: - check(cache, obj) - except (FileNotFoundError, ObjectFormatError): - if not quiet: - logger.warning( - "Cache '%s' not found. File '%s' won't be created.", - obj.hash_info, - path_info, - ) - _remove(path_info, fs, cache, force=force) - failed = path_info if failed or skip: if progress_callback and obj: diff --git a/dvc/dependency/repo.py b/dvc/dependency/repo.py index a9bbdcbf4..32f51b890 100644 --- a/dvc/dependency/repo.py +++ b/dvc/dependency/repo.py @@ -1,7 +1,8 @@ import os from collections import defaultdict -from typing import TYPE_CHECKING, Dict, Optional, Set +from typing import TYPE_CHECKING, Dict, Optional, Set, Tuple +from funcy import first from voluptuous import Required from dvc.path_info import PathInfo @@ -62,16 +63,18 @@ class RepoDependency(Dependency): def download(self, to, jobs=None): from dvc.checkout import checkout + from dvc.fs.memory import MemoryFileSystem from dvc.objects import save - from dvc.objects.db.git import GitObjectDB from dvc.repo.fetch import fetch_from_odb for odb, objs in self.get_used_objs().items(): - if not isinstance(odb, GitObjectDB): + if isinstance(odb.fs, MemoryFileSystem): + for obj in objs: + save(self.repo.odb.local, obj, jobs=jobs) + else: 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, @@ -96,10 +99,15 @@ class RepoDependency(Dependency): def get_used_objs( self, **kwargs ) -> Dict[Optional["ObjectDB"], Set["HashFile"]]: + used, _ = self._get_used_and_obj(**kwargs) + return used + + def _get_used_and_obj( + self, obj_only=False, **kwargs + ) -> Tuple[Dict[Optional["ObjectDB"], Set["HashFile"]], "HashFile"]: from dvc.config import NoRemoteError from dvc.exceptions import NoOutputOrStageError, PathMissingError - from dvc.objects.db.git import GitObjectDB - from dvc.objects.stage import stage + from dvc.objects.stage import get_staging, stage local_odb = self.repo.odb.local locked = kwargs.pop("locked", True) @@ -112,23 +120,25 @@ class RepoDependency(Dependency): 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 + if not obj_only: + 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 + odb.read_only = True + self._check_circular_import(objs) + used_objs[odb].update(objs) + except (NoRemoteError, NoOutputOrStageError): + pass try: staged_obj = stage( - local_odb, + None, path_info, repo.repo_fs, local_odb.fs.PARAM_CHECKSUM, @@ -137,45 +147,37 @@ class RepoDependency(Dependency): raise PathMissingError( self.def_path, self.def_repo[self.PARAM_URL] ) from exc + staging = get_staging() + staging.read_only = True 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 + used_objs[staging].add(staged_obj) + return used_objs, staged_obj - def _check_circular_import(self, odb): + def _check_circular_import(self, objs): from dvc.exceptions import CircularImportError from dvc.fs.repo import RepoFileSystem + from dvc.objects.tree import Tree - if not odb or not isinstance(odb.fs, RepoFileSystem): + obj = first(objs) + if isinstance(obj, Tree): + _, obj = first(obj) + if not isinstance(obj.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) + if obj.fs.repo_url is not None and obj.fs.repo_url == self_url: + raise CircularImportError(self, obj.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 + locked = kwargs.get("locked", True) + rev = self._get_rev(locked=locked) + if rev in self._staged_objs: + return self._staged_objs[rev] + _, obj = self._get_used_and_obj( + obj_only=True, filter_info=filter_info, **kwargs + ) + return obj def _make_repo(self, locked=True, **kwargs): from dvc.external_repo import external_repo diff --git a/dvc/fs/dvc.py b/dvc/fs/dvc.py index 5755ab7ca..10e27ec5e 100644 --- a/dvc/fs/dvc.py +++ b/dvc/fs/dvc.py @@ -51,6 +51,8 @@ class DvcFileSystem(BaseFileSystem): # pylint:disable=abstract-method # NOTE: use string paths here for performance reasons key = tuple(relpath(path_info, out.path_info).split(os.sep)) out.get_dir_cache(remote=remote) + if out.obj is None: + raise FileNotFoundError obj = out.obj.trie.get(key) if obj: return obj.hash_info diff --git a/dvc/fs/local.py b/dvc/fs/local.py index 996c53ff9..147cb06fa 100644 --- a/dvc/fs/local.py +++ b/dvc/fs/local.py @@ -90,7 +90,9 @@ class LocalFileSystem(BaseFileSystem): return os.stat(path) def move(self, from_info, to_info): - if from_info.scheme != "local" or to_info.scheme != "local": + if ( + isinstance(from_info, PathInfo) and from_info.scheme != "local" + ) or (isinstance(to_info, PathInfo) and to_info.scheme != "local"): raise NotImplementedError self.makedirs(to_info.parent) diff --git a/dvc/fs/memory.py b/dvc/fs/memory.py index 1175e9a0c..2c646da99 100644 --- a/dvc/fs/memory.py +++ b/dvc/fs/memory.py @@ -1,34 +1,35 @@ -from .base import BaseFileSystem +import threading + +from funcy import cached_property, wrap_prop + +from dvc.path_info import CloudURLInfo +from dvc.scheme import Schemes + +from .fsspec_wrapper import FSSpecWrapper -class MemoryFileSystem(BaseFileSystem): - scheme = "local" +class MemoryFileSystem(FSSpecWrapper): # pylint:disable=abstract-method + scheme = Schemes.MEMORY PARAM_CHECKSUM = "md5" + PATH_CLS = CloudURLInfo + TRAVERSE_PREFIX_LEN = 2 + DEFAULT_BLOCKSIZE = 4096 - def __init__(self, **kwargs): + def __eq__(self, other): + # NOTE: all fsspec MemoryFileSystem instances are equivalent and use a + # single global store + return isinstance(other, type(self)) + + __hash__ = FSSpecWrapper.__hash__ + + @wrap_prop(threading.Lock()) + @cached_property + def fs(self): from fsspec.implementations.memory import MemoryFileSystem as MemFS - super().__init__(**kwargs) + return MemFS(**self.fs_args) - self.fs = MemFS() - - def exists(self, path_info) -> bool: - return self.fs.exists(path_info.fspath) - - def open(self, path_info, mode="r", encoding=None, **kwargs): - return self.fs.open( - path_info.fspath, mode=mode, encoding=encoding, **kwargs - ) - - def info(self, path_info): - return self.fs.info(path_info.fspath) - - def stat(self, path_info): - import os - - info = self.fs.info(path_info.fspath) - - return os.stat_result((0, 0, 0, 0, 0, 0, info["size"], 0, 0, 0)) - - def walk_files(self, path_info, **kwargs): - raise NotImplementedError + def open(self, *args, **kwargs): + with super().open(*args, **kwargs) as fobj: + fobj.blocksize = self.DEFAULT_BLOCKSIZE + return fobj diff --git a/dvc/objects/__init__.py b/dvc/objects/__init__.py index f464ccb4b..bfe1321d4 100644 --- a/dvc/objects/__init__.py +++ b/dvc/objects/__init__.py @@ -7,6 +7,8 @@ from dvc.progress import Tqdm from .tree import Tree if TYPE_CHECKING: + from dvc.hash_info import HashInfo + from .db.base import ObjectDB from .file import HashFile @@ -20,6 +22,8 @@ def save( **kwargs, ): if isinstance(obj, Tree): + from .stage import get_staging + with ThreadPoolExecutor(max_workers=jobs) as executor: for future in Tqdm( as_completed( @@ -38,18 +42,26 @@ def save( ): future.result() + # if dir cache for this tree has already been staged, move the + # staged object when possible + staging = get_staging(odb) + if staging.exists(obj.hash_info): + obj = staging.get(obj.hash_info) + kwargs["move"] = True + + assert obj.path_info and obj.fs and obj.hash_info odb.add(obj.path_info, obj.fs, obj.hash_info, **kwargs) -def check(odb, obj): - odb.check(obj.hash_info) - +def check(odb: "ObjectDB", obj: "HashFile"): if isinstance(obj, Tree): for _, entry in obj: odb.check(entry.hash_info) + odb.check(obj.hash_info) -def load(odb, hash_info): + +def load(odb: "ObjectDB", hash_info: "HashInfo") -> "HashFile": if hash_info.isdir: return Tree.load(odb, hash_info) return odb.get(hash_info) diff --git a/dvc/objects/db/__init__.py b/dvc/objects/db/__init__.py index 8d4620514..38fb441be 100644 --- a/dvc/objects/db/__init__.py +++ b/dvc/objects/db/__init__.py @@ -26,6 +26,7 @@ def _get_odb(repo, settings): return None cls, config, path_info = get_cloud_fs(repo, **settings) + config["tmp_dir"] = repo.tmp_dir return get_odb(cls(**config), path_info, state=repo.state, **config) diff --git a/dvc/objects/db/base.py b/dvc/objects/db/base.py index c0419bc54..7c2b92657 100644 --- a/dvc/objects/db/base.py +++ b/dvc/objects/db/base.py @@ -3,13 +3,17 @@ import logging import os from concurrent.futures import ThreadPoolExecutor from copy import copy -from typing import Optional +from typing import TYPE_CHECKING, Optional -from dvc.objects.errors import ObjectFormatError +from dvc.objects.errors import ObjectDBPermissionError, ObjectFormatError from dvc.objects.file import HashFile -from dvc.objects.tree import Tree from dvc.progress import Tqdm +if TYPE_CHECKING: + from dvc.fs.base import BaseFileSystem + from dvc.hash_info import HashInfo + from dvc.types import AnyPath + logger = logging.getLogger(__name__) @@ -19,7 +23,7 @@ class ObjectDB: DEFAULT_CACHE_TYPES = ["copy"] CACHE_MODE: Optional[int] = None - def __init__(self, fs, path_info, **config): + def __init__(self, fs: "BaseFileSystem", path_info: "AnyPath", **config): from dvc.state import StateNoop self.fs = fs @@ -30,6 +34,7 @@ class ObjectDB: self.cache_type_confirmed = False self.slow_link_warning = config.get("slow_link_warning", True) self.tmp_dir = config.get("tmp_dir") + self.read_only = config.get("read_only", False) @property def config(self): @@ -39,21 +44,29 @@ class ObjectDB: "type": self.cache_types, "slow_link_warning": self.slow_link_warning, "tmp_dir": self.tmp_dir, + "read_only": self.read_only, } def __eq__(self, other): - return self.fs == other.fs and self.path_info == other.path_info + return ( + self.fs == other.fs + and self.path_info == other.path_info + and self.read_only == other.read_only + ) def __hash__(self): return hash((self.fs.scheme, self.path_info)) + def exists(self, hash_info: "HashInfo"): + return self.fs.exists(self.hash_to_path_info(hash_info.value)) + def move(self, from_info, to_info): self.fs.move(from_info, to_info) def makedirs(self, path_info): self.fs.makedirs(path_info) - def get(self, hash_info): + def get(self, hash_info: "HashInfo"): """get raw object""" return HashFile( # Prefer string path over PathInfo when possible due to performance @@ -62,7 +75,16 @@ class ObjectDB: hash_info, ) - def add(self, path_info, fs, hash_info, move=True, **kwargs): + def add( + self, + path_info: "AnyPath", + fs: "BaseFileSystem", + hash_info: "HashInfo", + move: bool = True, + **kwargs, + ): + if self.read_only: + raise ObjectDBPermissionError("Cannot add to read-only ODB") try: self.check(hash_info, check_hash=self.verify) return @@ -125,7 +147,11 @@ class ObjectDB: def set_exec(self, path_info): # pylint: disable=unused-argument pass - def check(self, hash_info, check_hash=True): + def check( + self, + hash_info: "HashInfo", + check_hash: bool = True, + ): """Compare the given hash with the (corresponding) actual one if check_hash is specified, or just verify the existence of the cache files on the filesystem. @@ -142,7 +168,7 @@ class ObjectDB: obj = self.get(hash_info) if self.is_protected(obj.path_info): - logger.trace( + logger.trace( # type: ignore[attr-defined] "Assuming '%s' is unchanged since it is read-only", obj.path_info, ) @@ -335,6 +361,10 @@ class ObjectDB: pass def gc(self, used, jobs=None): + from ..tree import Tree + + if self.read_only: + raise ObjectDBPermissionError("Cannot gc read-only ODB") used_hashes = set() for obj in used: used_hashes.add(obj.hash_info.value) diff --git a/dvc/objects/db/git.py b/dvc/objects/db/git.py deleted file mode 100644 index f5869d5dd..000000000 --- a/dvc/objects/db/git.py +++ /dev/null @@ -1,24 +0,0 @@ -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/objects/db/ssh.py b/dvc/objects/db/ssh.py index 1f981f100..1461fef06 100644 --- a/dvc/objects/db/ssh.py +++ b/dvc/objects/db/ssh.py @@ -79,7 +79,7 @@ class SSHObjectDB(ObjectDB): else: root = self.path_info.path with self.fs.ssh() as ssh: - if prefix and not ssh.exists(root): + if not ssh.exists(root): return # If we simply return an iterator then with above closes instantly if progress_callback: diff --git a/dvc/objects/errors.py b/dvc/objects/errors.py index c7616b65e..af87174ef 100644 --- a/dvc/objects/errors.py +++ b/dvc/objects/errors.py @@ -7,3 +7,11 @@ class ObjectError(DvcException): class ObjectFormatError(ObjectError): pass + + +class ObjectDBError(DvcException): + pass + + +class ObjectDBPermissionError(ObjectDBError): + pass diff --git a/dvc/objects/stage.py b/dvc/objects/stage.py index 63bcafa05..fdade1a12 100644 --- a/dvc/objects/stage.py +++ b/dvc/objects/stage.py @@ -1,13 +1,25 @@ from concurrent.futures import ThreadPoolExecutor from functools import partial +from typing import TYPE_CHECKING, Optional from dvc.exceptions import DvcIgnoreInCollectedDirError from dvc.hash_info import HashInfo from dvc.ignore import DvcIgnore -from dvc.objects.file import HashFile from dvc.progress import Tqdm from dvc.utils import file_md5 +from .file import HashFile + +if TYPE_CHECKING: + from dvc.fs.base import BaseFileSystem + from dvc.types import DvcPath + + from .db.base import ObjectDB + + +_STAGING_DIR = "staging" +_STAGING_MEMFS_PATH = f"dvc-{_STAGING_DIR}" + def _upload_file(path_info, fs, odb): from dvc.utils import tmp_fname @@ -59,11 +71,12 @@ def get_file_hash(path_info, fs, name, state=None): return hash_info -def _get_file_obj(path_info, fs, name, odb=None, state=None, upload=False): +def _get_file_obj(path_info, fs, name, odb=None, upload=False): if upload: assert odb and name == "md5" return _upload_file(path_info, fs, odb) + state = odb.state if odb else None obj = HashFile( path_info, fs, get_file_hash(path_info, fs, name, state=state) ) @@ -71,7 +84,13 @@ def _get_file_obj(path_info, fs, name, odb=None, state=None, upload=False): def _build_objects( - path_info, fs, name, odb, state, upload, dvcignore=None, **kwargs + path_info, + fs, + name, + odb=None, + upload=False, + dvcignore=None, + **kwargs, ): if dvcignore: walk_iterator = dvcignore.walk_files(fs, path_info) @@ -88,7 +107,6 @@ def _build_objects( fs=fs, name=name, odb=odb, - state=state, upload=upload, ) ) @@ -98,7 +116,7 @@ def _build_objects( yield from executor.map(worker, walk_iterator) -def _iter_objects(path_info, fs, name, odb, state, upload, **kwargs): +def _iter_objects(path_info, fs, name, upload=False, **kwargs): if not upload and name in fs.DETAIL_FIELDS: for details in fs.find(path_info, detail=True): file_info = path_info.replace(path=details["name"]) @@ -107,18 +125,14 @@ def _iter_objects(path_info, fs, name, odb, state, upload, **kwargs): return None - yield from _build_objects( - path_info, fs, name, odb, state, upload, **kwargs - ) + yield from _build_objects(path_info, fs, name, upload=upload, **kwargs) -def _build_tree(path_info, fs, name, odb, state, upload, **kwargs): +def _build_tree(path_info, fs, name, **kwargs): from .tree import Tree tree = Tree(None, None, None) - for file_info, obj in _iter_objects( - path_info, fs, name, odb, state, upload, **kwargs - ): + for file_info, obj in _iter_objects(path_info, fs, name, **kwargs): if DvcIgnore.DVCIGNORE_FILE == file_info.name: raise DvcIgnoreInCollectedDirError(file_info.parent) @@ -131,85 +145,143 @@ def _build_tree(path_info, fs, name, odb, state, upload, **kwargs): # Yes, this is a BUG, as long as we permit "/" in # filenames on Windows and "\" on Unix tree.add(file_info.relative_to(path_info).parts, obj) - tree.digest() - return tree -def _get_tree_obj(path_info, fs, name, odb, state, upload, **kwargs): +def _get_tree_obj(path_info, fs, name, odb=None, **kwargs): from .tree import Tree value = fs.info(path_info).get(name) - if value: + if odb and value: hash_info = HashInfo(name, value) try: - return Tree.load(odb, hash_info) + tree = Tree.load(odb, hash_info) + # NOTE: loaded entries are naive objects with hash_infos but no + # path_info. For staging trees, obj.path_info should be relative + # to the staging src `path_info` and src fs + for key, entry in tree: + entry.fs = fs + entry.path_info = path_info.joinpath(*key) + return tree except FileNotFoundError: pass - tree = _build_tree(path_info, fs, name, odb, state, upload, **kwargs) - - odb.add(tree.path_info, tree.fs, tree.hash_info) - if name != "md5": - # NOTE: used only for external outputs. Initial reasoning was to be - # able to validate .dir files right in the workspace (e.g. check s3 - # etag), but could be dropped for manual validation with regular md5, - # that would be universal for all clouds. - raw = odb.get(tree.hash_info) - hash_info = get_file_hash(raw.path_info, raw.fs, name, state) - tree.hash_info.name = hash_info.name - tree.hash_info.value = hash_info.value - if not tree.hash_info.value.endswith(".dir"): - tree.hash_info.value += ".dir" - odb.add(tree.path_info, tree.fs, tree.hash_info) - + tree = _build_tree(path_info, fs, name, odb=odb, **kwargs) return tree -def stage(odb, path_info, fs, name, upload=False, **kwargs): - assert path_info and ( - isinstance(path_info, str) or path_info.scheme == fs.scheme - ) +def get_staging(odb: Optional["ObjectDB"] = None) -> "ObjectDB": + """Return an ODB that can be used for staging objects. + + If odb.fs is local, .dvc/tmp/staging will be returned. Otherwise + the the global (temporary) memfs ODB will be returned. + """ + + from dvc.fs.memory import MemoryFileSystem + from dvc.path_info import CloudURLInfo, PathInfo + from dvc.scheme import Schemes + + from .db import get_odb + + if odb and odb.fs.scheme == Schemes.LOCAL and odb.tmp_dir: + fs = odb.fs + path_info: "DvcPath" = PathInfo(odb.tmp_dir) / _STAGING_DIR + config = odb.config + else: + fs = MemoryFileSystem() + path_info = CloudURLInfo(f"{Schemes.MEMORY}://{_STAGING_MEMFS_PATH}") + config = {} + return get_odb(fs, path_info, **config) + + +def _load_from_state(odb, staging, path_info, fs, name): + from . import load + from .errors import ObjectFormatError + from .tree import Tree + + state = odb.state + hash_info = state.get(path_info, fs) + if hash_info: + for odb_ in (odb, staging): + if odb_.exists(hash_info): + try: + obj = load(odb_, hash_info) + if isinstance(obj, Tree): + obj.hash_info.nfiles = len(obj) + for key, entry in obj: + entry.fs = fs + entry.path_info = path_info.joinpath(*key) + else: + obj.fs = fs + obj.path_info = path_info + assert obj.hash_info.name == name + obj.hash_info.size = hash_info.size + return obj + except ObjectFormatError: + pass + raise FileNotFoundError + + +def _stage_external_tree_info(odb, tree, name): + # NOTE: used only for external outputs. Initial reasoning was to be + # able to validate .dir files right in the workspace (e.g. check s3 + # etag), but could be dropped for manual validation with regular md5, + # that would be universal for all clouds. + assert odb and name != "md5" + + odb.add(tree.path_info, tree.fs, tree.hash_info) + raw = odb.get(tree.hash_info) + hash_info = get_file_hash(raw.path_info, raw.fs, name, state=odb.state) + tree.path_info = raw.path_info + tree.fs = raw.fs + tree.hash_info.name = hash_info.name + tree.hash_info.value = hash_info.value + if not tree.hash_info.value.endswith(".dir"): + tree.hash_info.value += ".dir" + return tree + + +def stage( + odb: Optional["ObjectDB"], + path_info: "DvcPath", + fs: "BaseFileSystem", + name: str, + upload: bool = False, + **kwargs, +) -> "HashFile": + assert path_info and path_info.scheme == fs.scheme details = fs.info(path_info) - state = odb.state - # pylint: disable=assignment-from-none - hash_info = state.get(path_info, fs) - - # If we have dir hash in state db, but dir cache file is lost, - # then we need to recollect the dir via .get_dir_hash() call below, - # see https://github.com/iterative/dvc/issues/2219 for context - if ( - hash_info - and hash_info.isdir - and not odb.fs.exists(odb.hash_to_path_info(hash_info.value)) - ): - hash_info = None - - if hash_info: - from . import load - from .tree import Tree - - obj = load(odb, hash_info) - if isinstance(obj, Tree): - obj.hash_info.nfiles = len(obj) - for key, entry in obj: - entry.fs = fs - entry.path_info = path_info.joinpath(*key) - else: - obj.fs = fs - obj.path_info = path_info - assert obj.hash_info.name == name - obj.hash_info.size = hash_info.size - return obj + staging = get_staging(odb) + if odb: + try: + return _load_from_state(odb, staging, path_info, fs, name) + except FileNotFoundError: + pass if details["type"] == "directory": - obj = _get_tree_obj(path_info, fs, name, odb, state, upload, **kwargs) + obj = _get_tree_obj( + path_info, fs, name, odb=odb, upload=upload, **kwargs + ) + if name == "md5": + if odb and odb.exists(obj.hash_info): + raw = odb.get(obj.hash_info) + else: + staging.add(obj.path_info, obj.fs, obj.hash_info) + raw = staging.get(obj.hash_info) + # cleanup unneeded memfs tmpfile and return obj based on staging + # ODB fs/path + if obj.fs != raw.fs: + obj.fs.remove(obj.path_info) + obj.fs = raw.fs + obj.path_info = raw.path_info + else: + obj = _stage_external_tree_info(odb, obj, name) else: - _, obj = _get_file_obj(path_info, fs, name, odb, state, upload) + _, obj = _get_file_obj(path_info, fs, name, odb=odb, upload=upload) - if obj.hash_info: - state.save(path_info, fs, obj.hash_info) + if odb and odb.state and obj.hash_info: + odb.state.save(path_info, fs, obj.hash_info) return obj diff --git a/dvc/objects/tree.py b/dvc/objects/tree.py index 0debbfc09..1e3d2e840 100644 --- a/dvc/objects/tree.py +++ b/dvc/objects/tree.py @@ -1,7 +1,7 @@ import json import logging import posixpath -from typing import TYPE_CHECKING, Optional, Tuple +from typing import Optional, Tuple from funcy import cached_property @@ -9,9 +9,6 @@ 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__) @@ -34,11 +31,11 @@ class Tree(HashFile): def digest(self): from dvc.fs.memory import MemoryFileSystem - from dvc.path_info import PathInfo + from dvc.path_info import CloudURLInfo from dvc.utils import tmp_fname memfs = MemoryFileSystem() - path_info = PathInfo(tmp_fname("")) + path_info = CloudURLInfo("memory://{}".format(tmp_fname(""))) with memfs.open(path_info, "wb") as fobj: fobj.write(self.as_bytes()) self.fs = memfs @@ -96,7 +93,6 @@ class Tree(HashFile): @classmethod def load(cls, odb, hash_info): - obj = odb.get(hash_info) try: @@ -121,39 +117,40 @@ class Tree(HashFile): return tree - def filter( - self, odb: "ObjectDB", prefix: Tuple[str], copy: bool = False - ) -> Optional[HashFile]: - """Return filter object(s) for this tree. + def filter(self, prefix: Tuple[str]) -> Optional["Tree"]: + """Return a filtered copy of this tree that only contains entries + inside prefix. - If copy is True, returned object will be a Tree containing - filtered entries, but with hash_info copied from the original tree. + The returned tree will contain the original tree's hash_info and + path_info. - If copy is False, returned object will be a raw HashFile or Tree with - newly computed hash_info for the filtered object. + Returns an empty tree if no object exists at the specified prefix. + """ + tree = Tree(self.path_info, self.fs, self.hash_info) + try: + for key, obj in self.trie.items(prefix): + tree.add(key, obj) + except KeyError: + pass + return tree + + def get(self, prefix: Tuple[str]) -> Optional[HashFile]: + """Return object at the specified prefix in this tree. + + Returns None if no object exists at the specified prefix. """ 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 - if copy: - tree = Tree(self.path_info, self.fs, self.hash_info) - depth = 0 - else: - tree = Tree(None, None, None) - depth = len(prefix) + 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 - if not copy: - tree.digest() - odb.add(tree.path_info, tree.fs, tree.hash_info) + tree.digest() return tree diff --git a/dvc/output.py b/dvc/output.py index 5b812faa5..f194a4940 100644 --- a/dvc/output.py +++ b/dvc/output.py @@ -566,14 +566,22 @@ class Output: assert self.hash_info if self.use_cache: - obj = ostage( - self.odb, - filter_info or self.path_info, - self.fs, - self.odb.fs.PARAM_CHECKSUM, - dvcignore=self.dvcignore, + granular = ( + self.is_dir_checksum + and filter_info + and filter_info != self.path_info ) - objects.save(self.odb, obj) + if granular: + obj = self._commit_granular_dir(filter_info) + else: + obj = ostage( + self.odb, + filter_info or self.path_info, + self.fs, + self.odb.fs.PARAM_CHECKSUM, + dvcignore=self.dvcignore, + ) + osave(self.odb, obj) checkout( filter_info or self.path_info, self.fs, @@ -585,6 +593,20 @@ class Output: ) self.set_exec() + def _commit_granular_dir(self, filter_info): + prefix = filter_info.relative_to(self.path_info).parts + save_obj = ostage( + self.odb, + self.path_info, + self.fs, + self.odb.fs.PARAM_CHECKSUM, + dvcignore=self.dvcignore, + ) + save_obj = save_obj.filter(prefix) + checkout_obj = save_obj.get(prefix) + osave(self.odb, save_obj) + return checkout_obj + def dumpd(self): ret = copy(self.hash_info.to_dict()) @@ -670,7 +692,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, **kwargs) + obj = obj.get(prefix) return obj @@ -802,7 +824,7 @@ class Output: obj = self.odb.get(self.hash_info) try: objects.check(self.odb, obj) - except (FileNotFoundError, ObjectFormatError): + except FileNotFoundError: self.repo.cloud.pull([obj], show_checksums=False, **kwargs) try: @@ -824,7 +846,7 @@ class Output: try: objects.check(self.odb, self.odb.get(self.hash_info)) - except (FileNotFoundError, ObjectFormatError): + except FileNotFoundError: msg = ( "Missing cache for directory '{}'. " "Cache for files inside will be lost. " @@ -837,7 +859,10 @@ class Output: ) return {} - obj = self.get_obj(filter_info=filter_info, copy=True) + obj = self.get_obj() + if filter_info and filter_info != self.path_info: + prefix = filter_info.relative_to(self.path_info).parts + obj = obj.filter(prefix) self._set_obj_names(obj) return {None: {obj}} @@ -894,7 +919,7 @@ class Output: return {} (dep,) = self.stage.deps - return dep.get_used_objs() + return dep.get_used_objs(**kwargs) 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 1d211484c..fff64b32e 100644 --- a/dvc/remote/base.py +++ b/dvc/remote/base.py @@ -165,8 +165,6 @@ class Remote: 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 @@ -197,7 +195,6 @@ class Remote: ) ) return self._make_status( - cache, objs, show_checksums, local_exists, @@ -207,7 +204,6 @@ class Remote: def _make_status( self, - cache, objs, show_checksums, local_exists, @@ -225,8 +221,6 @@ class Remote: file_status = {} dir_contents = {} for obj in objs: - if obj.fs.scheme != cache.fs.scheme: - continue hash_ = obj.hash_info.value if isinstance(obj, Tree): @@ -501,8 +495,6 @@ class Remote: if self.fs.scheme == "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): @@ -537,8 +529,6 @@ class Remote: 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: diff --git a/dvc/repo/diff.py b/dvc/repo/diff.py index 1b279e58c..74431c3b0 100644 --- a/dvc/repo/diff.py +++ b/dvc/repo/diff.py @@ -5,7 +5,6 @@ from typing import Dict, List from dvc.exceptions import PathMissingError from dvc.objects.stage import get_file_hash -from dvc.objects.stage import stage as ostage from dvc.repo import locked from dvc.repo.experiments.utils import fix_exp_head @@ -116,6 +115,7 @@ def _paths_checksums(repo, repo_fs, targets): def _output_paths(repo, repo_fs, targets): from dvc.fs.local import LocalFileSystem + from dvc.objects.stage import stage as ostage on_working_fs = isinstance(repo.fs, LocalFileSystem) diff --git a/dvc/repo/fetch.py b/dvc/repo/fetch.py index 688a95345..a0c3ee8f2 100644 --- a/dvc/repo/fetch.py +++ b/dvc/repo/fetch.py @@ -36,7 +36,7 @@ def fetch( config.NoRemoteError: thrown when downloading only local files and no remote is configured """ - from dvc.objects.db.git import GitObjectDB + from dvc.fs.memory import MemoryFileSystem if isinstance(targets, str): targets = [targets] @@ -76,7 +76,7 @@ def fetch( ) downloaded += d failed += f - elif isinstance(odb, GitObjectDB): + elif isinstance(odb.fs, MemoryFileSystem): # objs contains staged import objects which should be saved # last (after all other objects have been pulled) external.update(objs) diff --git a/dvc/repo/gc.py b/dvc/repo/gc.py index 46eb9e3fa..f98a37448 100644 --- a/dvc/repo/gc.py +++ b/dvc/repo/gc.py @@ -49,6 +49,7 @@ def gc( from contextlib import ExitStack + from dvc.objects.stage import get_staging from dvc.repo import Repo if not repos: @@ -83,6 +84,7 @@ def gc( ) if not removed: logger.info(f"No unused '{scheme}' cache to remove.") + get_staging(self.odb.local).gc(set(), jobs=jobs) if not cloud: return diff --git a/dvc/scheme.py b/dvc/scheme.py index 3c177dcd3..10425daa8 100644 --- a/dvc/scheme.py +++ b/dvc/scheme.py @@ -12,3 +12,4 @@ class Schemes: OSS = "oss" WEBDAV = "webdav" WEBDAVS = "webdavs" + MEMORY = "memory" diff --git a/tests/func/test_commit.py b/tests/func/test_commit.py index 958c1983c..21c231053 100644 --- a/tests/func/test_commit.py +++ b/tests/func/test_commit.py @@ -123,6 +123,8 @@ def test_commit_granular_output_dir(tmp_dir, dvc): def test_commit_granular_dir(tmp_dir, dvc): + from dvc.objects.stage import _STAGING_DIR + tmp_dir.gen( { "data": { @@ -135,8 +137,11 @@ def test_commit_granular_dir(tmp_dir, dvc): dvc.add("data", no_commit=True) cache = tmp_dir / ".dvc" / "cache" - assert set(cache.glob("*/*")) == { - cache / "1a" / "ca2c799df82929bbdd976557975546.dir" + staging = tmp_dir / ".dvc" / "tmp" / _STAGING_DIR + + assert set(cache.glob("*/*")) == set() + assert set(staging.glob("*/*")) == { + staging / "1a" / "ca2c799df82929bbdd976557975546.dir" } dvc.commit(os.path.join("data", "foo")) @@ -144,15 +149,26 @@ def test_commit_granular_dir(tmp_dir, dvc): cache / "1a" / "ca2c799df82929bbdd976557975546.dir", cache / "ac" / "bd18db4cc2f85cedef654fccc4a4d8", } + assert set(staging.glob("*/*")) == set() dvc.commit(os.path.join("data", "subdir")) assert set(cache.glob("*/*")) == { cache / "1a" / "ca2c799df82929bbdd976557975546.dir", - cache / "26" / "d6b64d96a660707412f523e8184b5f.dir", cache / "ac" / "bd18db4cc2f85cedef654fccc4a4d8", cache / "4c" / "e8d2a2cf314a52fa7f315ca37ca445", cache / "68" / "dde2c3c4e7953c2290f176bbdc9a54", } + assert set(staging.glob("*/*")) == set() + + dvc.commit(os.path.join("data")) + assert set(cache.glob("*/*")) == { + cache / "1a" / "ca2c799df82929bbdd976557975546.dir", + cache / "ac" / "bd18db4cc2f85cedef654fccc4a4d8", + cache / "4c" / "e8d2a2cf314a52fa7f315ca37ca445", + cache / "68" / "dde2c3c4e7953c2290f176bbdc9a54", + cache / "37" / "b51d194a7513e45b56f6524f2d51f2", + } + assert set(staging.glob("*/*")) == set() def test_commit_no_exec_missing_dep(tmp_dir, dvc): diff --git a/tests/func/test_data_cloud.py b/tests/func/test_data_cloud.py index fa055302e..c1c97632f 100644 --- a/tests/func/test_data_cloud.py +++ b/tests/func/test_data_cloud.py @@ -340,7 +340,7 @@ def test_pull_git_imports(tmp_dir, dvc, scm, erepo): os.makedirs(dvc.odb.local.cache_dir, exist_ok=True) clean_repos() - assert dvc.pull(force=True)["fetched"] == 2 + assert dvc.pull(force=True)["fetched"] == 3 assert (tmp_dir / "foo").exists() assert (tmp_dir / "foo").read_text() == "foo" diff --git a/tests/func/test_external_repo.py b/tests/func/test_external_repo.py index bd41e5205..ab7189c09 100644 --- a/tests/func/test_external_repo.py +++ b/tests/func/test_external_repo.py @@ -209,7 +209,7 @@ def test_subrepos_are_ignored(tmp_dir, erepo_dir): dvcignore=repo.dvcignore, ) save(repo.odb.local, obj) - assert set(cache_dir.glob("*/*")) == { + assert set(cache_dir.glob("??/*")) == { cache_dir / "e1" / "d9e8eae5374860ae025ec84cfd85c7.dir", cache_dir / "37" / "b51d194a7513e45b56f6524f2d51f2", cache_dir / "94" / "7d2b84e5aa88170e80dff467a5bfb6", diff --git a/tests/func/test_remote.py b/tests/func/test_remote.py index 14f3f2b32..2e9fc083e 100644 --- a/tests/func/test_remote.py +++ b/tests/func/test_remote.py @@ -155,6 +155,7 @@ def test_dir_hash_should_be_key_order_agnostic(tmp_dir, dvc): tree = Tree.from_list( [{"relpath": "1", "md5": "1"}, {"relpath": "2", "md5": "2"}] ) + tree.digest() with patch("dvc.objects.stage._get_tree_obj", return_value=tree): hash1 = stage( dvc.odb.local, path_info, dvc.odb.local.fs, "md5" @@ -163,6 +164,7 @@ def test_dir_hash_should_be_key_order_agnostic(tmp_dir, dvc): tree = Tree.from_list( [{"md5": "1", "relpath": "1"}, {"md5": "2", "relpath": "2"}] ) + tree.digest() with patch("dvc.objects.stage._get_tree_obj", return_value=tree): hash2 = stage( dvc.odb.local, path_info, dvc.odb.local.fs, "md5" diff --git a/tests/unit/objects/db/test_local.py b/tests/unit/objects/db/test_local.py index dc8f89e3e..14eb8330f 100644 --- a/tests/unit/objects/db/test_local.py +++ b/tests/unit/objects/db/test_local.py @@ -92,3 +92,69 @@ def test_set_exec_ignore_errors(tmp_dir, dvc, mocker, err): ) dvc.odb.local.set_exec(PathInfo("foo")) assert mock_chmod.called + + +def test_staging_file(tmp_dir, dvc): + from dvc.objects import check, save + from dvc.objects.stage import get_staging, stage + + tmp_dir.gen("foo", "foo") + fs = LocalFileSystem() + + local_odb = dvc.odb.local + staging_odb = get_staging(local_odb) + obj = stage(local_odb, tmp_dir / "foo", fs, "md5") + + for odb in (local_odb, staging_odb): + path_info = odb.hash_to_path_info(obj.hash_info.value) + assert not odb.fs.exists(path_info) + + # check for file after staging should fail since files are not added on + # stage() + with pytest.raises(FileNotFoundError): + check(local_odb, obj) + with pytest.raises(FileNotFoundError): + check(staging_odb, obj) + + save(local_odb, obj) + check(local_odb, obj) + with pytest.raises(FileNotFoundError): + check(staging_odb, obj) + + path_info = local_odb.hash_to_path_info(obj.hash_info.value) + assert fs.exists(path_info) + + +def test_staging_dir(tmp_dir, dvc): + from dvc.objects import check, save + from dvc.objects.stage import get_staging, stage + + tmp_dir.gen({"dir": {"foo": "foo", "bar": "bar"}}) + fs = LocalFileSystem() + local_odb = dvc.odb.local + staging_odb = get_staging(local_odb) + + obj = stage(local_odb, tmp_dir / "dir", fs, "md5") + + path_info = local_odb.hash_to_path_info(obj.hash_info.value) + assert not local_odb.fs.exists(path_info) + path_info = staging_odb.hash_to_path_info(obj.hash_info.value) + assert staging_odb.fs.exists(path_info) + + # check for raw object after staging should pass only when using the + # staging odb + raw = HashFile(obj.path_info, obj.fs, obj.hash_info) + with pytest.raises(FileNotFoundError): + check(local_odb, raw) + check(staging_odb, raw) + + # checking the entire tree should fail since individual file entries are + # not added on stage() + with pytest.raises(FileNotFoundError): + check(staging_odb, obj) + + save(local_odb, obj) + check(local_odb, obj) + + path_info = local_odb.hash_to_path_info(obj.hash_info.value) + assert fs.exists(path_info)