From 223f0f6c0fef273a5b6d815ccdd2c7fac5558037 Mon Sep 17 00:00:00 2001 From: "Mr. Outis" Date: Fri, 14 Jun 2019 12:43:27 -0500 Subject: [PATCH] status: show a progress bar while gathering info Fix #1762 Fix #1825 --- dvc/remote/azure.py | 1 - dvc/remote/base.py | 19 +++++++++++++------ dvc/remote/gs.py | 29 ++++++++++++----------------- dvc/remote/hdfs.py | 1 - dvc/remote/http.py | 19 +++++++------------ dvc/remote/local/__init__.py | 23 ++++------------------- dvc/remote/oss.py | 1 - dvc/remote/s3.py | 30 ++++++++++++------------------ dvc/remote/ssh/__init__.py | 17 ++++++++--------- 9 files changed, 56 insertions(+), 84 deletions(-) diff --git a/dvc/remote/azure.py b/dvc/remote/azure.py index dfda92f6b..316211d4c 100644 --- a/dvc/remote/azure.py +++ b/dvc/remote/azure.py @@ -43,7 +43,6 @@ class RemoteAZURE(RemoteBASE): REQUIRES = {"azure-storage-blob": BlockBlobService} PARAM_CHECKSUM = "etag" COPY_POLL_SECONDS = 5 - SUPPORTS_NO_TRAVERSE = False def __init__(self, repo, config): super(RemoteAZURE, self).__init__(repo, config) diff --git a/dvc/remote/base.py b/dvc/remote/base.py index 2d6acb12b..d2bfd44b2 100644 --- a/dvc/remote/base.py +++ b/dvc/remote/base.py @@ -15,7 +15,7 @@ from concurrent.futures import ThreadPoolExecutor import dvc.prompt as prompt from dvc.config import Config from dvc.exceptions import DvcException, ConfirmRemoveError -from dvc.progress import progress +from dvc.progress import progress, ProgressCallback from dvc.utils import LARGE_DIR_SIZE, tmp_fname, to_chunks, move from dvc.state import StateBase from dvc.path_info import PathInfo, URLInfo @@ -68,7 +68,6 @@ class RemoteBASE(object): scheme = "base" path_cls = URLInfo REQUIRES = {} - SUPPORTS_NO_TRAVERSE = True JOBS = 4 * cpu_count() PARAM_RELPATH = "relpath" @@ -500,7 +499,7 @@ class RemoteBASE(object): def copy(self, from_info, to_info, ctx=None): raise RemoteActionNotImplemented("copy", self.scheme) - def exists(self, path_infos): + def exists(self, path_info): raise NotImplementedError @classmethod @@ -618,6 +617,7 @@ class RemoteBASE(object): - No traverse: For each given checksum, run the `exists` method and filter the checksums that aren't on the remote. This is done in parallel threads. + It also shows a progress bar when performing the check. The reason for such an odd logic is that most of the remotes take much shorter time to just retrieve everything they have under @@ -628,13 +628,20 @@ class RemoteBASE(object): Returns: A list with checksums that were found in the remote """ - if self.no_traverse and self.SUPPORTS_NO_TRAVERSE: + progress_callback = ProgressCallback(len(checksums)) + + def exists_with_progress(chunks): + return self.batch_exists(chunks, callback=progress_callback) + + if self.no_traverse and hasattr(self, "batch_exists"): with ThreadPoolExecutor(max_workers=self.JOBS) as executor: path_infos = [self.checksum_to_path_info(x) for x in checksums] chunks = to_chunks(path_infos, self.JOBS) - results = executor.map(self.exists, chunks) + results = executor.map(exists_with_progress, chunks) in_remote = itertools.chain.from_iterable(results) - return list(itertools.compress(checksums, in_remote)) + ret = list(itertools.compress(checksums, in_remote)) + progress_callback.finish("") + return ret return list(set(checksums) & set(self.all())) diff --git a/dvc/remote/gs.py b/dvc/remote/gs.py index ac7f31835..04a9aaf34 100644 --- a/dvc/remote/gs.py +++ b/dvc/remote/gs.py @@ -96,28 +96,23 @@ class RemoteGS(RemoteBASE): def list_cache_paths(self): return self._list_paths(self.path_info.bucket, self.path_info.path) - def exists(self, path_infos): - single_path = False - - if not isinstance(path_infos, list): - single_path = True - path_infos = [path_infos] + def exists(self, path_info): + paths = set(self._list_paths(path_info.bucket, path_info.path)) + return any(path_info.path == path for path in paths) + def batch_exists(self, path_infos, callback): + paths = [] gs = self.gs - paths = itertools.chain.from_iterable( - self._list_paths(path_info.bucket, path_info.path, gs) - for path_info in path_infos - ) + for path_info in path_infos: + paths.append( + self._list_paths(path_info.bucket, path_info.path, gs) + ) + callback.update(str(path_info)) - paths = set(paths) + paths = set(itertools.chain.from_iterable(paths)) - results = [path_info.path in paths for path_info in path_infos] - - if single_path and results: - return all(results) - - return results + return [path_info.path in paths for path_info in path_infos] @contextmanager def transfer_context(self): diff --git a/dvc/remote/hdfs.py b/dvc/remote/hdfs.py index 23534c000..8c51d6604 100644 --- a/dvc/remote/hdfs.py +++ b/dvc/remote/hdfs.py @@ -20,7 +20,6 @@ logger = logging.getLogger(__name__) class RemoteHDFS(RemoteBASE): scheme = Schemes.HDFS REGEX = r"^hdfs://((?P.*)@)?.*$" - SUPPORTS_NO_TRAVERSE = False PARAM_CHECKSUM = "checksum" def __init__(self, repo, config): diff --git a/dvc/remote/http.py b/dvc/remote/http.py index 3ecfe90a8..56038e8e3 100644 --- a/dvc/remote/http.py +++ b/dvc/remote/http.py @@ -93,20 +93,15 @@ class RemoteHTTP(RemoteBASE): return fails - def exists(self, path_infos): - single_path = False + def exists(self, path_info): + return bool(self._request("HEAD", path_info.url)) - if not isinstance(path_infos, list): - single_path = True - path_infos = [path_infos] + def batch_exists(self, path_infos, callback): + results = [] - results = [ - bool(self._request("HEAD", path_info.url)) - for path_info in path_infos - ] - - if single_path and results: - return all(results) + for path_info in path_infos: + results.append(bool(self._request("HEAD", path_info.url))) + callback.update(str(path_info)) return results diff --git a/dvc/remote/local/__init__.py b/dvc/remote/local/__init__.py index 898f34ab4..94c48e1c0 100644 --- a/dvc/remote/local/__init__.py +++ b/dvc/remote/local/__init__.py @@ -250,7 +250,7 @@ class RemoteLOCAL(RemoteBASE): def cache_exists(self, md5s): return [ checksum - for checksum in md5s + for checksum in progress(md5s) if not self.changed_cache_file(checksum) ] @@ -308,24 +308,12 @@ class RemoteLOCAL(RemoteBASE): logger.info( "Preparing to collect status from {}".format(remote.path_info) ) - title = "Collecting information" - - ret = {} - - progress.set_n_total(1) - progress.update_target(title, 0, 100) - - progress.update_target(title, 10, 100) - - ret = self._group(checksum_infos, show_checksums=show_checksums) + ret = self._group(checksum_infos, show_checksums=show_checksums) or {} md5s = list(ret) - progress.update_target(title, 30, 100) - + logger.info("Collecting information from local cache...") local_exists = self.cache_exists(md5s) - progress.update_target(title, 40, 100) - # This is a performance optimization. We can safely assume that, # if the resources that we want to fetch are already cached, # there's no need to check the remote storage for the existance of @@ -333,12 +321,9 @@ class RemoteLOCAL(RemoteBASE): if download and sorted(local_exists) == sorted(md5s): remote_exists = local_exists else: + logger.info("Collecting information from remote cache...") remote_exists = list(remote.cache_exists(md5s)) - progress.update_target(title, 90, 100) - - progress.finish_target(title) - self._fill_statuses(ret, local_exists, remote_exists) self._log_missing_caches(ret) diff --git a/dvc/remote/oss.py b/dvc/remote/oss.py index 12fd90cc1..79912edc3 100644 --- a/dvc/remote/oss.py +++ b/dvc/remote/oss.py @@ -43,7 +43,6 @@ class RemoteOSS(RemoteBASE): path_cls = CloudURLInfo REQUIRES = {"oss2": oss2} PARAM_CHECKSUM = "etag" - SUPPORTS_NO_TRAVERSE = False COPY_POLL_SECONDS = 5 def __init__(self, repo, config): diff --git a/dvc/remote/s3.py b/dvc/remote/s3.py index c8bc6c9bf..1e27841f8 100644 --- a/dvc/remote/s3.py +++ b/dvc/remote/s3.py @@ -224,28 +224,22 @@ class RemoteS3(RemoteBASE): def list_cache_paths(self): return self._list_paths(self.path_info.bucket, self.path_info.path) - def exists(self, path_infos): - single_path = False - - if not isinstance(path_infos, list): - single_path = True - path_infos = [path_infos] + def exists(self, path_info): + paths = self._list_paths(path_info.bucket, path_info.path) + return any(path_info.path == path for path in paths) + def batch_exists(self, path_infos, callback): + paths = [] s3 = self.s3 - paths = itertools.chain.from_iterable( - self._list_paths(path_info.bucket, path_info.path, s3) - for path_info in path_infos - ) + for path_info in path_infos: + paths.append( + self._list_paths(path_info.bucket, path_info.path, s3) + ) + callback.update(str(path_info)) - paths = set(paths) - - results = [path_info.path in paths for path_info in path_infos] - - if single_path and results: - return all(results) - - return results + paths = set(itertools.chain.from_iterable(paths)) + return [path_info.path in paths for path_info in path_infos] @contextmanager def transfer_context(self): diff --git a/dvc/remote/ssh/__init__.py b/dvc/remote/ssh/__init__.py index e9d17fdcd..b75fd6477 100644 --- a/dvc/remote/ssh/__init__.py +++ b/dvc/remote/ssh/__init__.py @@ -128,18 +128,17 @@ class RemoteSSH(RemoteBASE): password=self.password, ) - def exists(self, path_infos): - single_path = False + def exists(self, path_info): + with self.ssh(path_info) as ssh: + return ssh.exists(path_info.path) - if not isinstance(path_infos, list): - single_path = True - path_infos = [path_infos] + def batch_exists(self, path_infos, callback): + results = [] with self.ssh(path_infos[0]) as ssh: - results = [ssh.exists(path_info.path) for path_info in path_infos] - - if single_path and results: - return all(results) + for path_info in path_infos: + results.append(ssh.exists(path_info.path)) + callback.update(str(path_info)) return results