diff --git a/dvc/dependency/base.py b/dvc/dependency/base.py index 62637cb14..4700e8fc8 100644 --- a/dvc/dependency/base.py +++ b/dvc/dependency/base.py @@ -60,7 +60,7 @@ class DependencyBase(object): @property def exists(self): - return self.remote.exists([self.path_info])[0] + return self.remote.exists(self.path_info) def changed(self): raise NotImplementedError diff --git a/dvc/remote/azure.py b/dvc/remote/azure.py index 228c68551..4b5ee2a68 100644 --- a/dvc/remote/azure.py +++ b/dvc/remote/azure.py @@ -49,7 +49,7 @@ class RemoteAzure(RemoteBase): or match.group('container_name') # backward compatibility or os.getenv('AZURE_STORAGE_CONTAINER_NAME')) - self.prefix = urlparse(self.url).path.lstrip('/') or '/' + self.prefix = urlparse(self.url).path.lstrip('/') self.connection_string = ( config.get(Config.SECTION_AZURE_CONNECTION_STRING) @@ -145,29 +145,39 @@ class RemoteAzure(RemoteBase): return [{ 'scheme': self.scheme, 'bucket': self.bucket, - 'key': '{}/{}/{}'.format(self.prefix, md5[0:2], md5[2:]) + 'key': posixpath.join(self.prefix, md5[0:2], md5[2:]) } for md5 in md5s] - def _all_keys(self, bucket, prefix): - return {blob.name - for blob in self.blob_service.list_blobs(bucket, - prefix=prefix)} + def _list_keys(self, bucket, prefix): + blob_service = self.blob_service + next_marker = None + while True: + blobs = blob_service.list_blobs(bucket, + prefix=prefix, + marker=next_marker) - def exists(self, path_infos): - ret = [] + for blob in blobs: + yield blob.name - if len(path_infos) == 0: - return ret + if not blobs.next_marker: + break - bucket = path_infos[0]['bucket'] - key = path_infos[0]['key'] + next_marker = blobs.next_marker - for path_info in path_infos: - assert path_info['scheme'] == self.scheme - assert path_info['bucket'] == bucket + def cache_exists(self, md5s): + assert isinstance(md5s, list) - key = path_info['key'] - ret.append(key in self._all_keys(bucket, key)) + if len(md5s) == 0: + return [] + + ret = len(md5s) * [False] + keys = [posixpath.join(self.prefix, md5[0:2], md5[2:]) for md5 in md5s] + for key in self._list_keys(self.bucket, self.prefix): + for i, k in enumerate(keys): + if k == key: + ret[i] = True + + assert len(ret) == len(keys) == len(md5s) return ret @@ -246,8 +256,14 @@ class RemoteAzure(RemoteBase): return posixpath.dirname(relpath) + posixpath.basename(relpath) def _all(self): - keys = self._all_keys(self.bucket, self.prefix) - return [self._path_to_etag(key) for key in keys] + # 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. + return ( + self._path_to_etag(key) + for key in self._list_keys(self.bucket, self.prefix) + ) def gc(self, cinfos): from dvc.remote.local import RemoteLOCAL diff --git a/dvc/remote/base.py b/dvc/remote/base.py index 02ab6b588..1d7f82128 100644 --- a/dvc/remote/base.py +++ b/dvc/remote/base.py @@ -131,6 +131,9 @@ class RemoteBase(object): def md5s_to_path_infos(self, md5s): raise NotImplementedError + def cache_exists(self, md5s): + raise NotImplementedError + def exists(self, path_infos): raise NotImplementedError diff --git a/dvc/remote/gs.py b/dvc/remote/gs.py index 23830071f..78a9f75eb 100644 --- a/dvc/remote/gs.py +++ b/dvc/remote/gs.py @@ -33,6 +33,10 @@ class RemoteGS(RemoteBase): self.url = config.get(Config.SECTION_REMOTE_URL, storagepath) self.projectname = config.get(Config.SECTION_GCP_PROJECTNAME, None) + parsed = urlparse(self.url) + self.bucket = parsed.netloc + self.prefix = parsed.path.lstrip('/') + @staticmethod def compat_config(config): ret = config.copy() @@ -40,14 +44,6 @@ class RemoteGS(RemoteBase): ret[Config.SECTION_REMOTE_URL] = url return ret - @property - def bucket(self): - return urlparse(self.url).netloc - - @property - def prefix(self): - return urlparse(self.url).path.lstrip('/') - @property def gs(self): return storage.Client() @@ -82,7 +78,7 @@ class RemoteGS(RemoteBase): cache = {'scheme': 'gs', 'bucket': self.bucket, 'key': key} if {self.PARAM_MD5: md5} != self.save_info(cache): - if self.exists([cache])[0]: + if self.exists(cache): msg = 'Corrupted cache file {}' logger.warn(msg.format(self.to_string(cache))) self.remove(cache) @@ -91,7 +87,7 @@ class RemoteGS(RemoteBase): return False def changed(self, path_info, checksum_info): - if not self.exists([path_info])[0]: + if not self.exists(path_info): return True md5 = checksum_info.get(self.PARAM_MD5, None) @@ -146,7 +142,7 @@ class RemoteGS(RemoteBase): logger.warn(msg.format(md5, self.to_string(path_info))) return - if self.exists([path_info])[0]: + if self.exists(path_info): msg = "Data '{}' exists. Removing before checkout." logger.warn(msg.format(self.to_string(path_info))) self.remove(path_info) @@ -179,26 +175,32 @@ class RemoteGS(RemoteBase): 'key': posixpath.join(self.prefix, md5[0:2], md5[2:])} for md5 in md5s] - def exists(self, path_infos): - ret = [] + def _list_keys(self, bucket, prefix): + for blob in self.gs.bucket(bucket).list_blobs(prefix=prefix): + yield blob.name - if len(path_infos) == 0: - return ret + def exists(self, path_info): + assert not isinstance(path_info, list) + assert path_info['scheme'] == 'gs' - gs = self.gs - bucket_name = path_infos[0]['bucket'] - bucket = gs.bucket(bucket_name) + keys = self._list_keys(path_info['bucket'], path_info['key']) + return any(path_info['key'] == key for key in keys) - for path_info in path_infos: - assert self.scheme == path_info['scheme'] - assert bucket_name == path_info['bucket'] + def cache_exists(self, md5s): + assert isinstance(md5s, list) + + if len(md5s) == 0: + return [] + + ret = len(md5s) * [False] + keys = [posixpath.join(self.prefix, md5[0:2], md5[2:]) for md5 in md5s] + for key in self._list_keys(self.bucket, self.prefix): + for i, k in enumerate(keys): + if k == key: + ret[i] = True + + assert len(ret) == len(keys) == len(md5s) - exists = False - key = path_info['key'] - keys = [blob.name for blob in bucket.list_blobs(prefix=key)] - if key in keys: - exists = True - ret.append(exists) return ret def upload(self, from_infos, to_infos, names=None): @@ -292,9 +294,14 @@ class RemoteGS(RemoteBase): return posixpath.dirname(relpath) + posixpath.basename(relpath) def _all_md5s(self): - blobs = self.gs.bucket(self.bucket).list_blobs(prefix=self.prefix) - blobs = list(blobs) - return [self._path_to_md5(blob.name) for blob in blobs] + # 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. + return ( + self._path_to_md5(key) + for key in self._list_keys(self.bucket, self.prefix) + ) def gc(self, cinfos): used = [info[self.PARAM_MD5] for info in cinfos['gs']] diff --git a/dvc/remote/hdfs.py b/dvc/remote/hdfs.py index dfc1f331c..3b7158b2d 100644 --- a/dvc/remote/hdfs.py +++ b/dvc/remote/hdfs.py @@ -93,7 +93,7 @@ class RemoteHDFS(RemoteBase): cache['url'] = posixpath.join(self.url, checksum[0:2], checksum[2:]) if {self.PARAM_CHECKSUM: checksum} != self.save_info(cache): - if self.exists([cache])[0]: + if self.exists(cache): msg = 'Corrupted cache file {}' logger.warn(msg.format(self.to_string(cache))) self.remove(cache) @@ -102,7 +102,7 @@ class RemoteHDFS(RemoteBase): return False def changed(self, path_info, checksum_info): - if not self.exists([path_info])[0]: + if not self.exists(path_info): return True checksum = checksum_info.get(self.PARAM_CHECKSUM, None) @@ -148,7 +148,7 @@ class RemoteHDFS(RemoteBase): logger.warn(msg.format(checksum, self.to_string(path_info))) return - if self.exists([path_info])[0]: + if self.exists(path_info): msg = "Data '{}' exists. Removing before checkout." logger.warn(msg.format(self.to_string(path_info))) self.remove(path_info) @@ -178,20 +178,28 @@ class RemoteHDFS(RemoteBase): 'url': posixpath.join(self.url, md5[0:2], md5[2:])} for md5 in md5s] - def exists(self, path_infos): - ret = [] + def exists(self, path_info): + assert not isinstance(path_info, list) + assert path_info['scheme'] == 'hdfs' - if len(path_infos) == 0: - return ret + try: + self.hadoop_fs('test -e {}'.format(path_info['url'])) + return True + except RemoteHDFSCmdError: + return False - for path_info in path_infos: - assert path_info['scheme'] == self.scheme - try: - self.hadoop_fs('test -e {}'.format(path_info['url'])) - exists = True - except RemoteHDFSCmdError: - exists = False - ret.append(exists) + def cache_exists(self, md5s): + assert isinstance(md5s, list) + + if len(md5s) == 0: + return [] + + existing_md5s = self._all_checksums() + ret = len(md5s) * [False] + for existing_md5 in existing_md5s: + for i, md5 in enumerate(md5s): + if md5 == existing_md5: + ret[i] = True return ret @@ -242,6 +250,11 @@ class RemoteHDFS(RemoteBase): return posixpath.dirname(relpath) + posixpath.basename(relpath) def _all_checksums(self): + try: + self.hadoop_fs('test -e {}'.format(self.url)) + except RemoteHDFSCmdError: + return [] + stdout = self.hadoop_fs('ls -R {}'.format(self.url)) lines = stdout.split('\n') flist = [] diff --git a/dvc/remote/http.py b/dvc/remote/http.py index 50bf1910c..4eb6c82a9 100644 --- a/dvc/remote/http.py +++ b/dvc/remote/http.py @@ -82,11 +82,19 @@ class RemoteHTTP(RemoteBase): if not no_progress_bar: progress.finish_target(name) - def exists(self, path_infos): - return [ - bool(self._request('HEAD', path_info.get('url'))) - for path_info in path_infos - ] + def exists(self, path_info): + assert not isinstance(path_info, list) + assert path_info['scheme'] in ['http', 'https'] + return bool(self._request('HEAD', path_info.get('url'))) + + def cache_exists(self, md5s): + assert isinstance(md5s, list) + ret = [] + for md5 in md5s: + url = posixpath.join(self.prefix, md5[0:2], md5[2:]) + exists = bool(self._request('HEAD', url)) + ret.append(exists) + return ret def save_info(self, path_info): if path_info['scheme'] not in ['http', 'https']: diff --git a/dvc/remote/local.py b/dvc/remote/local.py index 369820fab..d6a01599f 100644 --- a/dvc/remote/local.py +++ b/dvc/remote/local.py @@ -101,6 +101,11 @@ class RemoteLOCAL(RemoteBase): return True return False + def exists(self, path_info): + assert not isinstance(path_info, list) + assert path_info['scheme'] == 'local' + return os.path.exists(path_info['path']) + def changed_cache(self, md5): cache = self.get(md5) clist = [(cache, md5)] @@ -465,7 +470,7 @@ class RemoteLOCAL(RemoteBase): - The checkusm stored in the State is different from the given one - There's no file in the cache """ - if not self.exists([path_info])[0]: + if not self.exists(path_info): return True md5 = checksum_info.get(self.PARAM_MD5, None) @@ -494,11 +499,12 @@ class RemoteLOCAL(RemoteBase): 'path': os.path.join(self.prefix, md5[0:2], md5[2:])} for md5 in md5s] - def exists(self, path_infos): + def cache_exists(self, md5s): + assert isinstance(md5s, list) ret = [] - for path_info in path_infos: - assert path_info['scheme'] == 'local' - ret.append(os.path.exists(path_info['path'])) + for md5 in md5s: + path = os.path.join(self.prefix, md5[0:2], md5[2:]) + ret.append(os.path.exists(path)) return ret def upload(self, from_infos, to_infos, names=None): @@ -610,13 +616,9 @@ class RemoteLOCAL(RemoteBase): md5s, names = self._group(checksum_infos, show_checksums=show_checksums) - progress.update_target(title, 20, 100) - - path_infos = remote.md5s_to_path_infos(md5s) - progress.update_target(title, 30, 100) - remote_exists = remote.exists(path_infos) + remote_exists = remote.cache_exists(md5s) progress.update_target(title, 90, 100) @@ -663,7 +665,7 @@ class RemoteLOCAL(RemoteBase): # NOTE: dummy call to try to establish a connection # to see if we need to ask user for a password. - remote.exists(remote.md5s_to_path_infos(['000'])) + remote.cache_exists(['000']) progress.update_target(title, 70, 100) @@ -711,7 +713,7 @@ class RemoteLOCAL(RemoteBase): # NOTE: filter files that are already uploaded md5s = [i[self.PARAM_MD5] for i in checksum_infos] - exists = remote.exists(remote.md5s_to_path_infos(md5s)) + exists = remote.cache_exists(md5s) progress.update_target(title, 30, 100) diff --git a/dvc/remote/s3.py b/dvc/remote/s3.py index 3242f42a9..633568fb4 100644 --- a/dvc/remote/s3.py +++ b/dvc/remote/s3.py @@ -66,6 +66,10 @@ class RemoteS3(RemoteBase): if shared_creds: os.environ.setdefault('AWS_SHARED_CREDENTIALS_FILE', shared_creds) + parsed = urlparse(self.url) + self.bucket = parsed.netloc + self.prefix = parsed.path.lstrip('/') + @staticmethod def compat_config(config): ret = config.copy() @@ -73,14 +77,6 @@ class RemoteS3(RemoteBase): ret[Config.SECTION_REMOTE_URL] = url return ret - @property - def bucket(self): - return urlparse(self.url).netloc - - @property - def prefix(self): - return urlparse(self.url).path.lstrip('/') - @property def s3(self): session = boto3.session.Session(profile_name=self.profile, @@ -106,7 +102,7 @@ class RemoteS3(RemoteBase): path_info['key'])} def changed(self, path_info, checksum_info): - if not self.exists([path_info])[0]: + if not self.exists(path_info): return True etag = checksum_info.get(self.PARAM_ETAG, None) @@ -146,7 +142,7 @@ class RemoteS3(RemoteBase): cache = {'scheme': 's3', 'bucket': self.bucket, 'key': key} if {self.PARAM_ETAG: etag} != self.save_info(cache): - if self.exists([cache])[0]: + if self.exists(cache): msg = 'Corrupted cache file {}' logger.warn(msg.format(self.to_string(cache))) self.remove(cache) @@ -172,7 +168,7 @@ class RemoteS3(RemoteBase): logger.warn(msg.format(etag, self.to_string(path_info))) return - if self.exists([path_info])[0]: + if self.exists(path_info): msg = "Data '{}' exists. Removing before checkout." logger.warn(msg.format(self.to_string(path_info))) self.remove(path_info) @@ -202,10 +198,9 @@ class RemoteS3(RemoteBase): 'key': posixpath.join(self.prefix, md5[0:2], md5[2:])} for md5 in md5s] - def _all_keys(self, bucket, prefix): + def _list_keys(self, bucket, prefix): s3 = self.s3 - keys = [] kwargs = {'Bucket': bucket, 'Prefix': prefix} while True: @@ -219,7 +214,7 @@ class RemoteS3(RemoteBase): break for obj in contents: - keys.append(obj['Key']) + yield obj['Key'] token = resp.get('NextContinuationToken', None) if not token: @@ -227,24 +222,27 @@ class RemoteS3(RemoteBase): kwargs['ContinuationToken'] = token - return keys + def exists(self, path_info): + assert not isinstance(path_info, list) + assert path_info['scheme'] == 's3' - def exists(self, path_infos): - ret = [] + keys = self._list_keys(path_info['bucket'], path_info['key']) + return any(path_info['key'] == key for key in keys) - if len(path_infos) == 0: - return ret + def cache_exists(self, md5s): + assert isinstance(md5s, list) - bucket = path_infos[0]['bucket'] + if len(md5s) == 0: + return [] - for path_info in path_infos: - assert path_info['scheme'] == self.scheme - assert path_info['bucket'] == bucket - exists = False - key = path_info['key'] - if key in self._all_keys(bucket, key): - exists = True - ret.append(exists) + ret = len(md5s) * [False] + keys = [posixpath.join(self.prefix, md5[0:2], md5[2:]) for md5 in md5s] + for key in self._list_keys(self.bucket, self.prefix): + for i, k in enumerate(keys): + if k == key: + ret[i] = True + + assert len(keys) == len(ret) == len(md5s) return ret @@ -342,8 +340,14 @@ class RemoteS3(RemoteBase): return posixpath.dirname(relpath) + posixpath.basename(relpath) def _all(self): - keys = self._all_keys(self.bucket, self.prefix) - return [self._path_to_etag(key) for key in keys] + # 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. + return ( + self._path_to_etag(key) + for key in self._list_keys(self.bucket, self.prefix) + ) def gc(self, cinfos): used_etags = [info[self.PARAM_ETAG] for info in cinfos['s3']] diff --git a/dvc/remote/ssh.py b/dvc/remote/ssh.py index 755468728..922be2205 100644 --- a/dvc/remote/ssh.py +++ b/dvc/remote/ssh.py @@ -105,31 +105,33 @@ class RemoteSSH(RemoteBase): return ssh - def exists(self, path_infos): - ret = [] + def exists(self, path_info): + assert not isinstance(path_info, list) + assert path_info['scheme'] == 'ssh' - if len(path_infos) == 0: - return ret - - host = path_infos[0]['host'] - user = path_infos[0]['user'] - port = path_infos[0]['port'] - ssh = self.ssh(host=host, - user=user, - port=port) - - for path_info in path_infos: - assert host == path_info['host'] - assert user == path_info['user'] - assert port == path_info['port'] + with self.ssh(path_info['host'], + path_info['user'], + path_info['port']) as ssh: try: self._exec(ssh, 'test -e {}'.format(path_info['path'])) exists = True except RemoteSSHCmdError: exists = False - ret.append(exists) - ssh.close() + return exists + + def cache_exists(self, md5s): + assert isinstance(md5s, list) + + if len(md5s) == 0: + return [] + + existing_md5s = self._all_md5s() + ret = len(md5s) * [False] + for existing_md5 in existing_md5s: + for i, md5 in enumerate(md5s): + if md5 == existing_md5: + ret[i] = True return ret @@ -245,7 +247,7 @@ class RemoteSSH(RemoteBase): cache['path'] = posixpath.join(self.prefix, md5[0:2], md5[2:]) if {self.PARAM_MD5: md5} != self.save_info(cache): - if self.exists([cache])[0]: + if self.exists(cache): msg = 'Corrupted cache file {}' logger.warn(msg.format(self.to_string(cache))) self.remove(cache) @@ -254,7 +256,7 @@ class RemoteSSH(RemoteBase): return False def changed(self, path_info, checksum_info): - if not self.exists([path_info])[0]: + if not self.exists(path_info): return True md5 = checksum_info.get(self.PARAM_MD5, None) @@ -296,7 +298,7 @@ class RemoteSSH(RemoteBase): logger.warn(msg.format(md5, self.to_string(path_info))) return - if self.exists([path_info])[0]: + if self.exists(path_info): msg = "Data '{}' exists. Removing before checkout." logger.warn(msg.format(self.to_string(path_info))) self.remove(path_info)