RFC/perf: optimise push/pull/status

Signed-off-by: Ruslan Kuprieiev <[email protected]>

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))

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)

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)
 

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)

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']))