From 68897aa0def4e509a27d75f2e78106495d08e0c4 Mon Sep 17 00:00:00 2001 From: Batuhan Taskaya Date: Mon, 19 Apr 2021 18:09:35 +0300 Subject: [PATCH] update: only update the missing hashes with --to-remote (#5773) * objects: merge transfer() with stage() * objects: drop streamfile * objects: preserve the signature of get_file_hash * objects: override hash_jobs if jobs= passed * objects: support no progress bar directive * tests: spy on _get_file_hash * objects: update=... directive * update: only transfer missing caches with --to-remote * tests: remove local_remote from the primary test (due to it moves) * objects: add move=True/False option * objects: implement _get_file_obj on top of get_file_hash * objects: use threading on save() * objects: pass jobs= on to-cache as well * objects: use in place tqdm * objects: check for FileExistsError * Update dvc/objects/__init__.py Co-authored-by: Ruslan Kuprieiev Co-authored-by: Ruslan Kuprieiev --- dvc/data_cloud.py | 9 +++- dvc/objects/__init__.py | 23 +++++++-- dvc/objects/db/base.py | 34 ++++++++++--- dvc/objects/file.py | 3 +- dvc/objects/stage.py | 91 +++++++++++++++++++++++----------- dvc/objects/transfer.py | 92 ----------------------------------- dvc/remote/base.py | 31 ++++++++++-- dvc/repo/add.py | 15 ++++-- dvc/repo/diff.py | 1 + dvc/stage/imports.py | 2 +- tests/func/test_data_cloud.py | 2 +- tests/func/test_update.py | 65 ++++++++++++++++++++++++- 12 files changed, 225 insertions(+), 143 deletions(-) delete mode 100644 dvc/objects/transfer.py diff --git a/dvc/data_cloud.py b/dvc/data_cloud.py index d94811080..e77ab626e 100644 --- a/dvc/data_cloud.py +++ b/dvc/data_cloud.py @@ -92,12 +92,15 @@ class DataCloud: show_checksums=show_checksums, ) - def transfer(self, source, jobs=None, remote=None, command=None): + def transfer( + self, source, jobs=None, update=False, remote=None, command=None + ): """Transfer data items in a cloud-agnostic way. Args: source (str): url for the source location. jobs (int): number of jobs that can be running simultaneously. + update (bool): whether to update existing data or sync fully remote (dvc.remote.base.BaseRemote): optional remote to compare cache to. By default remote from core.remote config option is used. @@ -108,7 +111,9 @@ class DataCloud: from_fs = get_cloud_fs(self.repo, url=source) remote = self.get_remote(remote, command) - return remote.transfer(from_fs, from_fs.path_info, jobs=jobs) + return remote.transfer( + from_fs, from_fs.path_info, jobs=jobs, update=update + ) def status( self, diff --git a/dvc/objects/__init__.py b/dvc/objects/__init__.py index a580030bc..fafacf339 100644 --- a/dvc/objects/__init__.py +++ b/dvc/objects/__init__.py @@ -1,4 +1,5 @@ import logging +from concurrent.futures import ThreadPoolExecutor, as_completed from dvc.progress import Tqdm @@ -7,10 +8,26 @@ from .tree import Tree logger = logging.getLogger(__name__) -def save(odb, obj, **kwargs): +def save(odb, obj, jobs=None, **kwargs): if isinstance(obj, Tree): - for _, entry in Tqdm(obj): - odb.add(entry.path_info, entry.fs, entry.hash_info, **kwargs) + with ThreadPoolExecutor(max_workers=jobs) as executor: + for future in Tqdm( + as_completed( + executor.submit( + odb.add, + entry.path_info, + entry.fs, + entry.hash_info, + **kwargs + ) + for _, entry in obj + ), + desc="Saving files", + total=len(obj), + unit="file", + ): + future.result() + odb.add(obj.path_info, obj.fs, obj.hash_info, **kwargs) diff --git a/dvc/objects/db/base.py b/dvc/objects/db/base.py index 885f214cd..0f7a4b0c7 100644 --- a/dvc/objects/db/base.py +++ b/dvc/objects/db/base.py @@ -1,5 +1,6 @@ import itertools import logging +import os from concurrent.futures import ThreadPoolExecutor from copy import copy from typing import Optional @@ -42,7 +43,7 @@ class ObjectDB: hash_info, ) - def add(self, path_info, fs, hash_info, **kwargs): + def add(self, path_info, fs, hash_info, move=True, **kwargs): try: self.check(hash_info) return @@ -52,11 +53,32 @@ class ObjectDB: cache_info = self.hash_to_path_info(hash_info.value) # using our makedirs to create dirs with proper permissions self.makedirs(cache_info.parent) - if isinstance(fs, type(self.fs)): - self.fs.move(path_info, cache_info) - else: - with fs.open(path_info, mode="rb") as fobj: - self.fs.upload_fobj(fobj, cache_info) + use_move = isinstance(fs, type(self.fs)) and move + try: + if use_move: + self.fs.move(path_info, cache_info) + else: + with fs.open(path_info, mode="rb") as fobj: + self.fs.upload_fobj(fobj, cache_info) + except OSError as exc: + # If the target file already exists, we are going to simply + # ignore the exception (#4992). + # + # On Windows, it is not always guaranteed that you'll get + # FileExistsError (it might be PermissionError or a bare OSError) + # but all of those exceptions raised from the original + # FileExistsError so we have a separate check for that. + if isinstance(exc, FileExistsError) or ( + os.name == "nt" + and exc.__context__ + and isinstance(exc.__context__, FileExistsError) + ): + logger.debug("'%s' file already exists, skipping", path_info) + if use_move: + fs.remove(path_info) + else: + raise + self.protect(cache_info) self.fs.repo.state.save(cache_info, self.fs, hash_info) diff --git a/dvc/objects/file.py b/dvc/objects/file.py index 85b1ef916..6c5fdd66e 100644 --- a/dvc/objects/file.py +++ b/dvc/objects/file.py @@ -1,7 +1,6 @@ import logging from .errors import ObjectFormatError -from .stage import get_file_hash logger = logging.getLogger(__name__) @@ -37,6 +36,8 @@ class HashFile: ) def check(self, odb): + from .stage import get_file_hash + actual = get_file_hash( self.path_info, self.fs, self.hash_info.name, odb.repo.state ) diff --git a/dvc/objects/stage.py b/dvc/objects/stage.py index ac8313f46..ef29f9a4e 100644 --- a/dvc/objects/stage.py +++ b/dvc/objects/stage.py @@ -1,14 +1,31 @@ import errno import os from concurrent.futures import ThreadPoolExecutor +from functools import partial 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 +def _upload_file(path_info, fs, odb): + from dvc.utils import tmp_fname + from dvc.utils.stream import HashedStreamReader + + tmp_info = odb.fs.path_info / tmp_fname() + with fs.open(path_info, mode="rb", chunk_size=fs.CHUNK_SIZE) as stream: + stream = HashedStreamReader(stream) + odb.fs.upload_fobj( + stream, tmp_info, desc=path_info.name, total=fs.getsize(path_info) + ) + + obj = HashFile(tmp_info, odb.fs, stream.hash_info) + return path_info, obj + + def _get_file_hash(path_info, fs, name): info = fs.info(path_info) if name in info: @@ -49,43 +66,64 @@ def get_file_hash(path_info, fs, name, state=None): return hash_info -def _calculate_hashes(path_info, fs, name, state, **kwargs): - def _hash(file_info): - return file_info, get_file_hash(file_info, fs, name, state) +def _get_file_obj(path_info, fs, name, odb=None, state=None, upload=False): + if upload: + assert odb and name == "md5" + return _upload_file(path_info, fs, odb) + obj = HashFile( + path_info, fs, get_file_hash(path_info, fs, name, state=state) + ) + return path_info, obj + + +def _build_objects(path_info, fs, name, odb, state, upload, **kwargs): with Tqdm( - unit="md5", desc="Computing file/dir hashes (only done once)", + unit="md5", + desc="Computing file/dir hashes (only done once)", + disable=kwargs.pop("no_progress_bar", False), ) as pbar: - worker = pbar.wrap_fn(_hash) - with ThreadPoolExecutor(max_workers=fs.hash_jobs) as executor: - pairs = executor.map(worker, fs.walk_files(path_info, **kwargs)) - return dict(pairs) + worker = pbar.wrap_fn( + partial( + _get_file_obj, + fs=fs, + name=name, + odb=odb, + state=state, + upload=upload, + ) + ) + with ThreadPoolExecutor( + max_workers=kwargs.pop("jobs", fs.hash_jobs) + ) as executor: + yield from executor.map(worker, fs.walk_files(path_info, **kwargs)) -def _iter_hashes(path_info, fs, name, state, **kwargs): - if name in fs.DETAIL_FIELDS: +def _iter_objects(path_info, fs, name, odb, state, upload, **kwargs): + if not upload and name in fs.DETAIL_FIELDS: for details in fs.ls(path_info, recursive=True, detail=True): file_info = path_info.replace(path=details["name"]) hash_info = HashInfo( name, details[name], size=details.get("size"), ) - yield file_info, hash_info + yield file_info, HashFile(file_info, fs, hash_info) return None - yield from _calculate_hashes(path_info, fs, name, state, **kwargs).items() + yield from _build_objects( + path_info, fs, name, odb, state, upload, **kwargs + ) -def _build_tree(path_info, fs, name, state, **kwargs): - from .file import HashFile +def _build_tree(path_info, fs, name, odb, state, upload, **kwargs): from .tree import Tree tree = Tree(None, None, None) - for fi, hi in _iter_hashes(path_info, fs, name, state, **kwargs): - if DvcIgnore.DVCIGNORE_FILE == fi.name: - raise DvcIgnoreInCollectedDirError(fi.parent) - - obj = HashFile(fi, fs, hi) + for file_info, obj in _iter_objects( + path_info, fs, name, odb, state, upload, **kwargs + ): + if DvcIgnore.DVCIGNORE_FILE == file_info.name: + raise DvcIgnoreInCollectedDirError(file_info.parent) # NOTE: this is lossy transformation: # "hey\there" -> "hey/there" @@ -95,14 +133,14 @@ def _build_tree(path_info, fs, name, state, **kwargs): # # Yes, this is a BUG, as long as we permit "/" in # filenames on Windows and "\" on Unix - tree.add(fi.relative_to(path_info).parts, obj) + tree.add(file_info.relative_to(path_info).parts, obj) tree.digest() return tree -def _get_tree_obj(path_info, fs, name, odb, state, **kwargs): +def _get_tree_obj(path_info, fs, name, odb, state, upload, **kwargs): from .tree import Tree value = fs.info(path_info).get(name) @@ -113,7 +151,7 @@ def _get_tree_obj(path_info, fs, name, odb, state, **kwargs): except FileNotFoundError: pass - tree = _build_tree(path_info, fs, name, state, **kwargs) + tree = _build_tree(path_info, fs, name, odb, state, upload, **kwargs) odb.add(tree.path_info, tree.fs, tree.hash_info) if name != "md5": @@ -132,7 +170,7 @@ def _get_tree_obj(path_info, fs, name, odb, state, **kwargs): return tree -def stage(odb, path_info, fs, name, **kwargs): +def stage(odb, path_info, fs, name, upload=False, **kwargs): assert path_info and ( isinstance(path_info, str) or path_info.scheme == fs.scheme ) @@ -174,12 +212,9 @@ def stage(odb, path_info, fs, name, **kwargs): return obj if fs.isdir(path_info): - obj = _get_tree_obj(path_info, fs, name, odb, state, **kwargs) + obj = _get_tree_obj(path_info, fs, name, odb, state, upload, **kwargs) else: - from .file import HashFile - - hash_info = get_file_hash(path_info, fs, name, state) - obj = HashFile(path_info, fs, hash_info) + _, obj = _get_file_obj(path_info, fs, name, odb, state, upload) if obj.hash_info and fs.exists(path_info): state.save(path_info, fs, obj.hash_info) diff --git a/dvc/objects/transfer.py b/dvc/objects/transfer.py deleted file mode 100644 index c6fc0268e..000000000 --- a/dvc/objects/transfer.py +++ /dev/null @@ -1,92 +0,0 @@ -import itertools -import logging -from concurrent import futures - -from dvc.progress import Tqdm - -from . import save -from .file import HashFile -from .tree import Tree - -logger = logging.getLogger(__name__) - - -def _transfer_file(odb, from_fs, from_info): - from dvc.utils import tmp_fname - from dvc.utils.stream import HashedStreamReader - - tmp_info = odb.fs.path_info / tmp_fname() - with from_fs.open( - from_info, mode="rb", chunk_size=from_fs.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.fs.upload_fobj( - stream_reader, - tmp_info, - total=from_fs.getsize(from_info), - desc=from_info.name, - ) - - hash_info = stream_reader.hash_info - return tmp_info, hash_info - - -def _transfer_directory_contents(odb, from_fs, from_info, jobs, pbar): - rel_path_infos = {} - from_infos = from_fs.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_fs, 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_fs, from_info, jobs, no_progress_bar=False): - tree = Tree(None, None, None) - - 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_fs, from_info, jobs, pbar): - obj = HashFile(entry_tmp_info, odb.fs, entry_hash) - tree.add(entry_info.parts, obj) - - tree.digest() - save(odb, tree) - return tree.hash_info - - -def transfer(odb, from_fs, from_info, jobs=None, no_progress_bar=False): - jobs = jobs or min((from_fs.jobs, odb.fs.jobs)) - - if from_fs.isdir(from_info): - return _transfer_directory( - odb, - from_fs, - from_info, - jobs=jobs, - no_progress_bar=no_progress_bar, - ) - tmp_info, hash_info = _transfer_file(odb, from_fs, from_info) - odb.add(tmp_info, odb.fs, hash_info) - return hash_info diff --git a/dvc/remote/base.py b/dvc/remote/base.py index 583883980..204a81bba 100644 --- a/dvc/remote/base.py +++ b/dvc/remote/base.py @@ -8,6 +8,8 @@ from functools import partial, wraps from dvc.exceptions import DownloadError, UploadError from dvc.hash_info import HashInfo +from dvc.objects import save +from dvc.objects.stage import stage from ..progress import Tqdm from .index import RemoteIndex, RemoteIndexNoop @@ -505,17 +507,38 @@ class Remote: return ret - def transfer(self, from_fs, from_info, jobs=None, no_progress_bar=False): - from dvc.objects.transfer import transfer + def transfer( + self, + from_fs, + from_info, + jobs=None, + update=False, + no_progress_bar=False, + ): + # When running import-url --to-remote / add --to-remote/-o ... we + # assume that it is unlikely that the odb will contain majority of the + # hashes, so we transfer everything as is (even if that file might + # already be in the cache) and don't waste an upload to scan the layout + # of the source location. But when doing update --to-remote, there is + # a high probability that the odb might contain some of the hashes, so + # we first calculate all the hashes (but don't transfer anything) and + # then only update the missing cache files. - return transfer( + upload = not (update and from_fs.isdir(from_info)) + jobs = jobs or min((from_fs.jobs, self.odb.fs.jobs)) + obj = stage( self.odb, - from_fs, from_info, + from_fs, + "md5", + upload=upload, jobs=jobs, no_progress_bar=no_progress_bar, ) + save(self.odb, obj, jobs=jobs, move=upload) + return obj.hash_info + @staticmethod def _log_missing_caches(hash_info_dict): missing_caches = [ diff --git a/dvc/repo/add.py b/dvc/repo/add.py index f8176bb7b..6041ccf9c 100644 --- a/dvc/repo/add.py +++ b/dvc/repo/add.py @@ -173,12 +173,21 @@ def _process_stages( ) else: from dvc.fs import get_cloud_fs - from dvc.objects.transfer import transfer + from dvc.objects import save as osave + from dvc.objects.stage import stage as ostage from_fs = get_cloud_fs(repo, url=target) - out.hash_info = transfer( - out.odb, from_fs, from_fs.path_info, jobs=kwargs.get("jobs"), + jobs = kwargs.get("jobs", min((from_fs.jobs, out.odb.fs.jobs))) + obj = ostage( + out.odb, + from_fs.path_info, + from_fs, + "md5", + upload=True, + jobs=jobs, ) + osave(out.odb, obj, jobs=jobs, move=False) + out.hash_info = obj.hash_info out.checkout() Dvcfile(repo, stage.path).dump(stage) diff --git a/dvc/repo/diff.py b/dvc/repo/diff.py index 75c15cd96..75a8e39af 100644 --- a/dvc/repo/diff.py +++ b/dvc/repo/diff.py @@ -163,6 +163,7 @@ def _dir_output_paths(repo_fs, output, targets=None): if targets is None or any( fname.isin_or_eq(target) for target in targets ): + # pylint: disable=no-member yield str(fname), get_file_hash(fname, repo_fs, "md5").value except NoRemoteError: logger.warning("dir cache entry for '%s' is missing", output) diff --git a/dvc/stage/imports.py b/dvc/stage/imports.py index c5e4cb24b..3a5298033 100644 --- a/dvc/stage/imports.py +++ b/dvc/stage/imports.py @@ -14,7 +14,7 @@ def _update_import_on_remote(stage, remote, jobs): url = stage.deps[0].def_path stage.outs[0].hash_info = stage.repo.cloud.transfer( - url, jobs=jobs, remote=remote, command="update" + url, jobs=jobs, update=True, remote=remote, command="update" ) diff --git a/tests/func/test_data_cloud.py b/tests/func/test_data_cloud.py index 958c1c161..b5b7b229d 100644 --- a/tests/func/test_data_cloud.py +++ b/tests/func/test_data_cloud.py @@ -231,7 +231,7 @@ def test_warn_on_outdated_stage(tmp_dir, dvc, local_remote, caplog): def test_hash_recalculation(mocker, dvc, tmp_dir, local_remote): tmp_dir.gen({"foo": "foo"}) - test_get_file_hash = mocker.spy(dvc_module.objects.stage, "get_file_hash") + test_get_file_hash = mocker.spy(dvc_module.objects.stage, "_get_file_hash") ret = main(["config", "cache.type", "hardlink"]) assert ret == 0 ret = main(["add", "foo"]) diff --git a/tests/func/test_update.py b/tests/func/test_update.py index d4893e41a..f97ddcb16 100644 --- a/tests/func/test_update.py +++ b/tests/func/test_update.py @@ -358,12 +358,26 @@ def test_update_import_url_to_remote(tmp_dir, dvc, workspace, local_remote): indirect=True, ) def test_update_import_url_to_remote_directory( - tmp_dir, dvc, workspace, local_remote + mocker, tmp_dir, dvc, workspace, local_remote ): workspace.gen({"data": {"foo": "foo", "bar": {"baz": "baz"}}}) stage = dvc.imp_url("remote://workspace/data", to_remote=True) - workspace.gen({"data": {"foo2": "foo2", "bar": {"baz2": "baz2"}}}) + workspace.gen( + { + "data": { + "foo2": "foo2", + "bar": {"baz2": "baz2"}, + "repeated_hashes": { + "foo": "foo", + "baz": "baz", + "foo_with_different_name": "foo", + }, + } + } + ) + + upload_file_mock = mocker.spy(type(dvc.odb.local.fs), "upload_fobj") stage = dvc.update(stage.path, to_remote=True) dvc.pull("data") @@ -371,4 +385,51 @@ def test_update_import_url_to_remote_directory( "foo": "foo", "foo2": "foo2", "bar": {"baz": "baz", "baz2": "baz2"}, + "repeated_hashes": { + "foo": "foo", + "baz": "baz", + "foo_with_different_name": "foo", + }, + } + # 2 new hashes (foo2, baz2) + 1 .dir hash + assert upload_file_mock.mock.call_count == 3 + + +def test_update_import_url_to_remote_directory_changed_contents( + tmp_dir, dvc, local_cloud, local_remote +): + local_cloud.gen({"data": {"foo": "foo", "bar": {"baz": "baz"}}}) + stage = dvc.imp_url("remote://upstream/data", to_remote=True) + + local_cloud.gen( + {"data": {"foo": "not_foo", "foo2": "foo", "bar": {"baz2": "baz2"}}} + ) + stage = dvc.update(stage.path, to_remote=True) + + dvc.pull("data") + assert (tmp_dir / "data").read_text() == { + "foo": "not_foo", + "foo2": "foo", + "bar": {"baz": "baz", "baz2": "baz2"}, + } + + +def test_update_import_url_to_remote_directory_same_hash( + tmp_dir, dvc, local_cloud, local_remote +): + local_cloud.gen( + {"data": {"foo": "foo", "bar": {"baz": "baz"}, "same": "same"}} + ) + stage = dvc.imp_url("remote://upstream/data", to_remote=True) + + local_cloud.gen( + {"data": {"foo": "baz", "bar": {"baz": "foo"}, "same": "same"}} + ) + stage = dvc.update(stage.path, to_remote=True) + + dvc.pull("data") + assert (tmp_dir / "data").read_text() == { + "foo": "baz", + "bar": {"baz": "foo"}, + "same": "same", }