From f9db5b91d16db8978128ebfda4b3e5f2d2949d2b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?V=C3=ADt=20Novotn=C3=BD?= Date: Thu, 4 Jul 2019 12:53:40 +0200 Subject: [PATCH] remote: base: parallelize _collect_dir (#2197) * config: add core.checksum_jobs config option * remote: base: parallelize _collect_dir --- dvc/config.py | 5 +++++ dvc/remote/base.py | 55 ++++++++++++++++++++++++++++++---------------- 2 files changed, 41 insertions(+), 19 deletions(-) diff --git a/dvc/config.py b/dvc/config.py index a86777d4d..b0d280b81 100644 --- a/dvc/config.py +++ b/dvc/config.py @@ -152,6 +152,8 @@ class Config(object): # pylint: disable=too-many-instance-attributes SECTION_CORE_INTERACTIVE = "interactive" SECTION_CORE_ANALYTICS = "analytics" SECTION_CORE_ANALYTICS_SCHEMA = BOOL_SCHEMA + SECTION_CORE_CHECKSUM_JOBS = "checksum_jobs" + SECTION_CORE_CHECKSUM_JOBS_SCHEMA = And(Use(int), lambda x: x > 0) SECTION_CACHE = "cache" SECTION_CACHE_DIR = "dir" @@ -199,6 +201,9 @@ class Config(object): # pylint: disable=too-many-instance-attributes Optional( SECTION_CORE_ANALYTICS, default=True ): SECTION_CORE_ANALYTICS_SCHEMA, + Optional( + SECTION_CORE_CHECKSUM_JOBS, default=None + ): SECTION_CORE_CHECKSUM_JOBS_SCHEMA, # backward compatibility Optional(SECTION_CORE_CLOUD, default=""): SECTION_CORE_CLOUD_SCHEMA, Optional(SECTION_CORE_STORAGEPATH, default=""): str, diff --git a/dvc/remote/base.py b/dvc/remote/base.py index 58cf47995..bfe5ab874 100644 --- a/dvc/remote/base.py +++ b/dvc/remote/base.py @@ -39,6 +39,9 @@ STATUS_MAP = { } +CHECKSUM_JOBS = max(1, min(4, cpu_count() // 2)) + + class DataCloudError(DvcException): """ Data Cloud exception """ @@ -101,6 +104,11 @@ class RemoteBASE(object): ).format(url, missing, " ".join(missing), self.scheme) raise RemoteMissingDepsError(msg) + core = config.get(Config.SECTION_CORE, {}) + self.checksum_jobs = core.get( + Config.SECTION_CORE_CHECKSUM_JOBS, CHECKSUM_JOBS + ) + self.protected = False self.no_traverse = config.get(Config.SECTION_REMOTE_NO_TRAVERSE) self.state = StateBase() @@ -136,38 +144,47 @@ class RemoteBASE(object): def _collect_dir(self, path_info): dir_info = [] - for root, dirs, files in self.walk(path_info): - root_info = path_info / root + with ThreadPoolExecutor(max_workers=self.checksum_jobs) as executor: + for root, _dirs, files in self.walk(path_info): + root_info = path_info / root - if len(files) > LARGE_DIR_SIZE: - msg = ( - "Computing md5 for a large directory {}. " - "This is only done once." - ) - title = str(root_info) - logger.info(msg.format(title)) - files = progress(files, name=title) + if len(files) > LARGE_DIR_SIZE: + msg = ( + "Computing md5 for a large directory {}. " + "This is only done once." + ) + title = str(root_info) + logger.info(msg.format(title)) + files = progress(files, name=title) - for fname in files: - file_info = root_info / fname - relative_path = file_info.relative_to(path_info) - dir_info.append( + dir_info.extend( { # NOTE: this is lossy transformation: # "hey\there" -> "hey/there" # "hey/there" -> "hey/there" - # The latter is fine filename on Windows, - # which will transform to dir/file on back transform. + # The latter is fine filename on Windows, which + # will transform to dir/file on back transform. # # Yes, this is a BUG, as long as we permit "/" in # filenames on Windows and "\" on Unix - self.PARAM_RELPATH: relative_path.as_posix(), - self.PARAM_CHECKSUM: self.get_file_checksum(file_info), + self.PARAM_RELPATH: (root_info / fname) + .relative_to(path_info) + .as_posix(), + self.PARAM_CHECKSUM: executor.submit( + self.get_file_checksum, root_info / fname + ), } + for fname in files ) # NOTE: sorting the list by path to ensure reproducibility - return sorted(dir_info, key=itemgetter(self.PARAM_RELPATH)) + dir_info = sorted(dir_info, key=itemgetter(self.PARAM_RELPATH)) + + # NOTE: resolving futures + for entry in dir_info: + entry[self.PARAM_CHECKSUM] = entry[self.PARAM_CHECKSUM].result() + + return dir_info def get_dir_checksum(self, path_info): dir_info = self._collect_dir(path_info)