diff --git a/dvc/remote/base.py b/dvc/remote/base.py index 4c1bae8ee..c349a7a12 100644 --- a/dvc/remote/base.py +++ b/dvc/remote/base.py @@ -697,12 +697,13 @@ class RemoteBASE(object): def list_cache_paths(self, prefix=None, progress_callback=None): raise NotImplementedError - def all(self, *args, **kwargs): - # NOTE: The list might be way too big(e.g. 100M entries, md5 for each - # is 32 bytes, so ~3200Mb list) and we don't really need all of it at - # the same time, so it makes sense to use a generator to gradually - # iterate over it, without keeping all of it in memory. - for path in self.list_cache_paths(*args, **kwargs): + def cache_checksums(self, prefix=None, progress_callback=None): + """Iterate over remote cache checksums. + + If `prefix` is specified, only checksums which begin with `prefix` + will be returned. + """ + for path in self.list_cache_paths(prefix, progress_callback): try: yield self.path_to_checksum(path) except ValueError: @@ -710,14 +711,35 @@ class RemoteBASE(object): "'%s' doesn't look like a cache file, skipping", path ) - def gc(self, named_cache): + def all(self, jobs=None, name=None): + """Iterate over all checksums in the remote cache. + + Checksums will be fetched in parallel threads according to prefix + (except for small remotes) and a progress bar will be displayed. + """ + logger.debug( + "Fetching all checksums from '{}'".format( + name if name else "remote cache" + ) + ) + + if not self.CAN_TRAVERSE: + return self.cache_checksums() + + remote_size, remote_checksums = self._estimate_cache_size(name=name) + return self._cache_checksums_traverse( + remote_size, remote_checksums, jobs, name + ) + + def gc(self, named_cache, jobs=None): + logger.debug("named_cache: {} jobs: {}".format(named_cache, jobs)) used = self.extract_used_local_checksums(named_cache) if self.scheme != "": used.update(named_cache[self.scheme]) removed = False - for checksum in self.all(): + for checksum in self.all(jobs, str(self.path_info)): if checksum in used: continue path_info = self.checksum_to_path_info(checksum) @@ -850,7 +872,7 @@ class RemoteBASE(object): # Max remote size allowed for us to use traverse method remote_size, remote_checksums = self._estimate_cache_size( - checksums, name=name + checksums, name ) traverse_pages = remote_size / self.LIST_OBJECT_PAGE_SIZE @@ -875,32 +897,25 @@ class RemoteBASE(object): checksums - remote_checksums, jobs, name ) - if traverse_pages < 256 / self.JOBS: - # Threaded traverse will require making at least 255 more requests - # to the remote, so for small enough remotes, fetching the entire - # list at once will require fewer requests (but also take into - # account that this must be done sequentially rather than in - # parallel) - logger.debug( - "Querying {} checksums via default traverse".format( - len(checksums) - ) - ) - return list(checksums & set(self.all())) - - return self._cache_exists_traverse( - checksums, remote_checksums, remote_size, jobs, name + logger.debug( + "Querying {} checksums via traverse".format(len(checksums)) ) + remote_checksums = self._cache_checksums_traverse( + remote_size, remote_checksums, jobs, name + ) + return list(checksums & set(remote_checksums)) - def _all_with_limit(self, max_paths, prefix=None, progress_callback=None): + def _checksums_with_limit( + self, limit, prefix=None, progress_callback=None + ): count = 0 - for checksum in self.all(prefix, progress_callback): + for checksum in self.cache_checksums(prefix, progress_callback): yield checksum count += 1 - if count > max_paths: + if count > limit: logger.debug( - "`all()` returned max '{}' checksums, " - "skipping remaining results".format(max_paths) + "`cache_checksums()` returned max '{}' checksums, " + "skipping remaining results".format(limit) ) return @@ -913,33 +928,32 @@ class RemoteBASE(object): * self.LIST_OBJECT_PAGE_SIZE, ) - def _estimate_cache_size(self, checksums, short_circuit=True, name=None): + def _estimate_cache_size(self, checksums=None, name=None): """Estimate remote cache size based on number of entries beginning with "00..." prefix. """ prefix = "0" * self.TRAVERSE_PREFIX_LEN total_prefixes = pow(16, self.TRAVERSE_PREFIX_LEN) - if short_circuit: - max_remote_size = self._max_estimation_size(checksums) + if checksums: + max_checksums = self._max_estimation_size(checksums) else: - max_remote_size = None + max_checksums = None with Tqdm( desc="Estimating size of " + ("cache in '{}'".format(name) if name else "remote cache"), unit="file", - total=max_remote_size, ) as pbar: def update(n=1): pbar.update(n * total_prefixes) - if max_remote_size: - checksums = self._all_with_limit( - max_remote_size / total_prefixes, prefix, update + if max_checksums: + checksums = self._checksums_with_limit( + max_checksums / total_prefixes, prefix, update ) else: - checksums = self.all(prefix, update) + checksums = self.cache_checksums(prefix, update) remote_checksums = set(checksums) if remote_checksums: @@ -949,38 +963,60 @@ class RemoteBASE(object): logger.debug("Estimated remote size: {} files".format(remote_size)) return remote_size, remote_checksums - def _cache_exists_traverse( - self, checksums, remote_checksums, remote_size, jobs=None, name=None + def _cache_checksums_traverse( + self, remote_size, remote_checksums, jobs=None, name=None ): - logger.debug( - "Querying {} checksums via threaded traverse".format( - len(checksums) - ) - ) + """Iterate over all checksums in the remote cache. + Checksums are fetched in parallel according to prefix, except in + cases where the remote size is very small. - traverse_prefixes = ["{:02x}".format(i) for i in range(1, 256)] - if self.TRAVERSE_PREFIX_LEN > 2: - traverse_prefixes += [ - "{0:0{1}x}".format(i, self.TRAVERSE_PREFIX_LEN) - for i in range(1, pow(16, self.TRAVERSE_PREFIX_LEN - 2)) - ] + All checksums from the remote (including any from the size + estimation step passed via the `remote_checksums` argument) will be + returned. + + NOTE: For large remotes the list of checksums will be very + big(e.g. 100M entries, md5 for each is 32 bytes, so ~3200Mb list) + and we don't really need all of it at the same time, so it makes + sense to use a generator to gradually iterate over it, without + keeping all of it in memory. + """ + num_pages = remote_size / self.LIST_OBJECT_PAGE_SIZE + if num_pages < 256 / self.JOBS: + # Fetching prefixes in parallel requires at least 255 more + # requests, for small enough remotes it will be faster to fetch + # entire cache without splitting it into prefixes. + # + # NOTE: this ends up re-fetching checksums that were already + # fetched during remote size estimation + traverse_prefixes = [None] + initial = 0 + else: + yield from remote_checksums + initial = len(remote_checksums) + traverse_prefixes = ["{:02x}".format(i) for i in range(1, 256)] + if self.TRAVERSE_PREFIX_LEN > 2: + traverse_prefixes += [ + "{0:0{1}x}".format(i, self.TRAVERSE_PREFIX_LEN) + for i in range(1, pow(16, self.TRAVERSE_PREFIX_LEN - 2)) + ] with Tqdm( desc="Querying " + ("cache in '{}'".format(name) if name else "remote cache"), total=remote_size, - initial=len(remote_checksums), - unit="objects", + initial=initial, + unit="file", ) as pbar: def list_with_update(prefix): - return self.all(prefix=prefix, progress_callback=pbar.update) + return list( + self.cache_checksums( + prefix=prefix, progress_callback=pbar.update + ) + ) with ThreadPoolExecutor(max_workers=jobs or self.JOBS) as executor: in_remote = executor.map(list_with_update, traverse_prefixes,) - remote_checksums.update( - itertools.chain.from_iterable(in_remote) - ) - return list(checksums & remote_checksums) + yield from itertools.chain.from_iterable(in_remote) def _cache_object_exists(self, checksums, jobs=None, name=None): logger.debug( diff --git a/dvc/remote/local.py b/dvc/remote/local.py index 708480a13..fe0270ef3 100644 --- a/dvc/remote/local.py +++ b/dvc/remote/local.py @@ -27,6 +27,7 @@ class RemoteLOCAL(RemoteBASE): path_cls = PathInfo PARAM_CHECKSUM = "md5" PARAM_PATH = "path" + TRAVERSE_PREFIX_LEN = 2 UNPACKED_DIR_SUFFIX = ".unpacked" @@ -57,14 +58,17 @@ class RemoteLOCAL(RemoteBASE): return True def list_cache_paths(self, prefix=None, progress_callback=None): - assert prefix is None assert self.path_info is not None + if prefix: + path_info = self.path_info / prefix[:2] + else: + path_info = self.path_info if progress_callback: - for path in walk_files(self.path_info): + for path in walk_files(path_info): progress_callback() yield path else: - yield from walk_files(self.path_info) + yield from walk_files(path_info) def get(self, md5): if not md5: diff --git a/dvc/remote/ssh/__init__.py b/dvc/remote/ssh/__init__.py index ee4ad891c..7781669c4 100644 --- a/dvc/remote/ssh/__init__.py +++ b/dvc/remote/ssh/__init__.py @@ -4,6 +4,7 @@ import io import itertools import logging import os +import posixpath import threading from concurrent.futures import ThreadPoolExecutor from contextlib import closing, contextmanager @@ -44,6 +45,7 @@ class RemoteSSH(RemoteBASE): # paramiko stuff, so we would ideally have it double of server processors. # We use conservative setting of 4 instead to not exhaust max sessions. CHECKSUM_JOBS = 4 + TRAVERSE_PREFIX_LEN = 2 DEFAULT_CACHE_TYPES = ["copy"] @@ -258,15 +260,18 @@ class RemoteSSH(RemoteBASE): yield io.TextIOWrapper(fd, encoding=encoding) def list_cache_paths(self, prefix=None, progress_callback=None): - assert prefix is None + if prefix: + root = posixpath.join(self.path_info.path, prefix[:2]) + else: + root = self.path_info.path with self.ssh(self.path_info) as ssh: # If we simply return an iterator then with above closes instantly if progress_callback: - for path in ssh.walk_files(self.path_info.path): + for path in ssh.walk_files(root): progress_callback() yield path else: - yield from ssh.walk_files(self.path_info.path) + yield from ssh.walk_files(root) def walk_files(self, path_info): with self.ssh(path_info) as ssh: diff --git a/dvc/repo/gc.py b/dvc/repo/gc.py index 29f399240..8f1222fd8 100644 --- a/dvc/repo/gc.py +++ b/dvc/repo/gc.py @@ -8,8 +8,8 @@ from dvc.exceptions import InvalidArgumentError logger = logging.getLogger(__name__) -def _do_gc(typ, func, clist): - removed = func(clist) +def _do_gc(typ, func, clist, jobs=None): + removed = func(clist, jobs=jobs) if not removed: logger.info("No unused '{}' cache to remove.".format(typ)) @@ -74,22 +74,22 @@ def gc( ) ) - _do_gc("local", self.cache.local.gc, used) + _do_gc("local", self.cache.local.gc, used, jobs) if self.cache.s3: - _do_gc("s3", self.cache.s3.gc, used) + _do_gc("s3", self.cache.s3.gc, used, jobs) if self.cache.gs: - _do_gc("gs", self.cache.gs.gc, used) + _do_gc("gs", self.cache.gs.gc, used, jobs) if self.cache.ssh: - _do_gc("ssh", self.cache.ssh.gc, used) + _do_gc("ssh", self.cache.ssh.gc, used, jobs) if self.cache.hdfs: - _do_gc("hdfs", self.cache.hdfs.gc, used) + _do_gc("hdfs", self.cache.hdfs.gc, used, jobs) if self.cache.azure: - _do_gc("azure", self.cache.azure.gc, used) + _do_gc("azure", self.cache.azure.gc, used, jobs) if cloud: - _do_gc("remote", self.cloud.get_remote(remote, "gc -c").gc, used) + _do_gc("remote", self.cloud.get_remote(remote, "gc -c").gc, used, jobs) diff --git a/tests/unit/remote/test_base.py b/tests/unit/remote/test_base.py index ab799a78d..d5cad2870 100644 --- a/tests/unit/remote/test_base.py +++ b/tests/unit/remote/test_base.py @@ -49,18 +49,15 @@ class TestCmdError(TestCase, TestRemoteBASE): self.REMOTE_CLS(repo, config).remove("file") -@mock.patch.object(RemoteBASE, "_cache_exists_traverse") +@mock.patch.object(RemoteBASE, "_cache_checksums_traverse") @mock.patch.object(RemoteBASE, "_cache_object_exists") -@mock.patch.object( - RemoteBASE, "path_to_checksum", side_effect=lambda x: x, -) -def test_cache_exists(path_to_checksum, object_exists, traverse): +def test_cache_exists(object_exists, traverse): remote = RemoteBASE(None, {}) # remote does not support traverse remote.CAN_TRAVERSE = False with mock.patch.object( - remote, "list_cache_paths", return_value=list(range(256)) + remote, "cache_checksums", return_value=list(range(256)) ): checksums = list(range(1000)) remote.cache_exists(checksums) @@ -73,7 +70,7 @@ def test_cache_exists(path_to_checksum, object_exists, traverse): object_exists.reset_mock() traverse.reset_mock() with mock.patch.object( - remote, "list_cache_paths", return_value=list(range(256)) + remote, "cache_checksums", return_value=list(range(256)) ): checksums = list(range(1000)) remote.cache_exists(checksums) @@ -94,45 +91,56 @@ def test_cache_exists(path_to_checksum, object_exists, traverse): traverse.reset_mock() remote.JOBS = 16 with mock.patch.object( - remote, "list_cache_paths", return_value=list(range(256)) + remote, "cache_checksums", return_value=list(range(256)) ): checksums = list(range(1000000)) remote.cache_exists(checksums) object_exists.assert_not_called() traverse.assert_called_with( - frozenset(checksums), - set(range(256)), 256 * pow(16, remote.TRAVERSE_PREFIX_LEN), + set(range(256)), None, None, ) - # default traverse - object_exists.reset_mock() - traverse.reset_mock() - remote.TRAVERSE_WEIGHT_MULTIPLIER = 1 - with mock.patch.object(remote, "list_cache_paths", return_value=[0]): - checksums = set(range(1000000)) - remote.cache_exists(checksums) - traverse.assert_not_called() - object_exists.assert_not_called() - @mock.patch.object( - RemoteBASE, "list_cache_paths", return_value=[], + RemoteBASE, "cache_checksums", return_value=[], ) @mock.patch.object( RemoteBASE, "path_to_checksum", side_effect=lambda x: x, ) -def test_cache_exists_traverse(path_to_checksum, list_cache_paths): +def test_cache_checksums_traverse(path_to_checksum, cache_checksums): remote = RemoteBASE(None, {}) remote.path_info = PathInfo("foo") - remote._cache_exists_traverse({0}, set(), 4096) + + # parallel traverse + size = 256 / remote.JOBS * remote.LIST_OBJECT_PAGE_SIZE + list(remote._cache_checksums_traverse(size, {0})) for i in range(1, 16): - list_cache_paths.assert_any_call( + cache_checksums.assert_any_call( prefix="{:03x}".format(i), progress_callback=CallableOrNone ) for i in range(1, 256): - list_cache_paths.assert_any_call( + cache_checksums.assert_any_call( prefix="{:02x}".format(i), progress_callback=CallableOrNone ) + + # default traverse (small remote) + size -= 1 + cache_checksums.reset_mock() + list(remote._cache_checksums_traverse(size - 1, {0})) + cache_checksums.assert_called_with( + prefix=None, progress_callback=CallableOrNone + ) + + +def test_cache_checksums(): + remote = RemoteBASE(None, {}) + remote.path_info = PathInfo("foo") + + with mock.patch.object( + remote, "list_cache_paths", return_value=["12/3456", "bar"] + ): + checksums = list(remote.cache_checksums()) + assert checksums == ["123456"]