EVA-4318 - Split large files & add new fileParts to the JSON

eva_sub_cli/__init__.py

@@ -14,6 +14,7 @@
 SUBMISSION_WS_URL_VAR = 'SUBMISSION_WS_URL'
 ENA_WEBIN_ACCOUNT_VAR = 'ENA_WEBIN_ACCOUNT'
 ENA_WEBIN_PASSWORD_VAR = 'ENA_WEBIN_PASSWORD'
+UPLOAD_PART_SIZE_VAR = 'UPLOAD_PART_SIZE'
 
 
 try:

eva_sub_cli/etc/eva_schema.json

@@ -380,6 +380,32 @@
           "md5": {
             "type": "string",
             "description": "MD5 hash of the file"
+          },
+          "fileParts": {
+            "type": "array",
+            "description": "Parts the file was split into for upload, in order. Concatenated, they form the file.",
+            "minItems": 1,
+            "items": {
+              "type": "object",
+              "required": [
+                "fileName",
+                "fileSize"
+              ],
+              "properties": {
+                "fileName": {
+                  "type": "string",
+                  "description": "Name of the part"
+                },
+                "fileSize": {
+                  "type": "number",
+                  "description": "Size of the part in bytes"
+                },
+                "md5": {
+                  "type": "string",
+                  "description": "MD5 hash of the part"
+                }
+              }
+            }
           }
         }
       }

eva_sub_cli/exceptions.py

@@ -40,6 +40,18 @@ def __init__(self, message):
         super().__init__(self.message)
 
 
+class RetryableUploadException(Exception):
+    def __init__(self, message):
+        self.message = message
+        super().__init__(self.message)
+
+
+class FileTooLargeException(Exception):
+    def __init__(self, message):
+        self.message = message
+        super().__init__(self.message)
+
+
 class DirLockError(Exception):
     def __init__(self, message):
         self.message = message

eva_sub_cli/submit.py

@@ -1,26 +1,91 @@
 #!/usr/bin/env python
+import functools
+import hashlib
 import json
+import math
 import os
+import time
 
 import requests
 from ebi_eva_common_pyutils.config import WritableConfig
 from ebi_eva_common_pyutils.logger import AppLogger
 from retry import retry
 
-from eva_sub_cli import SUB_CLI_CONFIG_FILE, __version__
-from eva_sub_cli.exceptions import InvalidSubmissionException
+from eva_sub_cli import SUB_CLI_CONFIG_FILE, UPLOAD_PART_SIZE_VAR, __version__
+from eva_sub_cli.exceptions import InvalidSubmissionException, SubmissionUploadException, RetryableUploadException, \
+    FileTooLargeException
 from eva_sub_cli.submission_ws import SubmissionWSClient
 from eva_sub_cli.validators.validator import READY_FOR_SUBMISSION_TO_EVA
 
 SUB_CLI_CONFIG_KEY_SUBMISSION_ID = "submission_id"
 SUB_CLI_CONFIG_KEY_SUBMISSION_UPLOAD_URL = "submission_upload_url"
 SUB_CLI_CONFIG_KEY_COMPLETE = "submission_complete"
 
