ssh: increase jobs and parallelize batch_exists

- SSH number of jobs increased from 4 to 8.

- `batch_exists` now opens all the SFTP sessions available to
distribute the work between them. Number of sessions is restricted
by the `MaxSessions` option configurable by the `sshd_config`.

- Fix #2104

With the optimizations mentioned above, we increased the speed by 8x
(using the default `MaxSessions` configuration -- 10).

Shout-outs to @efiop, @shcheklein, and @dmpetrov for suggesting the
optimization techniques.

dvc/logger.py

@@ -176,7 +176,11 @@ def setup(level=logging.INFO):
                 "dvc": {
                     "level": level,
                     "handlers": ["console", "console_errors"],
-                }
+                },
+                "paramiko": {
+                    "level": logging.CRITICAL,
+                    "handlers": ["console", "console_errors"],
+                },
             },
         }
     )

dvc/main.py

@@ -34,6 +34,7 @@ def main(argv=None):
 
         elif args.verbose:
             logger.setLevel(logging.DEBUG)
+            logging.getLogger("paramiko").setLevel(logging.DEBUG)
 
         cmd = args.func(args)
         ret = cmd.run_cmd()

dvc/remote/ssh/__init__.py

@@ -3,6 +3,9 @@ from __future__ import unicode_literals
 import os
 import getpass
 import logging
+import itertools
+import errno
+from concurrent.futures import ThreadPoolExecutor
 
 try:
     import paramiko
@@ -12,6 +15,7 @@ except ImportError:
 import dvc.prompt as prompt
 from dvc.remote.ssh.connection import SSHConnection
 from dvc.config import Config
+from dvc.utils import to_chunks
 from dvc.utils.compat import urlparse, StringIO
 from dvc.remote.base import RemoteBASE
 from dvc.scheme import Schemes
@@ -24,7 +28,7 @@ class RemoteSSH(RemoteBASE):
     scheme = Schemes.SSH
     REQUIRES = {"paramiko": paramiko}
 
-    JOBS = 4
+    JOBS = 8
     PARAM_CHECKSUM = "md5"
     DEFAULT_PORT = 22
     TIMEOUT = 1800
@@ -133,14 +137,32 @@ class RemoteSSH(RemoteBASE):
             return ssh.exists(path_info.path)
 
     def batch_exists(self, path_infos, callback):
-        results = []
+        def _exists(chunk_and_channel):
+            chunk, channel = chunk_and_channel
+            ret = []
+            for path in chunk:
+                try:
+                    channel.stat(path)
+                    ret.append(True)
+                except IOError as exc:
+                    if exc.errno != errno.ENOENT:
+                        raise
+                    ret.append(False)
+                callback.update(path)
+            return ret
 
         with self.ssh(path_infos[0]) as ssh:
-            for path_info in path_infos:
-                results.append(ssh.exists(path_info.path))
-                callback.update(str(path_info))
+            with ssh.open_max_sftp_channels() as channels:
+                max_workers = len(channels)
 
-        return results
+                with ThreadPoolExecutor(max_workers=max_workers) as executor:
+                    paths = [path_info.path for path_info in path_infos]
+                    chunks = to_chunks(paths, max_workers)
+                    chunks_and_channels = zip(chunks, channels)
+                    outcome = executor.map(_exists, chunks_and_channels)
+                    results = list(itertools.chain.from_iterable(outcome))
+
+            return results
 
     def get_file_checksum(self, path_info):
         if path_info.scheme != self.scheme:

dvc/remote/ssh/connection.py

@@ -1,8 +1,8 @@
-import errno
 import os
 import posixpath
 import logging
-from stat import S_ISDIR
+import stat
+from contextlib import contextmanager
 
 try:
     import paramiko
@@ -10,6 +10,7 @@ except ImportError:
     paramiko = None
 
 from dvc.utils import tmp_fname
