dvc: introduce --jobs option and use map_progress

It replaces POOL_SIZE constant and by default equals to the
number of cpus available.

map_progress is a helper for 'sync' and 'import' commands to
run in a multi-threaded mode.

Signed-off-by: Ruslan Kuprieiev <[email protected]>

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:

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__':

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):
         """

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()