diff --git a/dvc/remote/base.py b/dvc/remote/base.py index d02ba83be..8ce47a9ff 100644 --- a/dvc/remote/base.py +++ b/dvc/remote/base.py @@ -10,7 +10,7 @@ import tempfile import itertools from operator import itemgetter from multiprocessing import cpu_count -from concurrent.futures import as_completed, ThreadPoolExecutor +from concurrent.futures import ThreadPoolExecutor import dvc.prompt as prompt from dvc.config import Config @@ -20,14 +20,7 @@ from dvc.exceptions import ( DvcIgnoreInCollectedDirError, ) from dvc.progress import progress, ProgressCallback -from dvc.utils import ( - LARGE_DIR_SIZE, - tmp_fname, - to_chunks, - move, - relpath, - as_future, -) +from dvc.utils import LARGE_DIR_SIZE, tmp_fname, to_chunks, move, relpath from dvc.state import StateBase from dvc.path_info import PathInfo, URLInfo @@ -147,64 +140,61 @@ class RemoteBASE(object): def get_file_checksum(self, path_info): raise NotImplementedError - def _get_file_checksum_future(self, file_info, executor): - checksum = self.state.get(file_info) - if checksum: - checksum = as_future(checksum) - else: - checksum = executor.submit(self.get_file_checksum, file_info) - return checksum - - def _get_dir_info(self, dir_path_info): - dir_info = {} + def _calculate_checksums(self, file_infos): + file_infos = list(file_infos) with ThreadPoolExecutor(max_workers=self.checksum_jobs) as executor: - for root, _dirs, files in self.walk(dir_path_info): - root_info = dir_path_info / root + tasks = executor.map(self.get_file_checksum, file_infos) - for fname in files: - - if fname == DvcIgnore.DVCIGNORE_FILE: - raise DvcIgnoreInCollectedDirError(root) - - file_info = root_info / fname - - checksum = self._get_file_checksum_future( - file_info, executor - ) - - relative_path = file_info.relative_to(dir_path_info) - dir_info[checksum] = { - # 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. - # - # Yes, this is a BUG, as long as we permit "/" in - # filenames on Windows and "\" on Unix - self.PARAM_RELPATH: relative_path.as_posix() - } - - checksums = as_completed(dir_info) - if len(dir_info) > LARGE_DIR_SIZE: + if len(file_infos) > LARGE_DIR_SIZE: msg = ( "Computing md5 for a large number of files. " "This is only done once." ) logger.info(msg) - checksums = progress(checksums, total=len(dir_info)) + tasks = progress(tasks, total=len(file_infos)) - # Resolving futures - for checksum in checksums: - entry = dir_info[checksum] - entry[self.PARAM_CHECKSUM] = checksum.result() - return dir_info + checksums = { + file_infos[index]: task for index, task in enumerate(tasks) + } + return checksums def _collect_dir(self, path_info): - dir_info = self._get_dir_info(path_info) + + file_infos = set() + for root, _dirs, files in self.walk(path_info): + + if DvcIgnore.DVCIGNORE_FILE in files: + raise DvcIgnoreInCollectedDirError(root) + + file_infos.update(path_info / root / fname for fname in files) + + checksums = {fi: self.state.get(fi) for fi in file_infos} + not_in_state = { + fi for fi, checksum in checksums.items() if checksum is None + } + + new_checksums = self._calculate_checksums(not_in_state) + + checksums.update(new_checksums) + + result = [ + { + self.PARAM_CHECKSUM: checksums[fi], + # 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. + # + # Yes, this is a BUG, as long as we permit "/" in + # filenames on Windows and "\" on Unix + self.PARAM_RELPATH: fi.relative_to(path_info).as_posix(), + } + for fi in file_infos + ] # Sorting the list by path to ensure reproducibility - return sorted(dir_info.values(), key=itemgetter(self.PARAM_RELPATH)) + return sorted(result, key=itemgetter(self.PARAM_RELPATH)) def get_dir_checksum(self, path_info): dir_info = self._collect_dir(path_info) diff --git a/dvc/utils/__init__.py b/dvc/utils/__init__.py index f91610211..ed04b8ace 100644 --- a/dvc/utils/__init__.py +++ b/dvc/utils/__init__.py @@ -2,8 +2,6 @@ from __future__ import unicode_literals -from concurrent.futures import Future - from dvc.utils.compat import ( str, builtin_str, @@ -394,9 +392,3 @@ def relpath(path, start=os.curdir): ): return path return os.path.relpath(path, start) - - -def as_future(anything): - future = Future() - future.set_result(anything) - return future