+# HTTP status codes that can be caused by transient conditions and are worth retrying.
+# 404 is included because the Globus HTTPS server can return it when its storage backend fails during an upload,
+RETRYABLE_HTTP_STATUS_CODES = {408, 404, 425, 429, 500, 502, 503, 504}
+MAX_RETRY_AFTER_SECONDS = 300
+
+# Files rejected by the upload server as too large are uploaded in parts of this size,
+# below the 1 GB limit of the upload server. It can be changed with the UPLOAD_PART_SIZE environment variable (bytes)
+DEFAULT_PART_SIZE = 512 * 1024 ** 2
+
+
+
+class FilePart:
+    """Read-only file-like view of a byte range of a file, which computes the md5 of the data read."""
+
+    def __init__(self, path, offset, size):
+        self.size = size
+        self._remaining = size
+        self._md5 = hashlib.md5()
+        self._file = open(path, 'rb')
+        self._file.seek(offset)
+
+    def read(self, amt=-1):
+        if self._remaining <= 0:
+            return b''
+        if amt is None or amt < 0 or amt > self._remaining:
+            amt = self._remaining
+        data = self._file.read(amt)
+        self._remaining -= len(data)
+        self._md5.update(data)
+        return data
+
+    def __len__(self):
+        # Used by requests to set the Content-Length header
+        return self.size
+
+    @property
+    def md5(self):
+        return self._md5.hexdigest()
+
+    def close(self):
+        self._file.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
+
+def get_part_size():
+    """Size in bytes of the parts when a file is uploaded in parts."""
+    part_size = os.environ.get(UPLOAD_PART_SIZE_VAR)
+    if not part_size:
+        return DEFAULT_PART_SIZE
+    if not part_size.isdigit() or int(part_size) == 0:
+        raise SubmissionUploadException(f'{UPLOAD_PART_SIZE_VAR} must be a positive number of bytes, not {part_size}')
+    return int(part_size)
+
 
 class StudySubmitter(AppLogger):
     def __init__(self, submission_dir, submission_config: WritableConfig = None, username=None, password=None):
         self.submission_ws_client = SubmissionWSClient(username, password)
         self.submission_dir = submission_dir
+        # Parts of the files uploaded in parts during this run, per file path
+        self.file_parts = {}
         if submission_config:
             self.sub_config = submission_config
         else:
@@ -58,28 +123,114 @@ def _upload_submission(self):
 
         submission_upload_url = self.sub_config[SUB_CLI_CONFIG_KEY_SUBMISSION_UPLOAD_URL]
 
-        self._upload_file(submission_upload_url, self.metadata_json_file)
-
         for f in self.vcf_files:
-            self._upload_file(submission_upload_url, f)
+            try:
+                self._upload_file(submission_upload_url, f)
+            except FileTooLargeException:
+                self.warning(f'{os.path.basename(f)} is too large to be uploaded in one go: uploading it in parts')
+                self._upload_file_in_parts(submission_upload_url, f)
+
+        # The metadata json is uploaded last so that it describes the files uploaded in parts
+        self._upload_file(submission_upload_url, self._write_metadata_for_submission())
 
-    @retry(tries=5, delay=10, backoff=5)
     def _upload_file(self, submission_upload_url, input_file):
         base_name = os.path.basename(input_file)
         self.debug(f'Transfer {base_name} to EVA FTP')
-        with open(input_file, 'rb') as f:
-            r = requests.put(os.path.join(submission_upload_url, base_name), data=f)
-        r.raise_for_status()
+        self._put(os.path.join(submission_upload_url, base_name), lambda: open(input_file, 'rb'), base_name)
         self.debug(f'Upload of {base_name} completed')
 
