Repository: iterative/dvc
Issue: ssh: check IOError key code

In some functions of `SSHConnection` we are catching `IOError` but we are only interested in catching an specific errno `ENOENT`.

Changed files:
dvc/logger.py, dvc/remote/ssh/__init__.py, dvc/main.py, dvc/utils/compat.py, dvc/remote/ssh/connection.py

diff --git a/dvc/logger.py b/dvc/logger.py
index 1c71a0e4f..35e873738 100644
--- a/dvc/logger.py
+++ b/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"],
+                },
             },
         }
     )
diff --git a/dvc/main.py b/dvc/main.py
index f00dc0ced..cb162fb7d 100644
--- a/dvc/main.py
+++ b/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()
diff --git a/dvc/remote/ssh/__init__.py b/dvc/remote/ssh/__init__.py
index b75fd6477..4c5cdaeb1 100644
--- a/dvc/remote/ssh/__init__.py
+++ b/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:
diff --git a/dvc/remote/ssh/connection.py b/dvc/remote/ssh/connection.py
index b5ea969d0..6bcd2ed3e 100644
--- a/dvc/remote/ssh/connection.py
+++ b/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
 
-    def isdir(self, path):
-        from stat import S_ISDIR
+        with ignore_file_not_found():
+            return self._sftp.stat(path).st_mode
 
-        self._sftp_connect()
-        try:
-            return S_ISDIR(self._sftp.stat(path).st_mode)
-        except IOError:
-            return False
+        return 0
 
-    def isfile(self, path):
-        from stat import S_ISREG
+    def exists(self, path, sftp=None):
+        return bool(self.st_mode(path))
 
-        self._sftp_connect()
-        try:
-            return S_ISREG(self._sftp.stat(path).st_mode)
-        except IOError:
-            return False
+    def isdir(self, path):
+        return stat.S_ISDIR(self.st_mode(path))
 
-    def islink(self, path):
-        from stat import S_ISLNK
+    def isfile(self, path):
+        return stat.S_ISREG(self.st_mode(path))
 
-        self._sftp_connect()
-        try:
-            return S_ISLNK(self._sftp.stat(path).st_mode)
-        except IOError:
-            return False
+    def islink(self, path):
+        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()
diff --git a/dvc/utils/compat.py b/dvc/utils/compat.py
index 0abb08c4a..4c4c95cf7 100644
--- a/dvc/utils/compat.py
+++ b/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