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 <[email protected]>

Co-authored-by: Ruslan Kuprieiev <[email protected]>
This commit is contained in:
Batuhan Taskaya
2021-04-19 18:09:35 +03:00
committed by GitHub
co-authored by Ruslan Kuprieiev
parent a8b7fda5a2
commit 68897aa0de
12 changed files with 225 additions and 143 deletions
+7 -2
View File
@@ -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,
+20 -3
View File
@@ -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)
+28 -6
View File
@@ -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)
+2 -1
View File
@@ -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
)
+63 -28
View File
@@ -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)
-92
View File
@@ -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
+27 -4
View File
@@ -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 = [
+12 -3
View File
@@ -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)
+1
View File
@@ -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)
+1 -1
View File
@@ -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"
)
+1 -1
View File
@@ -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"])
+63 -2
View File
@@ -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",
}