feat(_put_file): dynamically adjust chunksize based on file size (#990)

Co-authored-by: Martin Durant <[email protected]>

Author: skshetry

s3fs/core.py

@@ -3,6 +3,7 @@
 import errno
 import io
 import logging
+import math
 import mimetypes
 import os
 import socket
@@ -69,6 +70,8 @@ def setup_logging(level=None):
     ResponseParserError,
 )
 
+MAX_UPLOAD_PARTS = 10_000  # maximum number of parts for S3 multipart upload
+
 if ClientPayloadError is not None:
     S3_RETRYABLE_ERRORS += (ClientPayloadError,)
 
@@ -174,6 +177,18 @@ def _coalesce_version_id(*args):
         return version_ids.pop()
 
 
+def calculate_chunksize(filesize, chunksize=None, max_parts=MAX_UPLOAD_PARTS) -> int:
+    if chunksize is None:
+        chunksize = 50 * 2**20  # default chunksize set to 50 MiB
+        required_chunks = math.ceil(filesize / chunksize)
+        # increase chunksize to fit within the max_parts limit
+        if required_chunks > max_parts:
+            # S3 supports uploading objects up to 5 TiB in size,
+            # so each chunk can be up to ~524 MiB.
+            chunksize = math.ceil(filesize / max_parts)
+    return chunksize
+
+
 class S3FileSystem(AsyncFileSystem):
     """
     Access S3 as if it were a file system.
@@ -1242,7 +1257,7 @@ async def _put_file(
         lpath,
         rpath,
         callback=_DEFAULT_CALLBACK,
-        chunksize=50 * 2**20,
+        chunksize=None,
         max_concurrency=None,
         mode="overwrite",
         **kwargs,
@@ -1270,6 +1285,7 @@ async def _put_file(
             if content_type is not None:
                 kwargs["ContentType"] = content_type
 
+        chunksize = calculate_chunksize(size, chunksize=chunksize)
         with open(lpath, "rb") as f0:
             if size < min(5 * 2**30, 2 * chunksize):
                 chunk = f0.read()
@@ -1288,8 +1304,8 @@ async def _put_file(
                         key,
                         mpu,
                         f0,
+                        chunksize,
                         callback=callback,
-                        chunksize=chunksize,
                         max_concurrency=max_concurrency,
                     )
                     parts = [
@@ -1317,8 +1333,8 @@ async def _upload_file_part_concurrent(
         key,
         mpu,
         f0,
+        chunksize,
         callback=_DEFAULT_CALLBACK,
-        chunksize=50 * 2**20,
         max_concurrency=None,
     ):
         max_concurrency = max_concurrency or self.max_concurrency

s3fs/tests/test_s3fs.py

@@ -20,7 +20,7 @@
 
 import botocore
 import s3fs.core
-from s3fs.core import S3FileSystem
+from s3fs.core import MAX_UPLOAD_PARTS, S3FileSystem, calculate_chunksize
 from s3fs.utils import ignoring, SSEParams
 from botocore.exceptions import NoCredentialsError
 from fsspec.asyn import sync
@@ -2995,6 +2995,31 @@ def test_bucket_info(s3):
     assert info["name"] == test_bucket_name
 
 
+MB = 2**20
+GB = 2**30
+TB = 2**40
+
+
+@pytest.mark.parametrize(
+    ["filesize", "chunksize", "expected"],
+    [
+        # small file, use default chunksize
+        (1000, None, 50 * MB),
+        # exact boundary, use default chunksize
+        (50 * MB * MAX_UPLOAD_PARTS, None, 50 * MB),
+        # file requiring increased chunksize
+        (50 * MB * (MAX_UPLOAD_PARTS + 1), None, 52_434_043),
+        # very large files, expect increased chunksize
+        (1 * TB, None, 109_951_163),
+        (5 * TB, None, 549_755_814),
+        # respect explicit chunksize
+        (5 * GB, 10 * MB, 10 * MB),
+    ],
+)
+def test_calculate_chunksize(filesize, chunksize, expected):
+    assert calculate_chunksize(filesize, chunksize) == expected
+
+
 def test_find_ls_fail(s3):
     # beacuse of https://github.com/fsspec/s3fs/pull/989
     client = get_boto3_client()