diff --git a/dvc/remote/gs.py b/dvc/remote/gs.py index c6e03e2e7..5285f4b5d 100644 --- a/dvc/remote/gs.py +++ b/dvc/remote/gs.py @@ -102,14 +102,17 @@ class RemoteGS(RemoteBase): 'bucket': self.bucket, 'key': posixpath.join(self.prefix, md5[0:2], md5[2:])} for md5 in md5s] - def _exists(self, gs, path_info): - return gs.bucket(path_info['bucket']).get_blob(path_info['key']) != None - def exists(self, path_infos): ret = [] gs = self.gs + + keys = [blob.name for blob in gs.bucket(self.bucket).list_blobs(prefix=self.prefix)] + for path_info in path_infos: - ret.append(self._exists(gs, path_info)) + exists = False + if path_info['key'] in keys: + exists = True + ret.append(exists) return ret def upload(self, paths, path_infos, names=None): @@ -128,9 +131,6 @@ class RemoteGS(RemoteBase): if path_info['scheme'] != 'gs': raise NotImplementedError - if not os.path.exists(path) or self._exists(gs, path_info): - continue - Logger.debug("Uploading '{}' to '{}/{}'".format(path, path_info['bucket'], path_info['key'])) @@ -166,9 +166,6 @@ class RemoteGS(RemoteBase): if path_info['scheme'] != 'gs': raise NotImplementedError - if os.path.exists(path) or not self._exists(gs, path_info): - continue - Logger.debug("Downloading '{}/{}' to '{}'".format(path_info['bucket'], path_info['key'], path)) diff --git a/dvc/remote/hdfs.py b/dvc/remote/hdfs.py index e25d69910..ea19db0f1 100644 --- a/dvc/remote/hdfs.py +++ b/dvc/remote/hdfs.py @@ -61,13 +61,6 @@ class RemoteHDFS(RemoteBase): def rm(self, path_info): self.hadoop_fs('rm {}'.format(path_info['url']), user=path_info['user']) - def exists(self, path_info): - try: - self.hadoop_fs('test -e {}'.format(path_info['url']), user=path_info['user']) - return True - except DvcException: - return False - def save_info(self, path_info): if path_info['scheme'] != 'hdfs': raise NotImplementedError @@ -133,9 +126,6 @@ class RemoteHDFS(RemoteBase): if path_info['scheme'] != 'hdfs': raise NotImplementedError - if not os.path.exists(path) or self.exists(path_info): - continue - self.hadoop_fs('mkdir -p {}'.format(posixpath.dirname(path_info['url'])), user=path_info['user']) self.hadoop_fs('copyFromLocal {} {}'.format(path, path_info['url']), user=path_info['user']) @@ -153,9 +143,6 @@ class RemoteHDFS(RemoteBase): if path_info['scheme'] != 'hdfs': raise NotImplementedError - if os.path.exists(path) or not self.exists(path_info): - continue - dname = os.path.dirname(path) if not os.path.exists(dname): os.makedirs(dname) diff --git a/dvc/remote/local.py b/dvc/remote/local.py index 042c5d60e..7a6550881 100644 --- a/dvc/remote/local.py +++ b/dvc/remote/local.py @@ -267,9 +267,6 @@ class RemoteLOCAL(RemoteBase): if path_info['scheme'] != 'local': raise NotImplementedError - if not os.path.exists(path) or os.path.exists(path_info['path']): - continue - Logger.debug("Uploading '{}' to '{}'".format(path, path_info['path'])) if not name: @@ -296,9 +293,6 @@ class RemoteLOCAL(RemoteBase): if path_info['scheme'] != 'local': raise NotImplementedError - if os.path.exists(path) or not os.path.exists(path_info['path']): - continue - Logger.debug("Downloading '{}' to '{}'".format(path_info['path'], path)) if not name: @@ -339,28 +333,18 @@ class RemoteLOCAL(RemoteBase): def status(self, checksum_infos, remote, jobs=1): checksum_infos = self._collect(checksum_infos)[0] md5s = [info[self.PARAM_MD5] for info in checksum_infos] - chunks = to_chunks(md5s, jobs) + path_infos = remote.md5s_to_path_infos(md5s) + remote_exists = remote.exists(path_infos) + local_exists = [not self.changed(md5) for md5 in md5s] - def worker(chunk): - local_exists = [os.path.exists(self.get(md5)) for md5 in chunk] - path_infos = remote.md5s_to_path_infos(chunk) - remote_exists = remote.exists(path_infos) - return [(md5, STATUS_MAP[(l,r)]) for md5, l,r in zip(chunk, local_exists, remote_exists)] - - futures = [] - with ThreadPoolExecutor(max_workers=len(chunks)) as executor: - for chunk in chunks: - resp = executor.submit(worker, chunk) - futures.append(resp) - - status = [] - for future in futures: - status.extend(future.result()) - - return status + return [(md5, STATUS_MAP[l,r]) for md5, l, r in zip(md5s, local_exists, remote_exists)] def _do_pull(self, checksum_infos, remote, jobs=1, no_progress_bar=False): md5s = [info[self.PARAM_MD5] for info in checksum_infos] + + # NOTE: filter files that are not corrupted + md5s = list(filter(lambda md5: self.changed(md5), md5s)) + cache = [self.get(md5) for md5 in md5s] path_infos = remote.md5s_to_path_infos(md5s) @@ -393,8 +377,15 @@ class RemoteLOCAL(RemoteBase): def push(self, checksum_infos, remote, jobs=1): md5s = [info[self.PARAM_MD5] for info in self._collect(checksum_infos)[0]] + # NOTE: verifying that our cache is not corrupted md5s = list(filter(lambda md5: not self.changed(md5), md5s)) + + # NOTE: filter files that are already uploaded + path_infos = remote.md5s_to_path_infos(md5s) + md5s_exist = filter(lambda x: not x[1], list(zip(md5s, remote.exists(path_infos)))) + md5s = [md5 for md5, exists in md5s_exist] + cache = [self.get(md5) for md5 in md5s] path_infos = remote.md5s_to_path_infos(md5s) diff --git a/dvc/remote/s3.py b/dvc/remote/s3.py index 88c516bcc..b56586e86 100644 --- a/dvc/remote/s3.py +++ b/dvc/remote/s3.py @@ -117,19 +117,35 @@ class RemoteS3(RemoteBase): 'bucket': self.bucket, 'key': posixpath.join(self.prefix, md5[0:2], md5[2:])} for md5 in md5s] - def _exists(self, s3, path_info): - try: - s3.Object(path_info['bucket'], path_info['key']).get() - return True - except Exception: - return False - def exists(self, path_infos): ret = [] session = boto3.session.Session() - s3 = session.resource('s3') + s3 = session.client('s3') + + keys = [] + kwargs = {'Bucket': self.bucket, + 'Prefix': self.prefix} + while True: + resp = s3.list_objects_v2(**kwargs) + contents = resp.get('Contents', None) + if not contents: + break + + for obj in contents: + keys.append(obj['Key']) + + token = resp.get('NextContinuationToken', None) + if not token: + break + + kwargs['ContinuationToken'] = token + for path_info in path_infos: - ret.append(self._exists(s3, path_info)) + exists = False + if path_info['key'] in keys: + exists = True + ret.append(exists) + return ret def upload(self, paths, path_infos, names=None): @@ -143,15 +159,12 @@ class RemoteS3(RemoteBase): assert len(names) == len(paths) session = boto3.session.Session() - s3 = session.resource('s3') + s3 = session.client('s3') for path, path_info, name in zip(paths, path_infos, names): if path_info['scheme'] != 's3': raise NotImplementedError - if not os.path.exists(path) or self._exists(s3, path_info): - continue - Logger.debug("Uploading '{}' to '{}/{}'".format(path, path_info['bucket'], path_info['key'])) @@ -163,7 +176,7 @@ class RemoteS3(RemoteBase): cb = Callback(name, total) try: - s3.Object(path_info['bucket'], path_info['key']).upload_file(path, Callback=cb) + s3.upload_file(path, path_info['bucket'], path_info['key'], Callback=cb) except Exception as exc: Logger.error("Failed to upload '{}'".format(path), exc) continue @@ -181,15 +194,12 @@ class RemoteS3(RemoteBase): assert len(names) == len(fnames) session = boto3.session.Session() - s3 = session.resource('s3') + s3 = session.client('s3') for fname, path_info, name in zip(fnames, path_infos, names): if path_info['scheme'] != 's3': raise NotImplementedError - if os.path.exists(fname) or not self._exists(s3, path_info): - continue - Logger.debug("Downloading '{}/{}' to '{}'".format(path_info['bucket'], path_info['key'], fname)) @@ -201,14 +211,14 @@ class RemoteS3(RemoteBase): if no_progress_bar: cb = None else: - total = s3.Object(bucket_name=path_info['bucket'], - key=path_info['key']).content_length + total = s3.head_object(Bucket=path_info['bucket'], + Key=path_info['key'])['ContentLength'] cb = Callback(name, total) self._makedirs(fname) try: - s3.Object(path_info['bucket'], path_info['key']).download_file(tmp_file, Callback=cb) + s3.download_file(path_info['bucket'], path_info['key'], tmp_file, Callback=cb) except Exception as exc: Logger.error("Failed to download '{}/{}'".format(path_info['bucket'], path_info['key']), exc) diff --git a/dvc/remote/ssh.py b/dvc/remote/ssh.py index 749f0b750..e23de8e14 100644 --- a/dvc/remote/ssh.py +++ b/dvc/remote/ssh.py @@ -86,10 +86,17 @@ class RemoteSSH(RemoteBase): def exists(self, path_infos): ret = [] - ssh = self.ssh(path_infos[0]['host'], path_infos[0]['user']) - for path_info in path_infos: - ret.append(self._exists(ssh, path_info)) + ssh = self.ssh(host=self.host, user=self.user) + stdout = self._exec(ssh, 'find {} -type f -follow -print'.format(self.prefix)) + plist = stdout.split() ssh.close() + + for path_info in path_infos: + exists = False + if path_info['path'] in plist: + exists = True + ret.append(exists) + return ret def _exec(self, ssh, cmd): @@ -197,9 +204,6 @@ class RemoteSSH(RemoteBase): if path_info['scheme'] != 'ssh': raise NotImplementedError - if os.path.exists(path) or not self._exists(ssh, path_info): - continue - Logger.debug("Downloading '{}/{}' to '{}'".format(path_info['host'], path_info['path'], path)) @@ -238,9 +242,6 @@ class RemoteSSH(RemoteBase): if path_info['scheme'] != 'ssh': raise NotImplementedError - if not os.path.exists(path) or self._exists(ssh, path_info): - continue - Logger.debug("Uploading '{}' to '{}/{}'".format(path, path_info['host'], path_info['path']))