diff --git a/dvc/cli.py b/dvc/cli.py index 60153446c..a1be92d47 100644 --- a/dvc/cli.py +++ b/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) diff --git a/dvc/cloud/data_cloud.py b/dvc/cloud/data_cloud.py index 2e01b7726..5298855ea 100644 --- a/dvc/cloud/data_cloud.py +++ b/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) diff --git a/dvc/command/data_sync.py b/dvc/command/data_sync.py index ad9b91d1f..caa33c4ef 100644 --- a/dvc/command/data_sync.py +++ b/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) diff --git a/dvc/project.py b/dvc/project.py index 1e2b8ecfe..8050cc30d 100644 --- a/dvc/project.py +++ b/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() diff --git a/dvc/utils.py b/dvc/utils.py index 24e540443..55e4c037a 100644 --- a/dvc/utils.py +++ b/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: diff --git a/tests/test_utils.py b/tests/test_utils.py index e377a941d..b1344d89c 100644 --- a/tests/test_utils.py +++ b/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)