Limit S3 concurrent downloads to fix Java Heap OOM (#6402)

Signed-off-by: jorgee <[email protected]>
Signed-off-by: Jorge Ejarque <[email protected]>
Signed-off-by: Paolo Di Tommaso <[email protected]>
Co-authored-by: Chris Hakkaart <[email protected]>
Co-authored-by: Paolo Di Tommaso <[email protected]>
Co-authored-by: Claude <[email protected]>

Author: 4 people

docs/guides/aws-java-sdk-v2.md

@@ -22,9 +22,10 @@ The *S3 transfer manager* is a subsystem of SDK v2 that handles S3 uploads and d
 
 You can configure the concurrency and throughput of the S3 transfer manager manually using the `aws.client.maxConcurrency` and `aws.client.maxNativeMemory` configuration options. Alternatively, you can use the `aws.client.targetThroughputInGbps` option to set both values automatically based on a target throughput.
 
-## Multi-part uploads
+## Multi-part transfers
 
-Multi-part uploads are handled by the S3 transfer manager. You can use the `aws.client.minimumPartSize` and `aws.client.multipartThreshold` config options to control when and how multi-part uploads are performed.
+Multi-part transfer are handled by the S3 transfer manager. You can use the `aws.client.minimumPartSize` and `aws.client.multipartThreshold` config options to control when and how multi-part transfers are performed. 
+Concurrent multi-part downloads can consume large heap memory space due to the buffer size created per transfer. To avoid out of memory errors, the size consumed by these buffers is limited to 400 MB. You can use `aws.client.maxDownloadHeapMemory` to change this value.
 
 The following multi-part upload config options are no longer supported:
 

docs/reference/config.md

@@ -175,6 +175,9 @@ The following settings are available:
 `aws.client.maxConnections`
 : The maximum number of open HTTP connections used by the S3 transfer manager (default: `50`).
 
+`aws.client.maxDownloadHeapMemory`
+: The maximum size for the heap memory buffer used by concurrent downloads. It must be at least 10 times the `minimumPartSize` (default:`400 MB`).
+
 `aws.client.maxErrorRetry`
 : The maximum number of retry attempts for failed retryable requests (default: `-1`).
 

plugins/nf-amazon/build.gradle

@@ -55,20 +55,20 @@ dependencies {
     compileOnly 'org.pf4j:pf4j:3.12.0'
 
     api ('javax.xml.bind:jaxb-api:2.4.0-b180830.0359')
-    api ('software.amazon.awssdk:s3:2.31.64')
-    api ('software.amazon.awssdk:ec2:2.31.64')
-    api ('software.amazon.awssdk:batch:2.31.64')
-    api ('software.amazon.awssdk:iam:2.31.64')
-    api ('software.amazon.awssdk:ecs:2.31.64')
-    api ('software.amazon.awssdk:cloudwatchlogs:2.31.64')
-    api ('software.amazon.awssdk:codecommit:2.31.64')
-    api ('software.amazon.awssdk:sts:2.31.64')
-    api ('software.amazon.awssdk:ses:2.31.64')
-    api ('software.amazon.awssdk:sso:2.31.64')
-    api ('software.amazon.awssdk:ssooidc:2.31.64')
-    api ('software.amazon.awssdk:s3-transfer-manager:2.31.64')
-    api ('software.amazon.awssdk:apache-client:2.31.64')
-    api ('software.amazon.awssdk:aws-crt-client:2.31.64')
+    api ('software.amazon.awssdk:s3:2.33.2')
+    api ('software.amazon.awssdk:ec2:2.33.2')
+    api ('software.amazon.awssdk:batch:2.33.2')
+    api ('software.amazon.awssdk:iam:2.33.2')
+    api ('software.amazon.awssdk:ecs:2.33.2')
+    api ('software.amazon.awssdk:cloudwatchlogs:2.33.2')
+    api ('software.amazon.awssdk:codecommit:2.33.2')
+    api ('software.amazon.awssdk:sts:2.33.2')
+    api ('software.amazon.awssdk:ses:2.33.2')
+    api ('software.amazon.awssdk:sso:2.33.2')
+    api ('software.amazon.awssdk:ssooidc:2.33.2')
+    api ('software.amazon.awssdk:s3-transfer-manager:2.33.2')
+    api ('software.amazon.awssdk:apache-client:2.33.2')
+    api ('software.amazon.awssdk:aws-crt-client:2.33.2')
 
     constraints {
         api 'com.fasterxml.jackson.core:jackson-databind:2.12.7.1'

plugins/nf-amazon/src/main/nextflow/cloud/aws/config/AwsS3Config.groovy

@@ -70,6 +70,12 @@ class AwsS3Config implements ConfigScope {
     """)
     final Integer maxConnections
 
+    @ConfigOption
+    @Description("""
+        The maximum size for the heap memory buffer used by concurrent downloads. It must be at least 10 times the `minimumPartSize` (default:`400 MB`).
+    """)
+    final MemoryUnit maxDownloadHeapMemory
+
     @ConfigOption
     @Description("""
         The maximum number of retry attempts for failed retryable requests (default: `-1`).
@@ -215,6 +221,13 @@ class AwsS3Config implements ConfigScope {
     """)
     final String uploadStorageClass
 
+    private static final long _1MB = 1024 * 1024;
+    // According to CRT Async client docs https://sdk.amazonaws.com/java/api/latest/software/amazon/awssdk/services/s3/S3CrtAsyncClientBuilder.html
+    public static final long DEFAULT_PART_SIZE = 8 * _1MB;
+    public static final int DEFAULT_INIT_BUFFER_PARTS = 10;
+    // Maximum heap buffer size
+    public static final long DEFAULT_MAX_DOWNLOAD_BUFFER_SIZE = 400 * _1MB;
+
     AwsS3Config(Map opts) {
         this.anonymous = opts.anonymous as Boolean
         this.connectionTimeout = opts.connectionTimeout as Integer
@@ -224,6 +237,7 @@ class AwsS3Config implements ConfigScope {
             throw new IllegalArgumentException("S3 endpoint must begin with http:// or https:// prefix - offending value: '${endpoint}'")
         this.maxConcurrency = opts.maxConcurrency as Integer
         this.maxConnections = opts.maxConnections as Integer
+        this.maxDownloadHeapMemory = opts.maxDownloadHeapMemory as MemoryUnit
         this.maxErrorRetry = opts.maxErrorRetry as Integer
         this.maxNativeMemory = opts.maxNativeMemory as MemoryUnit
         this.minimumPartSize = opts.minimumPartSize as MemoryUnit
@@ -246,6 +260,7 @@ class AwsS3Config implements ConfigScope {
         this.uploadMaxAttempts = opts.uploadMaxAttempts as Integer
         this.uploadMaxThreads = opts.uploadMaxThreads as Integer
         this.uploadRetrySleep = opts.uploadRetrySleep as Duration
+        checkDownloadBufferParams()
     }
 
     private String parseStorageClass(String value) {
@@ -283,6 +298,7 @@ class AwsS3Config implements ConfigScope {
             connection_timeout: connectionTimeout?.toString(),
             max_concurrency: maxConcurrency?.toString(),
             max_connections: maxConnections?.toString(),
+            max_download_heap_memory: maxDownloadHeapMemory?.toBytes()?.toString(),
             max_error_retry: maxErrorRetry?.toString(),
             max_native_memory: maxNativeMemory?.toBytes()?.toString(),
             minimum_part_size: minimumPartSize?.toBytes()?.toString(),
@@ -306,4 +322,21 @@ class AwsS3Config implements ConfigScope {
             upload_storage_class: storageClass?.toString()
         ].findAll { k, v -> v != null }
     }
+
+    void checkDownloadBufferParams() {
+        if( maxDownloadHeapMemory != null  && maxDownloadHeapMemory.toBytes() == 0L ) {
+            throw new IllegalArgumentException("Configuration option `aws.client.maxDownloadHeapMemory` can't be 0")
+        }
+        if( minimumPartSize != null && minimumPartSize.toBytes() == 0L ) {
+            throw new IllegalArgumentException("Configuration option `aws.client.minimumPartSize` can't be 0")
+        }
+        if( maxDownloadHeapMemory != null || minimumPartSize != null ) {
+            final maxBuffer = maxDownloadHeapMemory ? maxDownloadHeapMemory.toBytes() : DEFAULT_MAX_DOWNLOAD_BUFFER_SIZE
+            final partSize = minimumPartSize ? minimumPartSize.toBytes() : DEFAULT_PART_SIZE
+            if( maxBuffer < DEFAULT_INIT_BUFFER_PARTS * partSize ) {
+                throw new IllegalArgumentException("Configuration option `aws.client.maxDownloadHeapMemory` must be at least " + DEFAULT_INIT_BUFFER_PARTS + " times `aws.client.minimumPartSize`")
+            }
+        }
+
+    }
 }