+from dvc.utils.compat import ignore_file_not_found
 from dvc.progress import progress
 from dvc.exceptions import DvcException
 from dvc.remote.base import RemoteCmdError
@@ -72,40 +73,25 @@ class SSHConnection:
 
         self._ssh.close()
 
-    def exists(self, path):
+    def st_mode(self, path):
         self._sftp_connect()
-        try:
-            return self._sftp.stat(path)
-        except IOError:
-            return False
-        pass
+
+        with ignore_file_not_found():
+            return self._sftp.stat(path).st_mode
+
+        return 0
+
+    def exists(self, path, sftp=None):
+        return bool(self.st_mode(path))
 
     def isdir(self, path):
-        from stat import S_ISDIR
-
-        self._sftp_connect()
-        try:
-            return S_ISDIR(self._sftp.stat(path).st_mode)
-        except IOError:
-            return False
+        return stat.S_ISDIR(self.st_mode(path))
 
     def isfile(self, path):
-        from stat import S_ISREG
-
-        self._sftp_connect()
-        try:
-            return S_ISREG(self._sftp.stat(path).st_mode)
-        except IOError:
-            return False
+        return stat.S_ISREG(self.st_mode(path))
 
     def islink(self, path):
-        from stat import S_ISLNK
-
-        self._sftp_connect()
-        try:
-            return S_ISLNK(self._sftp.stat(path).st_mode)
-        except IOError:
-            return False
+        return stat.S_ISLNK(self.st_mode(path))
 
     def makedirs(self, path):
         self._sftp_connect()
@@ -148,7 +134,7 @@ class SSHConnection:
         nondirs = []
         for entry in dir_entries:
             name = entry.filename
-            if S_ISDIR(entry.st_mode):
+            if stat.S_ISDIR(entry.st_mode):
                 dirs.append(name)
             else:
                 nondirs.append(name)
@@ -170,11 +156,8 @@ class SSHConnection:
                 yield posixpath.join(root, fname)
 
     def _remove_file(self, path):
-        try:
+        with ignore_file_not_found():
             self._sftp.remove(path)
-        except IOError as exc:
-            if exc.errno != errno.ENOENT:
-                raise
 
     def _remove_dir(self, path):
         for root, dirs, files in self.walk(path, topdown=False):
@@ -185,11 +168,9 @@ class SSHConnection:
             for dname in dirs:
                 path = posixpath.join(root, dname)
                 self._sftp.rmdir(dname)
-        try:
+
+        with ignore_file_not_found():
             self._sftp.rmdir(path)
-        except IOError as exc:
-            if exc.errno != errno.ENOENT:
-                raise
 
     def remove(self, path):
         self._sftp_connect()
@@ -311,3 +292,19 @@ class SSHConnection:
     def cp(self, src, dest):
         self.makedirs(posixpath.dirname(dest))
         self.execute("cp {} {}".format(src, dest))
+
+    @contextmanager
+    def open_max_sftp_channels(self):
+        try:
+            channels = []
+            while True:
+                try:
+                    channels.append(self._ssh.open_sftp())
+                except paramiko.ssh_exception.ChannelException:
+                    if not channels:
+                        raise
+                    break
+            yield channels
+        finally:
+            for channel in channels:
+                channel.close()

dvc/utils/compat.py

@@ -4,6 +4,7 @@ from __future__ import absolute_import
 import sys
 import os
 import errno
+from contextlib import contextmanager
 
 # Syntax sugar.
 _ver = sys.version_info
@@ -90,6 +91,15 @@ def _makedirs(name, mode=0o777, exist_ok=False):
             raise
 
 
+@contextmanager
+def ignore_file_not_found():
+    try:
+        yield
+    except IOError as exc:
+        if exc.errno != errno.ENOENT:
+            raise
+
+
 if is_py2:
     from urlparse import urlparse, urljoin, urlsplit, urlunsplit  # noqa: F401
     from BaseHTTPServer import HTTPServer  # noqa: F401