upload_dir_contents(): upload multiple files in parallel

py/utils/gs_utils.py

Index: py/utils/gs_utils.py
diff --git a/py/utils/gs_utils.py b/py/utils/gs_utils.py
index f53b4d1427bf9078463547fdf7d5751ef437a24d..8ab12c513271fe29ca6cb4030d52f178e964bdd2 100644
--- a/py/utils/gs_utils.py
+++ b/py/utils/gs_utils.py
@@ -19,10 +19,13 @@ API/library references:
 # System-level imports
 import errno
 import hashlib
+import logging
 import os
 import posixpath
 import re
 import sys
+from Queue import Queue
+from threading import Thread
 
 # Imports from third-party code
 TRUNK_DIRECTORY = os.path.abspath(os.path.join(
@@ -116,21 +119,24 @@ class GSUtils(object):
     IF_MODIFIED = 3 # if there is an existing file with the same name and
                     # contents, leave it alone
 
-  def __init__(self, boto_file_path=None):
+  def __init__(self, boto_file_path=None, logger=None):
     """Constructor.
 
     Params:
       boto_file_path: full path (local-OS-style) on local disk where .boto
           credentials file can be found.  If None, then the GSUtils object
           created will be able to access only public files in Google Storage.
+      logger: a logging.Logger instance to use for logging output; if None,
+          one will be created with default characteristics
 
     Raises an exception if no file is found at boto_file_path, or if the file
     found there is malformed.
     """
+    self.logger = logger or logging.getLogger(__name__)
     self._gs_access_key_id = None
     self._gs_secret_access_key = None
     if boto_file_path:
-      print 'Reading boto file from %s' % boto_file_path
+      self.logger.info('Reading boto file from %s' % boto_file_path)
       boto_dict = _config_file_as_dict(filepath=boto_file_path)
       self._gs_access_key_id = boto_dict['gs_access_key_id']
       self._gs_secret_access_key = boto_dict['gs_secret_access_key']
@@ -215,8 +221,8 @@ class GSUtils(object):
     if upload_if == self.UploadIf.IF_NEW:
       old_key = b.get_key(key_name=dest_path)
       if old_key:
-        print 'Skipping upload of existing file gs://%s/%s' % (
-            b.name, dest_path)
+        self.logger.info('Skipping upload of existing file gs://%s/%s' % (
+            b.name, dest_path))
         return
     elif upload_if == self.UploadIf.IF_MODIFIED:
       old_key = b.get_key(key_name=dest_path)
@@ -224,8 +230,9 @@ class GSUtils(object):
         if not local_md5:
           local_md5 = _get_local_md5(path=source_path)
         if ('"%s"' % local_md5) == old_key.etag:
-          print 'Skipping upload of unmodified file gs://%s/%s : %s' % (
-              b.name, dest_path, local_md5)
+          self.logger.info(
+              'Skipping upload of unmodified file gs://%s/%s : %s' % (
+                  b.name, dest_path, local_md5))
           return
     elif upload_if != self.UploadIf.ALWAYS:
       raise Exception('unknown value of upload_if: %s' % upload_if)
@@ -280,7 +287,7 @@ class GSUtils(object):
           id_type=id_type, id_value=id_value, permission=permission)
 
   def upload_dir_contents(self, source_dir, dest_bucket, dest_dir,
-                          upload_if=UploadIf.ALWAYS, **kwargs):
+                          num_threads=10, upload_if=UploadIf.ALWAYS, **kwargs):
     """Recursively upload contents of a local directory to Google Storage.
 
     params:
@@ -289,6 +296,7 @@ class GSUtils(object):
       dest_bucket: GS bucket to copy the files into
       dest_dir: full path (Posix-style) within that bucket; write the files into
           this directory.  If None, write into the root directory of the bucket.
+      num_threads: how many files to upload at once
       upload_if: one of the UploadIf values, describing in which cases we should
           upload the file
       kwargs: any additional keyword arguments "inherited" from upload_file()
@@ -310,6 +318,7 @@ class GSUtils(object):
       relative_dirpath = dirpath[prefix_length:]
       for filename in filenames:
         source_fileset.add(os.path.join(relative_dirpath, filename))
+    num_files_total = len(source_fileset)
 
     # If we are only uploading files conditionally, remove any unnecessary
     # files from source_fileset.
@@ -342,14 +351,49 @@ class GSUtils(object):
       else:
         raise Exception('unknown value of upload_if: %s' % upload_if)
 
+    # Set up a counter of files that have been uploaded.
+    # Python's Global Interpreter Lock should make this thread-safe; see
+    # http://www.gossamer-threads.com/lists/python/dev/273403
+    num_files_to_upload = len(source_fileset)
+    atomic_incrementer = iter(xrange(1, num_files_to_upload+1)).next
+
     # Upload any files still in source_fileset.
-    for rel_path in sorted(source_fileset):
-      self.upload_file(
-          source_path=os.path.join(source_dir, rel_path),
-          dest_bucket=b,
-          dest_path=posixpath.join(dest_dir, rel_path),
-          upload_if=self.UploadIf.ALWAYS,
-          **kwargs)
+    self.logger.info('Uploading %d files, skipping %d ...' % (
+        num_files_to_upload, num_files_total - num_files_to_upload))
+    if num_files_to_upload == 0:
+      return
+    if num_threads > num_files_to_upload:
+      num_threads = num_files_to_upload

[borenet, 2014/07/28 13:32:11]

Maybe line 358 should be moved here so that it doesn't have to run when
num_files_to_upload == 0?

[epoger, 2014/07/28 14:14:13]

Done.  I don't think there is significant performance impact either way, because
all that line does is set up a "pointer" to the iter().next function call.  I'm
happy with the code either way.

+    q = Queue(maxsize=2*num_threads)
+    end_of_queue = object()
+    def worker():
+      while True:
+        rel_path = q.get()
+        if rel_path is end_of_queue:
+          q.task_done()
+          return

[borenet, 2014/07/28 13:32:11]

Why not use q.get(False) or q.get_nowait() and try/except Queue.Empty?  That
avoids having to use a marker object.

[epoger, 2014/07/28 14:14:14]

The problem with checking for an empty queue is that we don't know WHY the queue
is empty: maybe the main thread finished writing tasks into it, or maybe it
hasn't had a chance to start writing tasks into it yet?  Or maybe it is in the
middle of writing tasks into the queue, and paused for a moment?

AFAIK the only ways to handle this without an end marker are:
1. Allow the worker threads to live forever, blocking on the queue, waiting for
tasks that never come.  Seems wasteful.
2. Make the main thread write all of its tasks into the queue before spinning up
the worker threads.  In this particular case that would probably be fine, but in
general it's nice to be able to start the workers while the main thread is still
filling the task queue.

Would you prefer one of those approaches here?  Or maybe there's another option
I haven't thought of?

[borenet, 2014/07/28 14:20:07]

I was thinking of #2, since you're dealing with the special case of never
needing to add things to the queue after the initial set.  My unfounded
intuition is that the performance difference will be minimal, and the code gets
a little simpler.  Really, I'm fine either way.

+        self.logger.info(' Uploading file %d/%d: %s' % (
+            atomic_incrementer(), num_files_to_upload, rel_path))
+        self.upload_file(
+            source_path=os.path.join(source_dir, rel_path),
+            dest_bucket=b,
+            dest_path=posixpath.join(dest_dir, rel_path),
+            upload_if=self.UploadIf.ALWAYS,
+            **kwargs)
+        q.task_done()
+    # Spin up the threads.
+    for _ in range(num_threads):
+      t = Thread(target=worker)
+      t.daemon = True
+      t.start()
+    # Start loading files into the work queue.
+    for rel_path in source_fileset:
+      q.put(rel_path)
+    # Tell all workers to go home.
+    for _ in range(num_threads):
+      q.put(end_of_queue)
+    # Block until all files have been uploaded and all workers have gone home.
+    q.join()
 
   def download_file(self, source_bucket, source_path, dest_path,
                     create_subdirs_if_needed=False):