Repository: DataBiosphere/toil

Issue:
Add checksum verification to AWS S3 jobstore
Large files downloaded from S3 can get corrupted for a number of reasons (e.g. due to network or disk issues) and we've seen this happen very rarely when running large scale pipelines. However, when it does happen, it can lead to unexpected failures, or worse, incorrect results [*].

S3 does not have a good checksum value for large files that get uploaded via multipart uploads (see: https://teppen.io/2018/06/23/aws_s3_etags/ ), so the easiest way is to add the `sha1` checksum as an additional value to the `SimpleDB` data and use that when downloading/uploading files to/from S3. I propose using `sha1` as it's also used in CWL.

Ideally other jobstores should also have this, but I've limited the scope to S3 for this issue.

[*] I had added a checksum check on our internal branch to see whether it would ever fire, and just saw an instance of it happening in a production Toil pipeline! So, decided to prioritize adding it to master as well.

┆Issue is synchronized with this [Jira Task](https://ucsc-cgl.atlassian.net/browse/TOIL-461)
┆Issue Number: TOIL-461



Discussion:

Resolving patch:
diff --git a/src/toil/common.py b/src/toil/common.py
--- a/src/toil/common.py
+++ b/src/toil/common.py
@@ -120,6 +120,7 @@ def __init__(self):
         # Misc
         self.disableCaching = True
         self.disableChaining = False
+        self.disableJobStoreChecksumVerification = False
         self.maxLogFileSize = 64000
         self.writeLogs = None
         self.writeLogsGzip = None
@@ -268,6 +269,7 @@ def parseIntList(s):
         setOption("maxLocalJobs", int)
         setOption("disableCaching")
         setOption("disableChaining")
+        setOption("disableJobStoreChecksumVerification")
         setOption("maxLogFileSize", h2b, iC(1))
         setOption("writeLogs")
         setOption("writeLogsGzip")
@@ -539,6 +541,11 @@ def _addOptions(addGroupFn, config):
     addOptionFn('--disableChaining', dest='disableChaining', action='store_true', default=False,
                 help="Disables chaining of jobs (chaining uses one job's resource allocation "
                 "for its successor job if possible).")
+    addOptionFn("--disableJobStoreChecksumVerification", dest="disableJobStoreChecksumVerification",
+                default=False, action="store_true",
+                help=("Disables checksum verification for files transferred to/from the job store. "
+                      "Checksum verification is a safety check to ensure the data is not corrupted "
+                      "during transfer. Currently only supported for non-streaming AWS files."))
     addOptionFn("--maxLogFileSize", dest="maxLogFileSize", default=None,
                 help=("The maximum size of a job log file to keep (in bytes), log files "
                       "larger than this will be truncated to the last X bytes. Setting "
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
@@ -531,7 +531,7 @@ def _supportsUrl(cls, url, export=False):
 
     def writeFile(self, localFilePath, jobStoreID=None, cleanup=False):
         info = self.FileInfo.create(jobStoreID if cleanup else None)
-        info.upload(localFilePath)
+        info.upload(localFilePath, not self.config.disableJobStoreChecksumVerification)
         info.save()
         log.debug("Wrote %r of from %r", info, localFilePath)
         return info.fileID
@@ -557,7 +557,7 @@ def writeSharedFileStream(self, sharedFileName, isProtected=None):
 
     def updateFile(self, jobStoreFileID, localFilePath):
         info = self.FileInfo.loadOrFail(jobStoreFileID)
-        info.upload(localFilePath)
+        info.upload(localFilePath, not self.config.disableJobStoreChecksumVerification)
         info.save()
         log.debug("Wrote %r from path %r.", info, localFilePath)
 
@@ -581,7 +581,7 @@ def getFileSize(self, jobStoreFileID):
     def readFile(self, jobStoreFileID, localFilePath, symlink=False):
         info = self.FileInfo.loadOrFail(jobStoreFileID)
         log.debug("Reading %r into %r.", info, localFilePath)
-        info.download(localFilePath)
+        info.download(localFilePath, not self.config.disableJobStoreChecksumVerification)
 
     @contextmanager
     def readFileStream(self, jobStoreFileID):
@@ -841,7 +841,7 @@ class FileInfo(SDBHelper):
         """
 
         def __init__(self, fileID, ownerID, encrypted,
-                     version=None, content=None, numContentChunks=0):
+                     version=None, content=None, numContentChunks=0,  checksum=None):
             """
             :type fileID: str
             :param fileID: the file's ID
@@ -862,6 +862,10 @@ def __init__(self, fileID, ownerID, encrypted,
 
             :type numContentChunks: int
             :param numContentChunks: the number of SDB domain attributes occupied by this files
+
+            :type checksum: str|None
+            :param checksum: the sha1 checksum of the file, if available.
+
             inlined content. Note that an inlined empty string still occupies one chunk.
             """
             super(AWSJobStore.FileInfo, self).__init__()
@@ -871,6 +875,7 @@ def __init__(self, fileID, ownerID, encrypted,
             self._version = version
             self._previousVersion = version
             self._content = content
+            self._checksum = checksum
             self._numContentChunks = numContentChunks
 
         @property
@@ -901,6 +906,14 @@ def previousVersion(self):
         def content(self):
             return self._content
 
+        @property
+        def checksum(self):
+            return self._checksum
+
+        @checksum.setter
+        def checksum(self, checksum):
+            self._checksum = checksum
+
         @content.setter
         def content(self, content):
             assert content is None or isinstance(content, bytes)
@@ -981,6 +994,7 @@ def strOrNone(s):
                 return None
             else:
                 version = strOrNone(item['version'])
+                checksum = strOrNone(item.get('checksum'))
                 encrypted = strict_bool(encrypted)
                 content, numContentChunks = cls.attributesToBinary(item)
                 if encrypted:
@@ -990,7 +1004,7 @@ def strOrNone(s):
                     if content is not None:
                         content = encryption.decrypt(content, sseKeyPath)
                 self = cls(fileID=item.name, ownerID=ownerID, encrypted=encrypted, version=version,
-                           content=content, numContentChunks=numContentChunks)
+                           content=content, numContentChunks=numContentChunks, checksum=checksum)
                 return self
 
         def toItem(self):
@@ -1014,7 +1028,8 @@ def toItem(self):
             numChunks = attributes['numChunks']
             attributes.update(dict(ownerID=self.ownerID,
                                    encrypted=self.encrypted,
-                                   version=self.version or ''))
+                                   version=self.version or '',
+                                   checksum=self.checksum or ''))
             return attributes, numChunks
 
         @classmethod
@@ -1056,17 +1071,27 @@ def save(self):
                 else:
                     raise
 
-        def upload(self, localFilePath):
+        def upload(self, localFilePath, calculateChecksum=True):
             file_size, file_time = fileSizeAndTime(localFilePath)
             if file_size <= self.maxInlinedSize():
                 with open(localFilePath, 'rb') as f:
                     self.content = f.read()
             else:
                 headers = self._s3EncryptionHeaders()
+                self.checksum = self._get_file_checksum(localFilePath) if calculateChecksum else None
                 self.version = uploadFromPath(localFilePath, partSize=self.outer.partSize,
                                               bucket=self.outer.filesBucket, fileID=compat_bytes(self.fileID),
                                               headers=headers)
 
+        def _get_file_checksum(self, localFilePath):
+            with open(localFilePath, 'rb') as f:
+                checksum = hashlib.sha1()
+                contents = f.read(1024 * 1024)
+                while contents != b'':
+                    checksum.update(contents)
+                    contents = f.read(1024 * 1024)
+                return 'sha1$%s' % checksum.hexdigest()
+
         @contextmanager
         def uploadStream(self, multipart=True, allowInlining=True):
             """
@@ -1255,7 +1280,7 @@ def copyTo(self, dstKey):
             else:
                 assert False
 
-        def download(self, localFilePath):
+        def download(self, localFilePath, verifyChecksum=True):
             if self.content is not None:
                 with open(localFilePath, 'wb') as f:
                     f.write(self.content)
@@ -1267,6 +1292,12 @@ def download(self, localFilePath):
                         key.get_contents_to_filename(localFilePath,
                                                      version_id=self.version,
                                                      headers=headers)
+                if verifyChecksum and self.checksum:
+                    downloadedChecksum = self._get_file_checksum(localFilePath)
+                    if self.checksum != downloadedChecksum:
+                        raise AssertionError(
+                            'Checksums do not match for file %s. Expected: %s Actual: %s' % (
+                                localFilePath, self.checksum, downloadedChecksum))
             else:
                 assert False
 
@@ -1350,6 +1381,7 @@ def __repr__(self):
                  ('version', r(self.version)),
                  ('previousVersion', r(self.previousVersion)),
                  ('content', r(self.content)),
+                 ('checksum', r(self.checksum)),
                  ('_numContentChunks', r(self._numContentChunks)))
             return "{}({})".format(type(self).__name__,
                                    ', '.join('%s=%s' % (k, v) for k, v in d))


Test patch: