Merge pull request #99 from efiop/master

Progress bar improvements, --jobs and progress bar for LOCAL cloud and AWS

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

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

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

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

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