From fa550f91c30cd8bd238fd4fcb76452c0c74669bb Mon Sep 17 00:00:00 2001 From: Ruslan Kuprieiev Date: Wed, 10 Feb 2021 07:49:27 +0200 Subject: [PATCH] cache: move unrelated helpers out of CloudCache class (#5439) * cache: move transfer to objects Moving it as is for now. * cache: move link() to checkout --- dvc/cache/base.py | 158 +---------------------------------------- dvc/cache/local.py | 11 --- dvc/checkout.py | 94 ++++++++++++++++++++++-- dvc/objects.py | 89 ++++++++++++++++++++++- dvc/remote/base.py | 10 ++- dvc/repo/add.py | 8 ++- dvc/system.py | 11 +-- dvc/tree/local.py | 2 +- dvc/utils/fs.py | 3 - tests/func/test_add.py | 4 +- tests/func/test_s3.py | 4 +- 11 files changed, 203 insertions(+), 191 deletions(-) diff --git a/dvc/cache/base.py b/dvc/cache/base.py index 860d746b7..905313560 100644 --- a/dvc/cache/base.py +++ b/dvc/cache/base.py @@ -1,21 +1,14 @@ import itertools import logging -from concurrent import futures from concurrent.futures import ThreadPoolExecutor from copy import copy from typing import Optional from funcy import decorator -from shortuuid import uuid -from dvc.dir_info import DirInfo -from dvc.exceptions import CacheLinkError, DvcException from dvc.progress import Tqdm -from dvc.remote.slow_link_detection import ( # type: ignore[attr-defined] - slow_link_guard, -) -from ..objects import HashFile, ObjectFormatError, Tree +from ..objects import HashFile, ObjectFormatError logger = logging.getLogger(__name__) @@ -43,60 +36,12 @@ class CloudCache: ) self.cache_type_confirmed = False - def link(self, from_info, to_info): - self._link(from_info, to_info, self.cache_types) - def move(self, from_info, to_info): self.tree.move(from_info, to_info) def makedirs(self, path_info): self.tree.makedirs(path_info) - def _link(self, from_info, to_info, link_types): - assert self.tree.isfile(from_info) - - self.makedirs(to_info.parent) - - self._try_links(from_info, to_info, link_types) - - def _verify_link(self, path_info, link_type): - if self.cache_type_confirmed: - return - - is_link = getattr(self.tree, f"is_{link_type}", None) - if is_link and not is_link(path_info): - self.tree.remove(path_info) - raise DvcException(f"failed to verify {link_type}") - - self.cache_type_confirmed = True - - @slow_link_guard - def _try_links(self, from_info, to_info, link_types): - while link_types: - link_method = getattr(self.tree, link_types[0]) - try: - self._do_link(from_info, to_info, link_method) - self._verify_link(to_info, link_types[0]) - return - - except DvcException as exc: - logger.debug( - "Cache type '%s' is not supported: %s", link_types[0], exc - ) - del link_types[0] - - raise CacheLinkError([to_info]) - - def _do_link(self, from_info, to_info, link_method): - if self.tree.exists(to_info): - raise DvcException(f"Link '{to_info}' already exists!") - - link_method(from_info, to_info) - - logger.debug( - "Created '%s': %s -> %s", self.cache_types[0], from_info, to_info, - ) - def get(self, hash_info): """ get raw object """ return HashFile( @@ -129,107 +74,6 @@ class CloudCache: if callback: callback(1) - def _transfer_file(self, from_tree, from_info): - from dvc.utils import tmp_fname - from dvc.utils.stream import HashedStreamReader - - tmp_info = self.tree.path_info / tmp_fname() - with from_tree.open( - from_info, mode="rb", chunk_size=from_tree.CHUNK_SIZE - ) as stream: - stream_reader = HashedStreamReader(stream) - # Since we don't know the hash beforehand, we'll - # upload it to a temporary location and then move - # it. - self.tree.upload_fobj( - stream_reader, - tmp_info, - total=from_tree.getsize(from_info), - desc=from_info.name, - ) - - hash_info = stream_reader.hash_info - return tmp_info, hash_info - - def _transfer_directory_contents(self, from_tree, from_info, jobs, pbar): - rel_path_infos = {} - from_infos = from_tree.walk_files(from_info) - - def create_tasks(executor, amount): - for entry_info in itertools.islice(from_infos, amount): - pbar.total += 1 - task = executor.submit( - pbar.wrap_fn(self._transfer_file), from_tree, entry_info - ) - rel_path_infos[task] = entry_info.relative_to(from_info) - yield task - - pbar.total = 0 - with ThreadPoolExecutor(max_workers=jobs) as executor: - tasks = set(create_tasks(executor, jobs * 5)) - - while tasks: - done, tasks = futures.wait( - tasks, return_when=futures.FIRST_COMPLETED - ) - tasks.update(create_tasks(executor, len(done))) - for task in done: - yield rel_path_infos.pop(task), task.result() - - def _transfer_directory( - self, from_tree, from_info, jobs, no_progress_bar=False - ): - dir_info = DirInfo() - - with Tqdm(total=1, unit="Files", disable=no_progress_bar) as pbar: - for ( - entry_info, - (entry_tmp_info, entry_hash), - ) in self._transfer_directory_contents( - from_tree, from_info, jobs, pbar - ): - self.add(entry_tmp_info, self.tree, entry_hash) - dir_info.trie[entry_info.parts] = entry_hash - - return Tree.save_dir_info(self, dir_info) - - def transfer(self, from_tree, from_info, jobs=None, no_progress_bar=False): - jobs = jobs or min((from_tree.jobs, self.tree.jobs)) - - if from_tree.isdir(from_info): - return self._transfer_directory( - from_tree, - from_info, - jobs=jobs, - no_progress_bar=no_progress_bar, - ) - tmp_info, hash_info = self._transfer_file(from_tree, from_info) - self.add(tmp_info, self.tree, hash_info) - return hash_info - - def cache_is_copy(self, path_info): - """Checks whether cache uses copies.""" - if self.cache_type_confirmed: - return self.cache_types[0] == "copy" - - if set(self.cache_types) <= {"copy"}: - return True - - workspace_file = path_info.with_name("." + uuid()) - test_cache_file = self.tree.path_info / ".cache_type_test_file" - if not self.tree.exists(test_cache_file): - self.makedirs(test_cache_file.parent) - with self.tree.open(test_cache_file, "wb") as fobj: - fobj.write(bytes(1)) - try: - self.link(test_cache_file, workspace_file) - finally: - self.tree.remove(workspace_file) - self.tree.remove(test_cache_file) - - self.cache_type_confirmed = True - return self.cache_types[0] == "copy" - # Override to return path as a string instead of PathInfo for clouds # which support string paths (see local) def hash_to_path(self, hash_): diff --git a/dvc/cache/local.py b/dvc/cache/local.py index 35c01c97e..733b4f3f6 100644 --- a/dvc/cache/local.py +++ b/dvc/cache/local.py @@ -55,11 +55,6 @@ class LocalCache(CloudCache): makedirs(path_info, exist_ok=True, mode=self._dir_mode) - def link(self, from_info, to_info): - super().link(from_info, to_info) - if self.cache_types[0] not in ("symlink", "hardlink"): - os.chmod(to_info, self._file_mode) - def hash_to_path(self, hash_): # NOTE: `self.cache_path` is already normalized so we can simply use # `os.sep` instead of `os.path.join`. This results in this helper @@ -85,12 +80,6 @@ class LocalCache(CloudCache): return ret - def _verify_link(self, path_info, link_type): - if link_type == "hardlink" and self.tree.getsize(path_info) == 0: - return - - super()._verify_link(path_info, link_type) - def _list_paths(self, prefix=None, progress_callback=None): assert self.tree.path_info is not None if prefix: diff --git a/dvc/checkout.py b/dvc/checkout.py index fb911179d..3727978c6 100644 --- a/dvc/checkout.py +++ b/dvc/checkout.py @@ -1,8 +1,18 @@ import logging +from shortuuid import uuid + import dvc.prompt as prompt -from dvc.exceptions import CheckoutError, ConfirmRemoveError +from dvc.exceptions import ( + CacheLinkError, + CheckoutError, + ConfirmRemoveError, + DvcException, +) from dvc.objects import ObjectFormatError, check, load +from dvc.remote.slow_link_detection import ( # type: ignore[attr-defined] + slow_link_guard, +) logger = logging.getLogger(__name__) @@ -61,6 +71,80 @@ def _remove(path_info, tree, cache, force=False): tree.remove(path_info) +def _verify_link(cache, path_info, link_type): + if link_type == "hardlink" and cache.tree.getsize(path_info) == 0: + return + + if cache.cache_type_confirmed: + return + + is_link = getattr(cache.tree, f"is_{link_type}", None) + if is_link and not is_link(path_info): + cache.tree.remove(path_info) + raise DvcException(f"failed to verify {link_type}") + + cache.cache_type_confirmed = True + + +def _do_link(cache, from_info, to_info, link_method): + if cache.tree.exists(to_info): + raise DvcException(f"Link '{to_info}' already exists!") + + link_method(from_info, to_info) + + logger.debug( + "Created '%s': %s -> %s", cache.cache_types[0], from_info, to_info, + ) + + +@slow_link_guard +def _try_links(cache, from_info, to_info, link_types): + while link_types: + link_method = getattr(cache.tree, link_types[0]) + try: + _do_link(cache, from_info, to_info, link_method) + _verify_link(cache, to_info, link_types[0]) + return + + except DvcException as exc: + logger.debug( + "Cache type '%s' is not supported: %s", link_types[0], exc + ) + del link_types[0] + + raise CacheLinkError([to_info]) + + +def _link(cache, from_info, to_info): + assert cache.tree.isfile(from_info) + cache.makedirs(to_info.parent) + _try_links(cache, from_info, to_info, cache.cache_types) + + +def _cache_is_copy(cache, path_info): + """Checks whether cache uses copies.""" + if cache.cache_type_confirmed: + return cache.cache_types[0] == "copy" + + if set(cache.cache_types) <= {"copy"}: + return True + + workspace_file = path_info.with_name("." + uuid()) + test_cache_file = cache.tree.path_info / ".cache_type_test_file" + if not cache.tree.exists(test_cache_file): + cache.makedirs(test_cache_file.parent) + with cache.tree.open(test_cache_file, "wb") as fobj: + fobj.write(bytes(1)) + try: + _link(cache, test_cache_file, workspace_file) + finally: + cache.tree.remove(workspace_file) + cache.tree.remove(test_cache_file) + + cache.cache_type_confirmed = True + return cache.cache_types[0] == "copy" + + def _checkout_file( path_info, tree, obj, cache, force, progress_callback=None, relink=False, ): @@ -71,15 +155,15 @@ def _checkout_file( if not relink and _changed(path_info, tree, obj, cache): modified = True _remove(path_info, tree, cache, force=force) - cache.link(cache_info, path_info) + _link(cache, cache_info, path_info) else: - if tree.iscopy(path_info) and cache.cache_is_copy(path_info): + if tree.iscopy(path_info) and _cache_is_copy(cache, path_info): cache.unprotect(path_info) else: _remove(path_info, tree, cache, force=force) - cache.link(cache_info, path_info) + _link(cache, cache_info, path_info) else: - cache.link(cache_info, path_info) + _link(cache, cache_info, path_info) modified = True tree.state.save(path_info, obj.hash_info) diff --git a/dvc/objects.py b/dvc/objects.py index 479176721..0d04e3816 100644 --- a/dvc/objects.py +++ b/dvc/objects.py @@ -1,5 +1,9 @@ +import itertools import json import logging +from concurrent import futures + +from dvc.progress import Tqdm from .dir_info import DirInfo from .exceptions import DvcException @@ -172,8 +176,6 @@ class Tree(HashFile): return cls(obj.path_info, obj.tree, hash_info) def save(self, odb, **kwargs): - from dvc.progress import Tqdm - assert self.src_hash_info.dir_info hi = self.save_dir_info( odb, self.src_hash_info.dir_info, self.hash_info @@ -253,3 +255,86 @@ def merge(odb, ancestor_info, our_info, their_info): hash_info = Tree.save_dir_info(odb, merged) hash_info.size = _get_dir_size(odb, merged) return hash_info + + +def _transfer_file(odb, from_tree, from_info): + from dvc.utils import tmp_fname + from dvc.utils.stream import HashedStreamReader + + tmp_info = odb.tree.path_info / tmp_fname() + with from_tree.open( + from_info, mode="rb", chunk_size=from_tree.CHUNK_SIZE + ) as stream: + stream_reader = HashedStreamReader(stream) + # Since we don't know the hash beforehand, we'll + # upload it to a temporary location and then move + # it. + odb.tree.upload_fobj( + stream_reader, + tmp_info, + total=from_tree.getsize(from_info), + desc=from_info.name, + ) + + hash_info = stream_reader.hash_info + return tmp_info, hash_info + + +def _transfer_directory_contents(odb, from_tree, from_info, jobs, pbar): + rel_path_infos = {} + from_infos = from_tree.walk_files(from_info) + + def create_tasks(executor, amount): + for entry_info in itertools.islice(from_infos, amount): + pbar.total += 1 + task = executor.submit( + pbar.wrap_fn(_transfer_file), odb, from_tree, entry_info + ) + rel_path_infos[task] = entry_info.relative_to(from_info) + yield task + + pbar.total = 0 + with futures.ThreadPoolExecutor(max_workers=jobs) as executor: + tasks = set(create_tasks(executor, jobs * 5)) + + while tasks: + done, tasks = futures.wait( + tasks, return_when=futures.FIRST_COMPLETED + ) + tasks.update(create_tasks(executor, len(done))) + for task in done: + yield rel_path_infos.pop(task), task.result() + + +def _transfer_directory( + odb, from_tree, from_info, jobs, no_progress_bar=False +): + dir_info = DirInfo() + + with Tqdm(total=1, unit="Files", disable=no_progress_bar) as pbar: + for ( + entry_info, + (entry_tmp_info, entry_hash), + ) in _transfer_directory_contents( + odb, from_tree, from_info, jobs, pbar + ): + odb.add(entry_tmp_info, odb.tree, entry_hash) + dir_info.trie[entry_info.parts] = entry_hash + + return Tree.save_dir_info(odb, dir_info) + + +def transfer(odb, from_tree, from_info, jobs=None, no_progress_bar=False): + jobs = jobs or min((from_tree.jobs, odb.tree.jobs)) + + if from_tree.isdir(from_info): + return _transfer_directory( + odb, + from_tree, + from_info, + jobs=jobs, + no_progress_bar=no_progress_bar, + ) + tmp_info, hash_info = _transfer_file(odb, from_tree, from_info) + odb.add(tmp_info, odb.tree, hash_info) + return hash_info diff --git a/dvc/remote/base.py b/dvc/remote/base.py index 25729b5d3..9714a6192 100644 --- a/dvc/remote/base.py +++ b/dvc/remote/base.py @@ -512,8 +512,14 @@ class Remote: return ret def transfer(self, from_tree, from_info, jobs=None, no_progress_bar=False): - return self.cache.transfer( - from_tree, from_info, jobs=jobs, no_progress_bar=no_progress_bar + from dvc.objects import transfer + + return transfer( + self.cache, + from_tree, + from_info, + jobs=jobs, + no_progress_bar=no_progress_bar, ) @staticmethod diff --git a/dvc/repo/add.py b/dvc/repo/add.py index 3b36a3593..7a0b4974c 100644 --- a/dvc/repo/add.py +++ b/dvc/repo/add.py @@ -165,11 +165,15 @@ def _process_stages( command="add", ) else: + from dvc.objects import transfer from dvc.tree import get_cloud_tree from_tree = get_cloud_tree(repo, url=target) - out.hash_info = out.cache.transfer( - from_tree, from_tree.path_info, jobs=kwargs.get("jobs"), + out.hash_info = transfer( + out.cache, + from_tree, + from_tree.path_info, + jobs=kwargs.get("jobs"), ) out.checkout() diff --git a/dvc/system.py b/dvc/system.py index 5d57a58ed..a7c27f62b 100644 --- a/dvc/system.py +++ b/dvc/system.py @@ -2,7 +2,6 @@ import errno import logging import os import platform -import shutil import sys from dvc.exceptions import DvcException @@ -27,10 +26,6 @@ class System: def is_unix(): return os.name != "nt" - @staticmethod - def copy(src, dest): - return shutil.copyfile(src, dest) - @staticmethod def hardlink(source, link_name): try: @@ -98,6 +93,8 @@ class System: @staticmethod def reflink(source, link_name): + from dvc.utils.fs import umask + source, link_name = os.fspath(source), os.fspath(link_name) system = platform.system() @@ -116,6 +113,10 @@ class System: if ret != 0: raise DvcException("reflink is not supported") + # NOTE: reflink has a new inode, but has the same mode as the src, + # so we need to chmod it to look like a normal copy. + os.chmod(link_name, 0o666 & ~umask) + @staticmethod def _getdirinfo(path): from collections import namedtuple diff --git a/dvc/tree/local.py b/dvc/tree/local.py index 11ed25037..70709d8e4 100644 --- a/dvc/tree/local.py +++ b/dvc/tree/local.py @@ -154,7 +154,7 @@ class LocalTree(BaseTree): def copy(self, from_info, to_info): tmp_info = to_info.parent / tmp_fname("") try: - System.copy(from_info, tmp_info) + copyfile(from_info, tmp_info) os.rename(tmp_info, to_info) except Exception: self.remove(tmp_info) diff --git a/dvc/utils/fs.py b/dvc/utils/fs.py index 5f04b9436..4227e29b9 100644 --- a/dvc/utils/fs.py +++ b/dvc/utils/fs.py @@ -195,9 +195,6 @@ def copyfile(src, dest, no_progress_bar=False, name=None): try: System.reflink(src, dest) - # NOTE: reflink has a new inode, but has the same mode as the src, - # so we need to chmod it to look like a normal copy. - os.chmod(dest, 0o666 & ~umask) except DvcException: with open(src, "rb") as fsrc, open(dest, "wb+") as fdest: with Tqdm.wrapattr( diff --git a/tests/func/test_add.py b/tests/func/test_add.py index 4690877d8..971c2f28d 100644 --- a/tests/func/test_add.py +++ b/tests/func/test_add.py @@ -937,8 +937,8 @@ def test_add_file_in_symlink_dir(make_tmp_dir, tmp_dir, dvc, external): def test_add_with_cache_link_error(tmp_dir, dvc, mocker, caplog): tmp_dir.gen("foo", "foo") - mocker.patch.object( - dvc.cache.local, "_do_link", side_effect=DvcException("link failed") + mocker.patch( + "dvc.checkout._do_link", side_effect=DvcException("link failed") ) with caplog.at_level(logging.WARNING, logger="dvc"): dvc.add("foo") diff --git a/tests/func/test_s3.py b/tests/func/test_s3.py index b148be84f..4cd6a9b0a 100644 --- a/tests/func/test_s3.py +++ b/tests/func/test_s3.py @@ -55,6 +55,8 @@ def test_copy_singlepart_preserve_etag(): [S3Tree.PATH_CLS("s3://bucket/"), S3Tree.PATH_CLS("s3://bucket/ns/")], ) def test_link_created_on_non_nested_path(base_info, tmp_dir, dvc, scm): + from dvc.checkout import _link + tree = S3Tree(dvc, {"url": str(base_info.parent)}) cache = CloudCache(tree) s3 = cache.tree.s3.meta.client @@ -62,7 +64,7 @@ def test_link_created_on_non_nested_path(base_info, tmp_dir, dvc, scm): s3.put_object( Bucket=base_info.bucket, Key=(base_info / "from").path, Body="data" ) - cache.link(base_info / "from", base_info / "to") + _link(cache, base_info / "from", base_info / "to") assert cache.tree.exists(base_info / "from") assert cache.tree.exists(base_info / "to")