+    @retry(exceptions=(RetryableUploadException, requests.ConnectionError, requests.Timeout),
+           tries=5, delay=10, backoff=5, max_delay=300)
+    def _put(self, url, open_data, name):
+        """Upload the data returned by open_data, which is called on each attempt. Returns the data object, closed."""
+        with open_data() as data:
+            r = requests.put(url, data=data)
+        try:
+            r.raise_for_status()
+        except requests.HTTPError as e:
+            self._handle_upload_error(r, name, e)
+        return data
+
+    def _upload_file_in_parts(self, submission_upload_url, input_file):
+        """Upload a file as consecutive parts of the size returned by get_part_size.
+        The parts are recorded so that they can be described in the metadata."""
+        base_name = os.path.basename(input_file)
+        file_size = os.path.getsize(input_file)
+        try:
+            parts = self._upload_parts(submission_upload_url, input_file, file_size, get_part_size())
+        except FileTooLargeException:
+            raise SubmissionUploadException(
+                f'{base_name} could not be uploaded even when split into parts. Please contact EVA Helpdesk.'
+            )
+        if sum(part['fileSize'] for part in parts) != file_size:
+            raise SubmissionUploadException(f'The parts of {base_name} do not add up to its size of {file_size} bytes')
+        self.file_parts[input_file] = parts
+        self.info(f'{base_name} uploaded in {len(parts)} parts')
+
+    def _upload_parts(self, submission_upload_url, input_file, file_size, part_size):
+        base_name = os.path.basename(input_file)
+        number_of_parts = max(1, math.ceil(file_size / part_size))
+        parts = []
+        for part_index in range(number_of_parts):
+            offset = part_index * part_size
+            size = min(part_size, file_size - offset)
+            # {base_name}.part001 - {base_name}.part999
+            part_name = f'{base_name}.part{part_index + 1:0{3}d}'
+            self.debug(f'Transfer {part_name} to EVA FTP')
+            # Pass a factory rather than a FilePart: each retry of _put needs a new reader starting at the
+            # beginning of the part, since a failed attempt consumes the reader and its md5
+            data = self._put(
+                os.path.join(submission_upload_url, part_name),
+                functools.partial(FilePart, input_file, offset, size),
+                part_name
+            )
+            parts.append({'fileName': part_name, 'fileSize': size, 'md5': data.md5})
+        return parts
+
+    def _metadata_for_submission(self):
+        """Metadata json with the description of the parts for the files uploaded in parts."""
+        metadata = self.metadata_json
+        parts_per_file_name = {
+            os.path.basename(vcf_file): parts
+            for vcf_file, parts in self.file_parts.items() if vcf_file in self.vcf_files
+        }
+        for file_dict in metadata.get('files', []):
+            file_name = os.path.basename(file_dict.get('fileName') or '')
+            if file_name in parts_per_file_name:
+                file_dict['fileParts'] = parts_per_file_name[file_name]
+            else:
+                file_dict.pop('fileParts', None)
+        return metadata
+
+    def _write_metadata_for_submission(self):
+        """Write the metadata to upload, without modifying the metadata json provided or generated during validation."""
+        upload_dir = os.path.join(self.submission_dir, '.metadata_upload_in_parts')
+        os.makedirs(upload_dir, exist_ok=True)
+        metadata_file = os.path.join(upload_dir, os.path.basename(self.metadata_json_file))
+        with open(metadata_file, 'w') as open_file:
+            json.dump(self._metadata_for_submission(), open_file)
+        return metadata_file
+
+    def _handle_upload_error(self, response, base_name, error):
+        if response.status_code in RETRYABLE_HTTP_STATUS_CODES:
+            retry_after = response.headers.get('Retry-After') # Wait for whatever the server says on top of the retry delay
+            if retry_after and retry_after.isdigit():
+                time.sleep(min(int(retry_after), MAX_RETRY_AFTER_SECONDS))
+            self.warning(f'Upload of {base_name} failed with error ({error}). It will be retried.')
+            raise RetryableUploadException(str(error)) from error
+        if response.status_code == 413:
+            raise FileTooLargeException(f'{base_name} is too large to be uploaded in one request') from error
+        raise SubmissionUploadException(f'Upload of {base_name} failed: {error}') from error
+
     def _initiate_submission(self):
         response_json = self.submission_ws_client.initiate_submission()
         self.debug(f'Submission ID {response_json["submissionId"]} received!!')
         # update config with submission id and upload url
         self.update_config_with_submission_id_and_upload_url(response_json["submissionId"], response_json["uploadUrl"])
 
     def _complete_submission(self):
-        self.submission_ws_client.mark_submission_uploaded(self.sub_config.get(SUB_CLI_CONFIG_KEY_SUBMISSION_ID), self.metadata_json)
+        self.submission_ws_client.mark_submission_uploaded(self.sub_config.get(SUB_CLI_CONFIG_KEY_SUBMISSION_ID),
+                                                           self._metadata_for_submission())
         self.debug("Submission ID {} Complete".format(self.sub_config.get(SUB_CLI_CONFIG_KEY_SUBMISSION_ID)))
         # update config with completion of the submission
         self.sub_config.set(SUB_CLI_CONFIG_KEY_COMPLETE, value=True)