diff --git a/dvc/config.py b/dvc/config.py index b6f68759f..a86777d4d 100644 --- a/dvc/config.py +++ b/dvc/config.py @@ -255,6 +255,7 @@ class Config(object): # pylint: disable=too-many-instance-attributes SECTION_REMOTE_TIMEOUT = "timeout" SECTION_REMOTE_PASSWORD = "password" SECTION_REMOTE_ASK_PASSWORD = "ask_password" + SECTION_REMOTE_NO_TRAVERSE = "no_traverse" SECTION_REMOTE_SCHEMA = { SECTION_REMOTE_URL: str, Optional(SECTION_AWS_REGION): str, @@ -278,6 +279,7 @@ class Config(object): # pylint: disable=too-many-instance-attributes Optional(SECTION_OSS_ACCESS_KEY_SECRET): str, Optional(SECTION_OSS_ENDPOINT): str, Optional(PRIVATE_CWD): str, + Optional(SECTION_REMOTE_NO_TRAVERSE, default=True): BOOL_SCHEMA, } SECTION_STATE = "state" diff --git a/dvc/remote/azure.py b/dvc/remote/azure.py index 913dc62c9..b31c65c5c 100644 --- a/dvc/remote/azure.py +++ b/dvc/remote/azure.py @@ -44,6 +44,7 @@ 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 dacb84e00..8f1be8257 100644 --- a/dvc/remote/base.py +++ b/dvc/remote/base.py @@ -6,14 +6,16 @@ import os import json import logging import tempfile +import itertools from operator import itemgetter from multiprocessing import cpu_count +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.utils import LARGE_DIR_SIZE, tmp_fname +from dvc.utils import LARGE_DIR_SIZE, tmp_fname, to_chunks from dvc.state import StateBase from dvc.path_info import PathInfo, URLInfo @@ -65,6 +67,7 @@ class RemoteBASE(object): scheme = "base" path_cls = URLInfo REQUIRES = {} + SUPPORTS_NO_TRAVERSE = True JOBS = 4 * cpu_count() PARAM_RELPATH = "relpath" @@ -99,8 +102,8 @@ class RemoteBASE(object): raise RemoteMissingDepsError(msg) self.protected = False + self.no_traverse = config.get(Config.SECTION_REMOTE_NO_TRAVERSE) self.state = StateBase() - self._dir_info = {} def __repr__(self): @@ -508,15 +511,36 @@ class RemoteBASE(object): return self.changed_cache_file(checksum) def cache_exists(self, checksums): - # NOTE: The reason for such an odd logic is that most of the remotes - # take much shorter time to just retrieve everything they have under - # a certain prefix(e.g. s3, gs, ssh, hdfs). Other remotes that can - # check if particular file exists much quicker, use their own - # implementation of cache_exists(see http, local). - # - # Result of all() might be way too big, so we should walk through - # it in one pass. - return list(filter(lambda checksum: checksum in checksums, self.all())) + """Check if the given checksums are stored in the remote. + + There are two ways of performing this check: + + - Traverse: Get a list of all the files in the remote + (traversing the cache directory) and compare it with + the given checksums. + + - 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. + + The reason for such an odd logic is that most of the remotes + take much shorter time to just retrieve everything they have under + a certain prefix (e.g. s3, gs, ssh, hdfs). Other remotes that can + check if particular file exists much quicker, use their own + implementation of cache_exists (see http, local). + + Returns: + A list with checksums that were found in the remote + """ + if self.no_traverse and self.SUPPORTS_NO_TRAVERSE: + 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) + in_remote = itertools.chain.from_iterable(results) + return list(itertools.compress(checksums, in_remote)) + + return list(set(checksums) & set(self.all())) def already_cached(self, path_info): current = self.get_checksum(path_info) diff --git a/dvc/remote/gs.py b/dvc/remote/gs.py index dbc40e3b6..519da7962 100644 --- a/dvc/remote/gs.py +++ b/dvc/remote/gs.py @@ -1,6 +1,7 @@ from __future__ import unicode_literals import logging +import itertools from dvc.scheme import Schemes @@ -90,19 +91,37 @@ class RemoteGS(RemoteBASE): blob.delete() - def _list_paths(self, bucket, prefix): - for blob in self.gs.bucket(bucket).list_blobs(prefix=prefix): + def _list_paths(self, bucket, prefix, gs=None): + gs = gs or self.gs + + for blob in gs.bucket(bucket).list_blobs(prefix=prefix): yield blob.name def list_cache_paths(self): return self._list_paths(self.path_info.bucket, self.path_info.path) - def exists(self, path_info): - assert not isinstance(path_info, list) - assert path_info.scheme == "gs" + def exists(self, path_infos): + single_path = False - paths = self._list_paths(path_info.bucket, path_info.path) - return any(path_info.path == path for path in paths) + if not isinstance(path_infos, list): + single_path = True + path_infos = [path_infos] + + gs = self.gs + + paths = itertools.chain.from_iterable( + self._list_paths(path_info.bucket, path_info.path, gs) + for path_info in path_infos + ) + + 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 def upload(self, from_infos, to_infos, names=None, no_progress_bar=False): names = self._verify_path_args(to_infos, from_infos, names) diff --git a/dvc/remote/hdfs.py b/dvc/remote/hdfs.py index 27d710786..206d61bda 100644 --- a/dvc/remote/hdfs.py +++ b/dvc/remote/hdfs.py @@ -20,6 +20,7 @@ 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 81734773c..20498cbf3 100644 --- a/dvc/remote/http.py +++ b/dvc/remote/http.py @@ -89,10 +89,22 @@ class RemoteHTTP(RemoteBASE): if not no_progress_bar: progress.finish_target(name) - def exists(self, path_info): - assert not isinstance(path_info, list) - assert path_info.scheme == self.scheme - return bool(self._request("HEAD", path_info.url)) + def exists(self, path_infos): + single_path = False + + if not isinstance(path_infos, list): + single_path = True + path_infos = [path_infos] + + results = [ + bool(self._request("HEAD", path_info.url)) + for path_info in path_infos + ] + + if single_path and results: + return all(results) + + return results def cache_exists(self, md5s): assert isinstance(md5s, list) diff --git a/dvc/remote/local/__init__.py b/dvc/remote/local/__init__.py index afe15a602..e33a59851 100644 --- a/dvc/remote/local/__init__.py +++ b/dvc/remote/local/__init__.py @@ -248,8 +248,11 @@ class RemoteLOCAL(RemoteBASE): move(tmp, outp) def cache_exists(self, md5s): - assert isinstance(md5s, list) - return list(filter(lambda md5: not self.changed_cache_file(md5), md5s)) + return [ + checksum + for checksum in md5s + if not self.changed_cache_file(checksum) + ] def upload(self, from_infos, to_infos, names=None, no_progress_bar=False): names = self._verify_path_args(to_infos, from_infos, names) @@ -367,7 +370,7 @@ class RemoteLOCAL(RemoteBASE): progress.update_target(title, 10, 100) ret = self._group(checksum_infos, show_checksums=show_checksums) - md5s = list(ret.keys()) + md5s = list(ret) progress.update_target(title, 30, 100) diff --git a/dvc/remote/oss.py b/dvc/remote/oss.py index 2270f9404..4dc1818e8 100644 --- a/dvc/remote/oss.py +++ b/dvc/remote/oss.py @@ -46,6 +46,7 @@ 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 de271c28a..e8087a8a9 100644 --- a/dvc/remote/s3.py +++ b/dvc/remote/s3.py @@ -3,6 +3,7 @@ from __future__ import unicode_literals import os import threading import logging +import itertools from dvc.scheme import Schemes @@ -207,9 +208,9 @@ class RemoteS3(RemoteBASE): logger.debug("Removing {}".format(path_info)) self.s3.delete_object(Bucket=path_info.bucket, Key=path_info.path) - def _list_paths(self, bucket, prefix): + def _list_paths(self, bucket, prefix, s3=None): """ Read config for list object api, paginate through list objects.""" - s3 = self.s3 + s3 = s3 or self.s3 kwargs = {"Bucket": bucket, "Prefix": prefix} if self.list_objects: list_objects_api = "list_objects" @@ -226,12 +227,28 @@ class RemoteS3(RemoteBASE): def list_cache_paths(self): return self._list_paths(self.path_info.bucket, self.path_info.path) - def exists(self, path_info): - assert not isinstance(path_info, list) - assert path_info.scheme == "s3" + def exists(self, path_infos): + single_path = False - paths = self._list_paths(path_info.bucket, path_info.path) - return any(path_info.path == path for path in paths) + if not isinstance(path_infos, list): + single_path = True + path_infos = [path_infos] + + s3 = self.s3 + + paths = itertools.chain.from_iterable( + self._list_paths(path_info.bucket, path_info.path, s3) + for path_info in path_infos + ) + + 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 def upload(self, from_infos, to_infos, names=None, no_progress_bar=False): names = self._verify_path_args(to_infos, from_infos, names) diff --git a/dvc/remote/ssh/__init__.py b/dvc/remote/ssh/__init__.py index 205163adf..a067db341 100644 --- a/dvc/remote/ssh/__init__.py +++ b/dvc/remote/ssh/__init__.py @@ -129,12 +129,20 @@ class RemoteSSH(RemoteBASE): password=self.password, ) - def exists(self, path_info): - assert not isinstance(path_info, list) - assert path_info.scheme == self.scheme + def exists(self, path_infos): + single_path = False - 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] + + 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) + + return results def get_file_checksum(self, path_info): if path_info.scheme != self.scheme: