Repository: DataBiosphere/toil

Issue:
Unable to import files
File import does not give any feedback.  It seems to be adding a file to the S3 bucket but gets stuck.  I hit ^C after about 15 minutes of no progress and no network activity.  It appears to waiting for a response:

```
$ cwltoil --logDebug --batchSystem=mesos --mesosMaster=mesos-master:5050 --jobStore=aws:us-west-2:cwltoil-test-pa2 main-bwa_freebayes_HG01977.cwl main-bwa_freebayes_HG01977-samples.json 
2016-05-11 14:17:31,458 INFO:toil.lib.bioio: Logging set at level: DEBUG
2016-05-11 14:18:04,334 INFO:toil.lib.bioio: Logging set at level: DEBUG
2016-05-11 14:18:04,398 DEBUG:toil.jobStores.aws.jobStore: Instantiating <class 'toil.jobStores.aws.jobStore.AWSJobStore'> for region us-west-2 and name prefix 'cwltoil-test-pa2'
2016-05-11 14:18:08,300 INFO:toil.jobStores.abstractJobStore: The workflow ID is: '75205970-692b-4e0e-b3b7-3f457c55e77b'
2016-05-11 14:18:08,648 DEBUG:toil.jobStores.aws.jobStore: Wrote FileInfo(fileID='a77581f6-ad82-5142-acb7-5a6b01eab3fb', ownerID='891f7db6-e4d9-4221-a58e-ab6cc4395f94', encrypted=False, version='', previousVersion='', content='\x80\x02ctoil.co...ryCountq6K\x00ub.', _numContentChunks=2) bound to <toil.jobStores.aws.jobStore.AWSJobStore object at 0x7f4fb88c3890> for shared file 'config.pickle'.
2016-05-11 14:18:08,734 INFO:toil.jobStores.abstractJobStore: Unable to import 'toil.jobStores.azureJobStore'

^CTraceback (most recent call last):
  File "/usr/local/bin/cwltoil", line 9, in <module>
    load_entry_point('toil==3.2.0a2', 'console_scripts', 'cwltoil')()
  File "/usr/local/lib/python2.7/dist-packages/toil/cwl/cwltoil.py", line 580, in main
    adjustFiles(builder.job, functools.partial(writeFile, toil.importFile, {}))
  File "/usr/local/lib/python2.7/dist-packages/cwltool/process.py", line 128, in adjustFiles
    adjustFiles(rec[d], op)
  File "/usr/local/lib/python2.7/dist-packages/cwltool/process.py", line 131, in adjustFiles
    adjustFiles(d, op)
  File "/usr/local/lib/python2.7/dist-packages/cwltool/process.py", line 131, in adjustFiles
    adjustFiles(d, op)
  File "/usr/local/lib/python2.7/dist-packages/cwltool/process.py", line 126, in adjustFiles
    rec["path"] = op(rec["path"])
  File "/usr/local/lib/python2.7/dist-packages/toil/cwl/cwltoil.py", line 162, in writeFile
    index[x] = (writeFunc(rp), os.path.basename(x))
  File "/usr/local/lib/python2.7/dist-packages/toil/common.py", line 653, in importFile
    return self._jobStore.importFile(srcUrl)
  File "/usr/local/lib/python2.7/dist-packages/toil/jobStores/abstractJobStore.py", line 259, in importFile
    return self._importFile(findJobStoreForUrl(url), url)
  File "/usr/local/lib/python2.7/dist-packages/toil/jobStores/aws/jobStore.py", line 279, in _importFile
    info.copyFrom(srcKey)
  File "/usr/local/lib/python2.7/dist-packages/toil/jobStores/aws/jobStore.py", line 846, in copyFrom
    headers=self._s3EncryptionHeaders()).version_id
  File "/usr/local/lib/python2.7/dist-packages/toil/jobStores/aws/jobStore.py", line 878, in _copyKey
    headers=headers)
  File "/usr/local/lib/python2.7/dist-packages/toil/jobStores/aws/jobStore.py", line 85, in copyKeyMultipart
    headers=headers)
  File "/usr/local/lib/python2.7/dist-packages/boto/s3/multipart.py", line 306, in copy_part_from_key
    query_args=query_args)
  File "/usr/local/lib/python2.7/dist-packages/boto/s3/bucket.py", line 872, in copy_key
    body = response.read()
  File "/usr/local/lib/python2.7/dist-packages/boto/connection.py", line 410, in read
    self._cached_response = http_client.HTTPResponse.read(self)
  File "/usr/lib/python2.7/httplib.py", line 588, in read
    return self._read_chunked(amt)
  File "/usr/lib/python2.7/httplib.py", line 630, in _read_chunked
    line = self.fp.readline(_MAXLINE + 1)
  File "/usr/lib/python2.7/socket.py", line 480, in readline
    data = self._sock.recv(self._rbufsize)
  File "/usr/lib/python2.7/ssl.py", line 734, in recv
    return self.read(buflen)
  File "/usr/lib/python2.7/ssl.py", line 621, in read
    v = self._sslobj.read(len or 1024)
KeyboardInterrupt
```

