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 e25039c0a..98285b522 100644 --- a/dvc/command/data_sync.py +++ b/dvc/command/data_sync.py @@ -6,8 +6,7 @@ from dvc.exceptions import DvcException from dvc.runtime import Runtime from dvc.system import System from dvc.data_cloud import DataCloud - -POOL_SIZE = 4 +from dvc.utils import map_progress class DataSyncError(DvcException): @@ -28,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: @@ -44,7 +42,7 @@ class CmdDataSync(CmdBase): else: raise DataSyncError('File "{}" does not exit'.format(target)) - pool.map(cloud.sync, targets) + 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 5c6aa6eac..65292d81e 100644 --- a/dvc/command/import_file.py +++ b/dvc/command/import_file.py @@ -1,8 +1,6 @@ import os -from shutil import copyfile import re import requests -from multiprocessing.pool import ThreadPool from dvc.command.base import CmdBase, DvcLock from dvc.data_cloud import sizeof_fmt @@ -11,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.progress import progress +from dvc.utils import copyfile, map_progress class ImportFileError(DvcException): def __init__(self, msg): @@ -178,21 +175,18 @@ class CmdImportFile(CmdBase): Logger.debug('Downloaded {}'.format(sizeof_fmt(downloaded))) # update progress bar - self.progress.update_target(name, downloaded, total_length) + progress.update_target(name, downloaded, total_length) f.write(chunk) # tell progress bar that this target is finished downloading - self.progress.finish_target(name) + progress.finish_target(name) def download_targets(self, targets): """ Download targets in a number of threads. """ - self.progress = Progress(len(targets)) - p = ThreadPool(processes=POOL_SIZE) - p.map(self.download_target, targets) - self.progress.finish() + map_progress(self.download_target, targets, self.parsed_args.jobs) def create_state_files(self, targets, lock): """ diff --git a/dvc/data_cloud.py b/dvc/data_cloud.py index a14221bf6..c636e967f 100644 --- a/dvc/data_cloud.py +++ b/dvc/data_cloud.py @@ -3,7 +3,6 @@ import hashlib import os import threading import configparser -import shutil from boto.s3.connection import S3Connection from google.cloud import storage as gc @@ -12,7 +11,8 @@ from dvc.logger import Logger from dvc.exceptions import DvcException from dvc.config import ConfigError from dvc.settings import Settings - +from dvc.progress import progress +from dvc.utils import copyfile class DataCloudError(DvcException): def __init__(self, msg): @@ -27,8 +27,16 @@ def sizeof_fmt(num, suffix='B'): return "%.1f%s%s" % (num, 'Y', suffix) -def percent_cb(complete, total): - Logger.debug('{} transferred out of {}'.format(sizeof_fmt(complete), sizeof_fmt(total))) +def percent_cb(name, complete, total): + Logger.debug('{}: {} transferred out of {}'.format( + name, + sizeof_fmt(complete), + sizeof_fmt(total))) + progress.update_target(os.path.basename(name), complete, total) + + +def create_cb(name): + return (lambda cur,tot: percent_cb(name, cur, tot)) def file_md5(fname): @@ -123,11 +131,11 @@ class DataCloudBase(object): class DataCloudLOCAL(DataCloudBase): def sync_to_cloud(self, item): Logger.debug('sync to cloud ' + item.resolved_cache.dvc + " " + self.storage_path) - shutil.copy(item.resolved_cache.dvc, self.storage_path) + copyfile(item.resolved_cache.dvc, self.storage_path) def sync_from_cloud(self, item): Logger.debug('sync from cloud ' + self.storage_path + " " + item.resolved_cache.dvc) - shutil.copy(self.storage_path, item.resolved_cache.dvc) + copyfile(self.storage_path, item.resolved_cache.dvc) def remove_from_cloud(self, item): Logger.debug('rm from cloud ' + item.resolved_cache.dvc) @@ -237,7 +245,8 @@ class DataCloudAWS(DataCloudBase): raise DataCloudError('File "{}" does not exist in the cloud'.format(key_name)) Logger.info('Downloading cache file from S3 "{}/{}"'.format(bucket.name, key_name)) - key.get_contents_to_filename(item.resolved_cache.relative, cb=percent_cb) + key.get_contents_to_filename(item.resolved_cache.relative, + cb=create_cb(item.resolved_cache.relative)) Logger.info('Downloading completed') def sync_to_cloud(self, data_item): @@ -259,7 +268,8 @@ class DataCloudAWS(DataCloudBase): Logger.info('Uploading cache file "{}" to S3 "{}"'.format(data_item.resolved_cache.relative, aws_key)) key = bucket.new_key(aws_key) - key.set_contents_from_filename(data_item.resolved_cache.relative, cb=percent_cb) + key.set_contents_from_filename(data_item.resolved_cache.relative, + cb=create_cb(data_item.resolved_cache.relative)) Logger.info('Uploading completed') def remove_from_cloud(self, data_item): diff --git a/dvc/progress.py b/dvc/progress.py index 31c45b682..f4026db3a 100644 --- a/dvc/progress.py +++ b/dvc/progress.py @@ -6,11 +6,14 @@ class Progress(object): """ Simple multi-target progress bar. """ - def __init__(self, total): - self._n_total = total + def __init__(self): + self._n_total = 0 self._n_finished = 0 self._lock = threading.Lock() + def set_n_total(self, total): + self._n_total = total + def _clearln(self): print('\r\x1b[K', end='') @@ -73,3 +76,5 @@ class Progress(object): name = target_name[:name_len] if len(target_name) > name_len else target_name return num + bar + percent + name + +progress = Progress() diff --git a/dvc/utils.py b/dvc/utils.py index 8f73b897c..991634c48 100644 --- a/dvc/utils.py +++ b/dvc/utils.py @@ -2,6 +2,12 @@ import os import stat import shutil +from multiprocessing.pool import ThreadPool + +from dvc.progress import progress + +LOCAL_CHUNK_SIZE = 1024*1024*1024 + def cached_property(f): def get(self): try: @@ -24,3 +30,39 @@ def rmtree(dir): os.chmod(dir, stat.S_IWUSR) shutil.rmtree(dir, ignore_errors=True) + + +def copyfile(src, dest): + copied = 0 + name = os.path.basename(src) + total = os.stat(src).st_size + + f = open(src, 'rb') + + if os.path.isdir(dest): + o = open(dest + '/' + name, 'wb+') + else: + o = open(dest, 'wb+') + + while True: + buf = f.read(LOCAL_CHUNK_SIZE) + if not buf: + break + o.write(buf) + copied += len(buf) + progress.update_target(name, copied, total) + + progress.finish_target(name) + + 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()