Revert "cloud: process targets one-by-one"

This reverts commit 61b469144e.

dvc/cli.py

@@ -123,6 +123,15 @@ def parse_args(argv=None):
                         help='Command or command file to execute')
     run_parser.set_defaults(func=CmdRun)
 
+    # Parent parser used in pull/push/status
+    parent_cache_parser = argparse.ArgumentParser(
+                        add_help=False,
+                        parents=[parent_parser])
+    parent_cache_parser.add_argument('-j',
+                        '--jobs',
+                        type=int,
+                        default=cpu_count(),
+                        help='Number of jobs to run simultaneously.')
     # Cache
     cache_parser = subparsers.add_parser(
                         'cache',
@@ -135,21 +144,21 @@ def parse_args(argv=None):
     # Pull
     pull_parser = cache_subparsers.add_parser(
                         'pull',
-                        parents=[parent_parser],
+                        parents=[parent_cache_parser],
                         help='Pull data files from the cloud')
     pull_parser.set_defaults(func=CmdDataPull)
 
     # Push
     push_parser = cache_subparsers.add_parser(
                         'push',
-                        parents=[parent_parser],
+                        parents=[parent_cache_parser],
                         help='Push data files to the cloud')
     push_parser.set_defaults(func=CmdDataPush)
 
     # Status
     status_parser = cache_subparsers.add_parser(
                         'status',
-                        parents=[parent_parser],
+                        parents=[parent_cache_parser],
                         help='Show mismatches between local cache and cloud cache')
     status_parser.set_defaults(func=CmdDataStatus)
 

dvc/cloud/data_cloud.py

@@ -2,7 +2,7 @@ from dvc.cloud.instance_manager import CloudSettings
 from dvc.logger import Logger
 from dvc.exceptions import DvcException
 from dvc.config import ConfigError
-from dvc.progress import progress
+from dvc.utils import map_progress
 
 from dvc.cloud.aws import DataCloudAWS
 from dvc.cloud.gcp import DataCloudGCP
@@ -79,36 +79,26 @@ class DataCloud(object):
 
         self._cloud.sanity_check()
 
-    def _process_targets(self, func, targets):
+    def _map_targets(self, func, targets, jobs):
         """
-        Process targets with a progress bar.
+        Process targets as data items in parallel.
         """
-        ret = []
-        progress.set_n_total(len(targets))
-        try:
-            for target in targets:
-                ret += [func(target)]
-        except Exception as exc:
-            raise
-        finally:
-            progress.finish()
+        return map_progress(func, targets, jobs)
 
-        return list(zip(targets, ret))
-
-    def push(self, targets):
+    def push(self, targets, jobs=1):
         """
         Push data items in a cloud-agnostic way.
         """
-        return self._process_targets(self._cloud.push, targets)
+        return self._map_targets(self._cloud.push, targets, jobs)
 
-    def pull(self, targets):
+    def pull(self, targets, jobs=1):
         """
         Pull data items in a cloud-agnostic way.
         """
-        return self._process_targets(self._cloud.pull, targets)
+        return self._map_targets(self._cloud.pull, targets, jobs)
 
-    def status(self, targets):
+    def status(self, targets, jobs=1):
         """
         Check status of data items in a cloud-agnostic way.
         """
-        return self._process_targets(self._cloud.status, targets)
+        return self._map_targets(self._cloud.status, targets, jobs)

dvc/command/data_sync.py

@@ -9,7 +9,7 @@ from dvc.exceptions import DvcException
 class CmdDataPull(CmdBase):
     def run(self):
         try:
-            self.project.pull()
+            self.project.pull(self.args.jobs)
         except DvcException as exc:
             self.project.logger.error('Failed to pull data from the cloud', exc)
             return 1
@@ -19,7 +19,7 @@ class CmdDataPull(CmdBase):
 class CmdDataPush(CmdBase):
     def run(self):
         try:
-            self.project.push()
+            self.project.push(self.args.jobs)
         except DvcException as exc:
             self.project.logger.error('Failed to push data to the cloud', exc)
             return 1
@@ -46,7 +46,7 @@ class CmdDataStatus(CmdBase):
 
     def run(self):
         try:
-            status = self.project.status()
+            status = self.project.status(self.args.jobs)
             self._show(status)
         except DvcException as exc:
             self.project.logger.error('Failed to obtain data status', exc)

dvc/project.py

@@ -195,15 +195,15 @@ class Project(object):
             self._remove_cache(cache)
             self.logger.info(u'\'{}\' was removed'.format(self.to_dvc_path(cache)))
 
-    def push(self):
-        self.cloud.push(self._used_cache())
+    def push(self, jobs=1):
+        self.cloud.push(self._used_cache(), jobs)
 
-    def pull(self):
-        self.cloud.pull(self._used_cache())
+    def pull(self, jobs=1):
+        self.cloud.pull(self._used_cache(), jobs)
         self.checkout()
 
-    def status(self):
-        return self.cloud.status(self._used_cache())
+    def status(self, jobs=1):
+        return self.cloud.status(self._used_cache(), jobs)
 
     def graph(self):
         G = nx.DiGraph()

dvc/utils.py

@@ -72,6 +72,24 @@ def copyfile(src, dest):
     fdest.close()
 
 
+def map_progress(func, targets, n_threads):
+    """
+    Process targets in multi-threaded mode with progress bar
+    """
+    progress.set_n_total(len(targets))
+    pool = ThreadPool(processes=n_threads)
+    ret = []
+
+    try:
+        ret = pool.map(func, targets)
+    except Exception as exc:
+        raise
+    finally:
+        progress.finish()
+
+    return list(zip(targets, ret))
+
+
 def parse_target_metric_file(file_name):
     with open(file_name, 'r') as fd:
         try:

tests/test_utils.py

@@ -26,3 +26,14 @@ class TestUtils(TestDvc):
         shutil.rmtree(dest_dir)
         os.remove(src)
         os.remove(dest)
+
+    def test_map_progress(self):
+        def f(target):
+            with open(target, 'w+') as o:
+                o.write(target)
+
+        targets = ['map{}'.format(i) for i in range(1, 10)]
+        n_threads = [1, 10, 20]
+
+        for n in n_threads:
+            utils.map_progress(f, targets, n)