Unable to import files
File import does not give any feedback.  It seems to be adding a file to the S3 bucket but gets stuck.  I hit ^C after about 15 minutes of no progress and no network activity.  It appears to waiting for a response:

```
$ cwltoil --logDebug --batchSystem=mesos --mesosMaster=mesos-master:5050 --jobStore=aws:us-west-2:cwltoil-test-pa2 main-bwa_freebayes_HG01977.cwl main-bwa_freebayes_HG01977-samples.json 
2016-05-11 14:17:31,458 INFO:toil.lib.bioio: Logging set at level: DEBUG
2016-05-11 14:18:04,334 INFO:toil.lib.bioio: Logging set at level: DEBUG
2016-05-11 14:18:04,398 DEBUG:toil.jobStores.aws.jobStore: Instantiating <class 'toil.jobStores.aws.jobStore.AWSJobStore'> for region us-west-2 and name prefix 'cwltoil-test-pa2'
2016-05-11 14:18:08,300 INFO:toil.jobStores.abstractJobStore: The workflow ID is: '75205970-692b-4e0e-b3b7-3f457c55e77b'
2016-05-11 14:18:08,648 DEBUG:toil.jobStores.aws.jobStore: Wrote FileInfo(fileID='a77581f6-ad82-5142-acb7-5a6b01eab3fb', ownerID='891f7db6-e4d9-4221-a58e-ab6cc4395f94', encrypted=False, version='', previousVersion='', content='\x80\x02ctoil.co...ryCountq6K\x00ub.', _numContentChunks=2) bound to <toil.jobStores.aws.jobStore.AWSJobStore object at 0x7f4fb88c3890> for shared file 'config.pickle'.
2016-05-11 14:18:08,734 INFO:toil.jobStores.abstractJobStore: Unable to import 'toil.jobStores.azureJobStore'

^CTraceback (most recent call last):
  File "/usr/local/bin/cwltoil", line 9, in <module>
    load_entry_point('toil==3.2.0a2', 'console_scripts', 'cwltoil')()
  File "/usr/local/lib/python2.7/dist-packages/toil/cwl/cwltoil.py", line 580, in main
    adjustFiles(builder.job, functools.partial(writeFile, toil.importFile, {}))
  File "/usr/local/lib/python2.7/dist-packages/cwltool/process.py", line 128, in adjustFiles
    adjustFiles(rec[d], op)
  File "/usr/local/lib/python2.7/dist-packages/cwltool/process.py", line 131, in adjustFiles
    adjustFiles(d, op)
  File "/usr/local/lib/python2.7/dist-packages/cwltool/process.py", line 131, in adjustFiles
    adjustFiles(d, op)
  File "/usr/local/lib/python2.7/dist-packages/cwltool/process.py", line 126, in adjustFiles
    rec["path"] = op(rec["path"])
  File "/usr/local/lib/python2.7/dist-packages/toil/cwl/cwltoil.py", line 162, in writeFile
    index[x] = (writeFunc(rp), os.path.basename(x))
  File "/usr/local/lib/python2.7/dist-packages/toil/common.py", line 653, in importFile
    return self._jobStore.importFile(srcUrl)
  File "/usr/local/lib/python2.7/dist-packages/toil/jobStores/abstractJobStore.py", line 259, in importFile
    return self._importFile(findJobStoreForUrl(url), url)
  File "/usr/local/lib/python2.7/dist-packages/toil/jobStores/aws/jobStore.py", line 279, in _importFile
    info.copyFrom(srcKey)
  File "/usr/local/lib/python2.7/dist-packages/toil/jobStores/aws/jobStore.py", line 846, in copyFrom
    headers=self._s3EncryptionHeaders()).version_id
  File "/usr/local/lib/python2.7/dist-packages/toil/jobStores/aws/jobStore.py", line 878, in _copyKey
    headers=headers)
  File "/usr/local/lib/python2.7/dist-packages/toil/jobStores/aws/jobStore.py", line 85, in copyKeyMultipart
    headers=headers)
  File "/usr/local/lib/python2.7/dist-packages/boto/s3/multipart.py", line 306, in copy_part_from_key
    query_args=query_args)
  File "/usr/local/lib/python2.7/dist-packages/boto/s3/bucket.py", line 872, in copy_key
    body = response.read()
  File "/usr/local/lib/python2.7/dist-packages/boto/connection.py", line 410, in read
    self._cached_response = http_client.HTTPResponse.read(self)
  File "/usr/lib/python2.7/httplib.py", line 588, in read
    return self._read_chunked(amt)
  File "/usr/lib/python2.7/httplib.py", line 630, in _read_chunked
    line = self.fp.readline(_MAXLINE + 1)
  File "/usr/lib/python2.7/socket.py", line 480, in readline
    data = self._sock.recv(self._rbufsize)
  File "/usr/lib/python2.7/ssl.py", line 734, in recv
    return self.read(buflen)
  File "/usr/lib/python2.7/ssl.py", line 621, in read
    v = self._sslobj.read(len or 1024)
KeyboardInterrupt
```



