From dda8368f20e85a296ebfabc79d5ce4e4469c0b57 Mon Sep 17 00:00:00 2001 From: "Mr. Outis" Date: Thu, 13 Jun 2019 20:36:42 -0500 Subject: [PATCH] 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 | 6 ++- dvc/main.py | 1 + dvc/remote/ssh/__init__.py | 34 ++++++++++++++--- dvc/remote/ssh/connection.py | 73 +++++++++++++++++------------------- dvc/utils/compat.py | 10 +++++ 5 files changed, 79 insertions(+), 45 deletions(-) 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 + + 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() 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