Stone_Red/dvc
1
0
Fork 0
mirror of https://github.com/Stone-Red-Code/dvc.git synced 2026-11-10 15:28:34 +01:00
Code Issues Packages Projects Releases Wiki Activity

remote: use progress bar for remote cache query status during dvc gc (#3559)

Browse Source 
* remote: use progress bar for remote cache query status during `dvc gc`

* test that cache_checksums() handles unexpected paths

* Fix SSHMocked tests
This commit is contained in:
Peter Rowlands (변기호)
2020-04-01 18:27:54 +03:00
committed by GitHub
parent b02babee9e
commit ee0066fae4
5 changed files with 151 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dvc/remote/base.py
+94 -58
View File
@@ -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(
dvc/remote/local.py
+7 -3
View File
@@ -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:
dvc/remote/ssh/__init__.py
+8 -3
View File
@@ -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:
dvc/repo/gc.py
+9 -9
View File
@@ -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)
tests/unit/remote/test_base.py
+33 -25
View File
@@ -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"]