Discussion:
Is the destination job store in the same AWS region as the source object's bucket?

How big is the file?

And by "no network activity" you mean very little activity, or literally none? When importing from S3, the copying takes place server-side (on S3) so I would expect there to be little activity, but not no activity at all.

@JakeNarkizian, at the very least we we should add a debug log statement for every part.

@JakeNarkizian we should also consider multi-threading the copy such that we copy several parts at once. This would be reasonably easy to do with Python's [futures backport](https://pypi.python.org/pypi/futures).

Hi @hannes-ucsc, thanks for a quick response.  It is entirely possible data is being transferred correctly and it is just impossible to tell due to lack of logging.  I said "no data" because the network device stats were much too low for a file transfer going through the node, but that is understandable for a server to server copy.  Also, no files appeared in the S3 bucket via the web interface, but that could be due to eventual consistency or need to "commit" the file?

It is about 75 GiB of data.  I don't know what region the original buckets are in, I will have to check.

At minimum the logging (at info level, I think) should say which files are being imported and when each transfer begins & ends.

Is the destination job store in the same AWS region as the source object's bucket?

How big is the file?

And by "no network activity" you mean very little activity, or literally none? When importing from S3, the copying takes place server-side (on S3) so I would expect there to be little activity, but not no activity at all.

@JakeNarkizian, at the very least we we should add a debug log statement for every part.

@JakeNarkizian we should also consider multi-threading the copy such that we copy several parts at once. This would be reasonably easy to do with Python's [futures backport](https://pypi.python.org/pypi/futures).

