upload_dir_contents(): upload multiple files in parallel

py/utils/gs_utils.py

 #!/usr/bin/python
 
 # pylint: disable=C0301
 """
 Copyright 2014 Google Inc.
 
 Use of this source code is governed by a BSD-style license that can be
 found in the LICENSE file.
 
 Utilities for accessing Google Cloud Storage, using the boto library (wrapper
 for the XML API).
 
 API/library references:
 - https://developers.google.com/storage/docs/reference-guide
 - http://googlecloudstorage.blogspot.com/2012/09/google-cloud-storage-tutorial-using-boto.html
 """
 # pylint: enable=C0301
 
 # 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(
     os.path.dirname(__file__), os.pardir, os.pardir))
 for import_subdir in ['boto']:
   import_dirpath = os.path.join(
       TRUNK_DIRECTORY, 'third_party', 'externals', import_subdir)
   if import_dirpath not in sys.path:
     # We need to insert at the beginning of the path, to make sure that our
     # imported versions are favored over others that might be in the path.
@@ skipping 73 matching lines @@
     take longer than just uploading the file.
     See http://skbug.com/2778 ('gs_utils: when uploading IF_NEW, batch up
     checks for existing files within a single remote directory')
     """
     ALWAYS = 1      # always upload the file
     IF_NEW = 2      # if there is an existing file with the same name,
                     # leave it alone
     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']
     # Which field we get/set in ACL entries, depending on IdType.
     self._field_by_id_type = {
         self.IdType.GROUP_BY_DOMAIN: 'domain',
         self.IdType.GROUP_BY_EMAIL:  'email_address',
         self.IdType.GROUP_BY_ID:     'id',
         self.IdType.USER_BY_EMAIL:   'email_address',
         self.IdType.USER_BY_ID:      'id',
@@ skipping 64 matching lines @@
     so that HTTP downloads of the file would be unzipped automatically.
     See https://developers.google.com/storage/docs/gsutil/addlhelp/
         WorkingWithObjectMetadata#content-encoding
     """
     b = self._connect_to_bucket(bucket=dest_bucket)
     local_md5 = None  # filled in lazily
 
     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)
       if old_key:
         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)
 
     # Upload the file using a temporary name at first, in case the transfer
     # is interrupted partway through.
     if not local_md5:
       local_md5 = _get_local_md5(path=source_path)
     initial_key = Key(b)
     initial_key.name = dest_path + '-uploading-' + local_md5
@@ skipping 34 matching lines @@
 
     # Set ACLs on the file.
     # We do this *after* copy_key(), because copy_key's preserve_acl
     # functionality would incur a performance hit.
     for (id_type, id_value, permission) in fine_grained_acl_list or []:
       self.set_acl(
           bucket=b, path=final_key.name,
           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:
       source_dir: full path (local-OS-style) on local disk of directory to copy
           contents of
       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()
 
     The copy operates as a merge: any files in source_dir will be "overlaid" on
     top of the existing content in dest_dir.  Existing files with the same names
     may or may not be overwritten, depending on the value of upload_if.
 
     TODO(epoger): Upload multiple files simultaneously to reduce latency.
     """
     b = self._connect_to_bucket(bucket=dest_bucket)
     if not dest_dir:
       dest_dir = ''
 
     # Create a set of all files within source_dir.
     source_fileset = set()
     prefix_length = len(source_dir)+1
     for dirpath, _, filenames in os.walk(source_dir):
       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.
     if upload_if == self.UploadIf.ALWAYS:
       pass  # there are no shortcuts... upload them all
     else:
       # Create a mapping of filename to Key for existing files within dest_dir
       existing_dest_filemap = {}
       prefix = dest_dir
       if prefix and not prefix.endswith('/'):
@@ skipping 12 matching lines @@
       elif upload_if == self.UploadIf.IF_MODIFIED:
         for rel_path in files_in_common:
           local_md5 = '"%s"' % _get_local_md5(path=os.path.join(
               source_dir, rel_path))
           key = existing_dest_filemap[rel_path]
           if local_md5 == key.etag:
             source_fileset.remove(rel_path)
       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):
     """Downloads a single file from Google Cloud Storage to local disk.
 
     Args:
       source_bucket: GS bucket to download the file from
       source_path: full path (Posix-style) within that bucket
       dest_path: full path (local-OS-style) on local disk to copy the file to
       create_subdirs_if_needed: boolean; whether to create subdirectories as
@@ skipping 240 matching lines @@
 
 def _get_local_md5(path):
   """Returns the MD5 hash of a file on local disk."""
   hasher = hashlib.md5()
   with open(path, 'rb') as f:
     while True:
       data = f.read(64*1024)
       if not data:
         return hasher.hexdigest()
       hasher.update(data)