diff --git a/dvc/command/base.py b/dvc/command/base.py index 167bae691..51a8e8ce3 100644 --- a/dvc/command/base.py +++ b/dvc/command/base.py @@ -1,6 +1,8 @@ import argparse import fasteners +from multiprocessing import cpu_count + from dvc.config import ConfigError from dvc.exceptions import DvcException from dvc.logger import Logger @@ -89,6 +91,8 @@ class CmdBase(object): def define_common_args(self, parser): parser.add_argument('--quiet', '-q', action='store_true', default=False, help='Be quiet.') parser.add_argument('--verbose', '-v', action='store_true', default=False, help='Be verbose.') + parser.add_argument('--jobs', '-j', type=int, default=cpu_count(), + help='Number of jobs to run simultaneously.') def process_common_args(self): if self._parsed_args.quiet and not self._parsed_args.verbose: diff --git a/dvc/command/data_sync.py b/dvc/command/data_sync.py index 88006ac67..98285b522 100644 --- a/dvc/command/data_sync.py +++ b/dvc/command/data_sync.py @@ -6,9 +6,7 @@ from dvc.exceptions import DvcException from dvc.runtime import Runtime from dvc.system import System from dvc.data_cloud import DataCloud -from dvc.progress import progress - -POOL_SIZE = 4 +from dvc.utils import map_progress class DataSyncError(DvcException): @@ -29,7 +27,6 @@ class CmdDataSync(CmdBase): def run(self): with DvcLock(self.is_locker, self.git): cloud = DataCloud(self.settings) - pool = ThreadPool(processes=POOL_SIZE) targets = [] if len(self.parsed_args.targets) == 0: @@ -45,9 +42,7 @@ class CmdDataSync(CmdBase): else: raise DataSyncError('File "{}" does not exit'.format(target)) - progress.set_n_total(len(targets)) - pool.map(cloud.sync, targets) - progress.finish() + map_progress(cloud.sync, targets, self.parsed_args.jobs) pass if __name__ == '__main__': diff --git a/dvc/command/import_file.py b/dvc/command/import_file.py index afe9e5fa8..65292d81e 100644 --- a/dvc/command/import_file.py +++ b/dvc/command/import_file.py @@ -1,7 +1,6 @@ import os import re import requests -from multiprocessing.pool import ThreadPool from dvc.command.base import CmdBase, DvcLock from dvc.data_cloud import sizeof_fmt @@ -10,9 +9,8 @@ from dvc.exceptions import DvcException from dvc.runtime import Runtime from dvc.state_file import StateFile from dvc.system import System -from dvc.command.data_sync import POOL_SIZE from dvc.progress import progress -from dvc.utils import copyfile +from dvc.utils import copyfile, map_progress class ImportFileError(DvcException): def __init__(self, msg): @@ -188,10 +186,7 @@ class CmdImportFile(CmdBase): """ Download targets in a number of threads. """ - progress.set_n_total(len(targets)) - p = ThreadPool(processes=POOL_SIZE) - p.map(self.download_target, targets) - progress.finish() + map_progress(self.download_target, targets, self.parsed_args.jobs) def create_state_files(self, targets, lock): """ diff --git a/dvc/utils.py b/dvc/utils.py index f7e4dc805..991634c48 100644 --- a/dvc/utils.py +++ b/dvc/utils.py @@ -2,6 +2,8 @@ import os import stat import shutil +from multiprocessing.pool import ThreadPool + from dvc.progress import progress LOCAL_CHUNK_SIZE = 1024*1024*1024 @@ -54,3 +56,13 @@ def copyfile(src, dest): f.close() o.close() + + +def map_progress(func, targets, n_threads): + """ + Process targets in multi-threaded mode with progress bar + """ + progress.set_n_total(len(targets)) + p = ThreadPool(processes=n_threads) + p.map(func, targets) + progress.finish()