Hi @hannes-ucsc, thanks for a quick response.  It is entirely possible data is being transferred correctly and it is just impossible to tell due to lack of logging.  I said "no data" because the network device stats were much too low for a file transfer going through the node, but that is understandable for a server to server copy.  Also, no files appeared in the S3 bucket via the web interface, but that could be due to eventual consistency or need to "commit" the file?

It is about 75 GiB of data.  I don't know what region the original buckets are in, I will have to check.

At minimum the logging (at info level, I think) should say which files are being imported and when each transfer begins & ends.


Resolving patch:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -25,7 +25,7 @@
     author_email='benedict@soe.usc.edu',
     url="https://github.com/BD2KGenomics/toil",
     install_requires=[
-        'bd2k-python-lib==1.13.dev19',
+        'bd2k-python-lib==1.13.dev21',
         'dill==0.2.5'],
     extras_require={
         'mesos': [
diff --git a/src/toil/jobStores/aws/jobStore.py b/src/toil/jobStores/aws/jobStore.py
--- a/src/toil/jobStores/aws/jobStore.py
+++ b/src/toil/jobStores/aws/jobStore.py
@@ -17,6 +17,10 @@
 from StringIO import StringIO
 from contextlib import contextmanager, closing
 import logging
+from multiprocessing import cpu_count
+
+import collections
+
 import os
 import re
 import uuid
@@ -37,8 +41,10 @@
 from boto.sdb.item import Item
 import boto.s3
 import boto.sdb
-from boto.exception import SDBResponseError, S3ResponseError, S3CreateError
+from boto.exception import S3CreateError
 from boto.s3.key import Key
+from boto.exception import SDBResponseError, S3ResponseError
+from concurrent.futures import ThreadPoolExecutor
 
 from toil.jobStores.abstractJobStore import (AbstractJobStore,
                                              NoSuchJobException,
@@ -57,55 +63,82 @@
 
 log = logging.getLogger(__name__)
 
-defaultPartSize = 2 ** 20 * 50
-
 
-def copyKeyMultipart(srcKey, dstBucketName, dstKeyName, headers=None):
+def copyKeyMultipart(srcKey, dstBucketName, dstKeyName, partSize, headers=None):
     """
     Copies a key from a source key to a destination key in multiple parts. Note that if the
     destination key exists it will be overwritten implicitly, and if it does not exist a new
-    key will be created.
+    key will be created. If the destination bucket does not exist an error will be raised.
 
     :param boto.s3.key.Key srcKey: The source key to be copied from.
     :param str dstBucketName: The name of the destination bucket for the copy.
     :param str dstKeyName: The name of the destination key that will be created or overwritten.
+    :param int partSize: The size of each individual part, must be >= 5 MiB but large enough to
+           not exceed 10k parts for the whole file
     :param dict headers: Any headers that should be passed.
 
     :rtype: boto.s3.multipart.CompletedMultiPartUpload
     :return: An object representing the completed upload.
     """
-    partSize = defaultPartSize
+    def copyPart(partIndex):
+        if exceptions:
+            return None
+        try:
+            for attempt in retry_s3():
+                with attempt:
+                    start = partIndex * partSize
+                    end = min(start + partSize, totalSize)
+                    part = upload.copy_part_from_key(src_bucket_name=srcKey.bucket.name,
+                                                     src_key_name=srcKey.name,
+                                                     src_version_id=srcKey.version_id,
+                                                     # S3 part numbers are 1-based
+                                                     part_num=partIndex + 1,
+                                                     # S3 range intervals are closed at the end
+                                                     start=start, end=end - 1,
+                                                     headers=headers)
+        except Exception as e:
+            if len(exceptions) < 5:
+                exceptions.append(e)
+                log.error('Failed to copy part number %d:', partIndex, exc_info=True)
+            else:
+                log.warn('Also failed to copy part number %d due to %s.', partIndex, e)
+            return None
+        else:
+            log.debug('Successfully copied part %d of %d.', partIndex, totalParts)
+            # noinspection PyUnboundLocalVariable
+            return part
+
+    totalSize = srcKey.size
+    totalParts = (totalSize + partSize - 1) / partSize
+    exceptions = []
     # We need a location-agnostic connection to S3 so we can't use the one that we
     # normally use for interacting with the job store bucket.
     with closing(boto.connect_s3()) as s3:
-        headers = headers or {}
-        totalSize = srcKey.size
         for attempt in retry_s3():
             with attempt:
                 dstBucket = s3.get_bucket(dstBucketName)
                 upload = dstBucket.initiate_multipart_upload(dstKeyName, headers=headers)
+        log.info("Initiated multipart copy from 's3://%s/%s' to 's3://%s/%s'.",
+                 srcKey.bucket.name, srcKey.name, dstBucketName, dstKeyName)
         try:
-            start = 0
-            partIndex = itertools.count()
-            while start < totalSize:
-                end = min(start + partSize, totalSize)
-                for attempt in retry_s3():
-                    with attempt:
-                        upload.copy_part_from_key(src_bucket_name=srcKey.bucket.name,
-                                                  src_key_name=srcKey.name,
-                                                  src_version_id=srcKey.version_id,
-                                                  part_num=next(partIndex) + 1,
-                                                  start=start,
-                                                  end=end - 1,
-                                                  headers=headers)
-                start += partSize
+            # We can oversubscribe cores by at least a factor of 16 since each copy task just
+            # blocks, waiting on the server. Limit # of threads to 128, since threads aren't
+            # exactly free either. Lastly, we don't need more threads than we have parts.
+            with ThreadPoolExecutor(max_workers=min(cpu_count() * 16, totalParts, 128)) as executor:
+                parts = list(executor.map(copyPart, xrange(0, totalParts)))
+                if exceptions:
+                    raise RuntimeError('Failed to copy at least %d part(s)' % len(exceptions))
+                assert len(filter(None, parts)) == totalParts
         except:
             with panic(log=log):
                 upload.cancel_upload()
         else:
             for attempt in retry_s3():
                 with attempt:
-                    return upload.complete_upload()
+                    completed = upload.complete_upload()
+                    log.info("Completed copy from 's3://%s/%s' to 's3://%s/%s'.",
+                             srcKey.bucket.name, srcKey.name, dstBucketName, dstKeyName)
+                    return completed
 
 
 class AWSJobStore(AbstractJobStore):
@@ -148,7 +181,7 @@ def loadOrCreateJobStore(cls, jobStoreString, config=None, **kwargs):
 
     # Do not invoke the constructor, use the factory method above.
 
-    def __init__(self, region, namePrefix, config=None, partSize=defaultPartSize):
+    def __init__(self, region, namePrefix, config=None, partSize=50 << 20):
         """
         Create a new job store in AWS or load an existing one from there.
 
@@ -158,6 +191,10 @@ def __init__(self, region, namePrefix, config=None, partSize=defaultPartSize):
 
         :param config: the config object to written to this job store. Must be None for existing
         job stores. Must not be None for new job stores.
+
+        :param int partSize: The size of each individual part used for multipart operations like
+               upload and copy, must be >= 5 MiB but large enough to not exceed 10k parts for the
+               whole file
         """
         log.debug("Instantiating %s for region %s and name prefix '%s'",
                   self.__class__, region, namePrefix)
@@ -351,26 +388,39 @@ def _getKeyForUrl(url, existing=None):
         :rtype: Key
         """
         # Get the bucket's region to avoid a redirect per request
-        with closing(boto.connect_s3()) as s3:
-            region = bucket_location_to_region(s3.get_bucket(url.netloc).get_location())
+        try:
+            with closing(boto.connect_s3()) as s3:
+                region = bucket_location_to_region(s3.get_bucket(url.netloc).get_location())
+        except S3ResponseError as e:
+            if e.error_code == 'AccessDenied':
+                log.warn("Could not determine location of bucket hosting URL '%s', reverting "
+                         "to generic S3 endpoint.", url.geturl())
+                s3 = boto.connect_s3()
+            else:
+                raise
+        else:
+            # Note that caller is responsible for closing the connection
+            s3 = boto.s3.connect_to_region(region)
 
-        # Note that caller is responsible for closing the connection
-        s3 = boto.s3.connect_to_region(region)
         try:
-            bucket = s3.get_bucket(url.netloc)
-            key = bucket.get_key(url.path[1:])
+            keyName = url.path[1:]
+            bucketName = url.netloc
+            bucket = s3.get_bucket(bucketName)
+            key = bucket.get_key(keyName)
             if existing is True:
                 if key is None:
-                    raise RuntimeError('Key does not exist.')
+                    raise RuntimeError("Key '%s' does not exist in bucket '%s'." %
+                                       (keyName, bucketName))
             elif existing is False:
                 if key is not None:
-                    raise RuntimeError('Key exists.')
+                    raise RuntimeError("Key '%s' exists in bucket '%s'." %
+                                       (keyName, bucketName))
             elif existing is None:
                 pass
             else:
                 assert False
             if key is None:
-                key = bucket.new_key(url.path[1:])
+                key = bucket.new_key(keyName)
         except:
             with panic():
                 s3.close()
@@ -974,6 +1024,7 @@ def _copyKey(self, srcKey, dstBucketName, dstKeyName, headers=None):
                 return copyKeyMultipart(srcKey=srcKey,
                                         dstBucketName=dstBucketName,
                                         dstKeyName=dstKeyName,
+                                        partSize=self.outer.partSize,
                                         headers=headers)
             else:
                 # We need a location-agnostic connection to S3 so we can't use the one that we


Test patch:
diff --git a/src/toil/test/jobStores/jobStoreTest.py b/src/toil/test/jobStores/jobStoreTest.py
--- a/src/toil/test/jobStores/jobStoreTest.py
+++ b/src/toil/test/jobStores/jobStoreTest.py
@@ -16,6 +16,7 @@
 
 import hashlib
 import logging
+
 import os
 import shutil
 import tempfile
@@ -25,12 +26,13 @@
 import uuid
 from Queue import Queue
 from abc import abstractmethod, ABCMeta
-from itertools import chain, islice
+from itertools import chain, islice, count
 from threading import Thread
 from unittest import skip
 
 from bd2k.util import memoize
 from bd2k.util.exceptions import panic
+from mock import patch
 
 from toil.common import Config
 from toil.jobStores.abstractJobStore import (AbstractJobStore,
@@ -813,6 +815,44 @@ def testInlinedFiles(self):
                 with master.readSharedFileStream('foo') as f:
                     self.assertEqual(s, f.read())
 
+    def testInaccessableLocation(self):
+        # This could break if 1000genomes enables GetBucketLocation in their S3 bucket policy,
+        # see https://rt.sanger.ac.uk/Ticket/Display.html?id=534925. If that happens, we'll have
+        # to set up or mock a bucket that disallows it.
+        url = 's3://1000genomes/20131219.populations.tsv'
+        with patch('toil.jobStores.aws.jobStore.log') as mock_log:
+            jobStoreID = self.master.importFile(url)
+            self.assertTrue(self.master.fileExists(jobStoreID))
+            args, kwargs = mock_log.warn.call_args
+            self.assertTrue('Could not determine location' in args[0])
+
+    def testMultiPartImportFailures(self):
+        # This should be less than the number of threads in the pool used by the MP copy.
+        num_parts = 10
+        i = count()
+
+        # noinspection PyUnusedLocal
+        def fail(*args, **kwargs):
+            # The sleep ensures that all tasks are scheduled in the thread pool. Without it,
+            # there is a chance that one task fails before another is scheduled, causing the
+            # latter to bail out immediatly and failing the assertion that ensure the number of
+            # failing tasks.
+            time.sleep(.25)
+            if i.next() % 2 == 0:
+                raise RuntimeError()
+
+        with patch('boto.s3.multipart.MultiPartUpload.copy_part_from_key',
+                   new_callable=lambda: fail):
+            self.master.partSize = self.mpTestPartSize
+            bucket = self._createExternalStore()
+            url, md5 = self._prepareTestFile(bucket, self.mpTestPartSize * num_parts)
+            try:
+                self.master.importFile(url)
+            except RuntimeError as e:
+                self.assertEquals(e.message, 'Failed to copy at least %d part(s)' % (num_parts / 2))
+            else:
+                self.fail('Expected a RuntimeError to be raised')
+
     def _prepareTestFile(self, bucket, size=None):
         fileName = 'testfile_%s' % uuid.uuid4()
         url = 's3://%s/%s' % (bucket.name, fileName)
diff --git a/src/toil/test/src/importExportFileTest.py b/src/toil/test/src/importExportFileTest.py
--- a/src/toil/test/src/importExportFileTest.py
+++ b/src/toil/test/src/importExportFileTest.py
@@ -19,7 +19,7 @@
 from toil.common import Toil
 from toil.job import Job
 from toil.leader import FailedJobsException
-from toil.test import ToilTest, make_tests
+from toil.test import ToilTest
 
 
 class ImportExportFileTest(ToilTest):
diff --git a/src/toil/test/src/multipartTransferTest.py b/src/toil/test/src/multipartTransferTest.py
new file mode 100644
--- /dev/null
+++ b/src/toil/test/src/multipartTransferTest.py
@@ -0,0 +1,82 @@
+# Copyright (C) 2015 UCSC Computational Genomics Lab
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import absolute_import
+
+import logging
+import os
+import uuid
+from contextlib import contextmanager, closing
+
+import boto
+import boto.s3
+
+from toil.jobStores.aws.jobStore import copyKeyMultipart
+from toil.jobStores.aws.utils import region_to_bucket_location
+from toil.test import ToilTest, make_tests
+
+partSize = 2 ** 20 * 5
+logger = logging.getLogger(__name__)
+
+
+@contextmanager
+def openS3(keySize=None):
+    """
+    Creates an AWS bucket. If keySize is given, a key of random bytes is created and its handle
+    is yielded. If no keySize is given an empty bucket handle is yielded. The bucket and all
+    created keys are cleaned up automatically.
+
+    :param int keySize: Size of key to be created.
+    """
+    if keySize is not None and keySize < 0:
+        raise ValueError('Key size must be greater than zero')
+    with closing(boto.s3.connect_to_region(AWSMultipartCopyTest.region)) as s3:
+        bucket = s3.create_bucket('multipart-transfer-test-%s' % uuid.uuid4(),
+                                  location=region_to_bucket_location(AWSMultipartCopyTest.region))
+        try:
+            keyName = 'test'
+            if keySize is None:
+                yield bucket
+            else:
+                key = bucket.new_key(keyName)
+                content = os.urandom(keySize)
+                key.set_contents_from_string(content)
+
+                yield bucket.get_key(keyName)
+        finally:
+            for key in bucket.list():
+                key.delete()
+            bucket.delete()
+
+
+class AWSMultipartCopyTest(ToilTest):
+    region = 'us-west-2'
+
+    @classmethod
+    def makeTests(cls):
+        def multipartCopy(self, threadPoolSize):
+            # key size is padded to ensure some threads are reused
+            keySize = int((threadPoolSize * partSize) * 1.3)
+            with openS3(keySize) as srcKey:
+                with openS3() as dstBucket:
+                    copyKeyMultipart(srcKey, dstBucket, 'test', partSize)
+                    self.assertEqual(srcKey.get_contents_as_string(),
+                                     dstBucket.get_key('test').get_contents_as_string())
+
+        make_tests(multipartCopy,
+                   targetClass=AWSMultipartCopyTest,
+                   threadPoolSize={str(x): x for x in (1, 2, 16)})
+
+
+AWSMultipartCopyTest.makeTests()