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 Queue
 import re
 import sys
+import threading
 
 # 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.
     sys.path.insert(0, import_dirpath)
 from boto.exception import BotoServerError
 from boto.gs import acl
 from boto.gs.bucket import Bucket
 from boto.gs.connection import GSConnection
 from boto.gs.key import Key
 from boto.s3.bucketlistresultset import BucketListResultSet
 from boto.s3.connection import SubdomainCallingFormat
 from boto.s3.prefix import Prefix
 
+# How many files to upload at once, by default.
+# TODO(epoger): Is there a way to compute this intelligently?  To some extent
+# it is a function of how many cores are on the machine, and how many other
+# processes it is running; but it's probably more a function of how much time
+# each core sits idle waiting for network I/O to complete.
+DEFAULT_UPLOAD_THREADS = 10
+
 
 class AnonymousGSConnection(GSConnection):
   """GSConnection class that allows anonymous connections.
 
   The GSConnection class constructor in
   https://github.com/boto/boto/blob/develop/boto/gs/connection.py doesn't allow
   for anonymous connections (connections without credentials), so we have to
   override it.
   """
   def __init__(self):
@@ skipping 53 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,
+                          num_threads=DEFAULT_UPLOAD_THREADS,
                           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 13 matching lines @@
         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)
 
     # 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)
+    num_files_to_upload = len(source_fileset)
+    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
+
+    # Create a work queue with all files that need to be uploaded.
+    q = Queue.Queue(maxsize=num_files_to_upload)
+    for rel_path in source_fileset:
+      q.put(rel_path)
+
+    # Spin up worker threads to read from the task queue.
+    def worker():
+      while True:
+        try:
+          rel_path = q.get(block=False)
+        except Queue.Empty:
+          return  # no more tasks in the queue, so exit
+        self.logger.info(' Uploading file %d/%d: %s' % (
+            num_files_to_upload - q.qsize(), 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()
+    for _ in range(num_threads):
+      t = threading.Thread(target=worker)
+      t.daemon = True
+      t.start()
+
+    # Block until all files have been uploaded and all workers have exited.
+    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)