From 11d1fb502a240ba7d9c1ca3f8e1bc6466e73ffe8 Mon Sep 17 00:00:00 2001 From: Ruslan Kuprieiev Date: Wed, 21 Jun 2017 03:37:42 +0300 Subject: [PATCH 1/4] progress: use global variable 'progress' There is no point in having more than just one global progress bar. Signed-off-by: Ruslan Kuprieiev --- dvc/command/import_file.py | 10 +++++----- dvc/progress.py | 9 +++++++-- 2 files changed, 12 insertions(+), 7 deletions(-) diff --git a/dvc/command/import_file.py b/dvc/command/import_file.py index 5c6aa6eac..6e80bfe4d 100644 --- a/dvc/command/import_file.py +++ b/dvc/command/import_file.py @@ -12,7 +12,7 @@ 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 class ImportFileError(DvcException): @@ -178,21 +178,21 @@ 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)) + progress.set_n_total(len(targets)) p = ThreadPool(processes=POOL_SIZE) p.map(self.download_target, targets) - self.progress.finish() + progress.finish() def create_state_files(self, targets, lock): """ 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() From bdb30cc602478d8673f93b5cc0574e66ac5f9b5b Mon Sep 17 00:00:00 2001 From: Ruslan Kuprieiev Date: Wed, 21 Jun 2017 03:51:06 +0300 Subject: [PATCH 2/4] sync: add progress bar for AWS Signed-off-by: Ruslan Kuprieiev --- dvc/command/data_sync.py | 3 +++ dvc/data_cloud.py | 20 +++++++++++++++----- 2 files changed, 18 insertions(+), 5 deletions(-) diff --git a/dvc/command/data_sync.py b/dvc/command/data_sync.py index e25039c0a..88006ac67 100644 --- a/dvc/command/data_sync.py +++ b/dvc/command/data_sync.py @@ -6,6 +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 @@ -44,7 +45,9 @@ 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() pass if __name__ == '__main__': diff --git a/dvc/data_cloud.py b/dvc/data_cloud.py index a14221bf6..c19e9c55f 100644 --- a/dvc/data_cloud.py +++ b/dvc/data_cloud.py @@ -12,7 +12,7 @@ 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 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: return percent_cb(name, cur, tot)) def file_md5(fname): @@ -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): From 7c6765c1960b04542a2b3ebd21cb4ba73d05bae9 Mon Sep 17 00:00:00 2001 From: Ruslan Kuprieiev Date: Sun, 25 Jun 2017 01:30:06 +0300 Subject: [PATCH 3/4] import/sync: add progress bar for local copy Currently progress bar is only supported for downloading from web(aws,http* etc), but now it is also supported for local files(i.e. LOCAL for 'data sync' and just local files for 'import') Signed-off-by: Ruslan Kuprieiev --- dvc/command/import_file.py | 3 +-- dvc/data_cloud.py | 8 ++++---- dvc/utils.py | 30 ++++++++++++++++++++++++++++++ 3 files changed, 35 insertions(+), 6 deletions(-) diff --git a/dvc/command/import_file.py b/dvc/command/import_file.py index 6e80bfe4d..afe9e5fa8 100644 --- a/dvc/command/import_file.py +++ b/dvc/command/import_file.py @@ -1,5 +1,4 @@ import os -from shutil import copyfile import re import requests from multiprocessing.pool import ThreadPool @@ -13,7 +12,7 @@ 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 class ImportFileError(DvcException): def __init__(self, msg): diff --git a/dvc/data_cloud.py b/dvc/data_cloud.py index c19e9c55f..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 @@ -13,6 +12,7 @@ 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): @@ -36,7 +36,7 @@ def percent_cb(name, complete, total): def create_cb(name): - return (lambda cur,tot: return percent_cb(name, cur, tot)) + return (lambda cur,tot: percent_cb(name, cur, tot)) def file_md5(fname): @@ -131,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) diff --git a/dvc/utils.py b/dvc/utils.py index 8f73b897c..f7e4dc805 100644 --- a/dvc/utils.py +++ b/dvc/utils.py @@ -2,6 +2,10 @@ import os import stat import shutil +from dvc.progress import progress + +LOCAL_CHUNK_SIZE = 1024*1024*1024 + def cached_property(f): def get(self): try: @@ -24,3 +28,29 @@ 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() From 76db9db5c2403c73e9184c5bad0f74ee92fe6ce3 Mon Sep 17 00:00:00 2001 From: Ruslan Kuprieiev Date: Sun, 25 Jun 2017 04:02:07 +0300 Subject: [PATCH 4/4] 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 --- dvc/command/base.py | 4 ++++ dvc/command/data_sync.py | 9 ++------- dvc/command/import_file.py | 9 ++------- dvc/utils.py | 12 ++++++++++++ 4 files changed, 20 insertions(+), 14 deletions(-) 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()