Merge pull request #203 from Aiven-Open/s3-handle-max-parts-limit-for-mulitpart-upload

s3: respect max num of parts per upload

Author: giacomo-alzetta-aiven

rohmu/object_storage/config.py

@@ -62,17 +62,24 @@ def calculate_azure_max_block_size() -> int:
 LOCAL_CHUNK_SIZE: Final[int] = 1024 * 1024
 
 
+S3_MAX_NUM_PARTS_PER_UPLOAD: Final[int] = 10000
+S3_MIN_PART_SIZE_MB: Final[int] = 5
+S3_MAX_PART_SIZE_MB: Final[int] = 524
+S3_MIN_PART_SIZE_BYTES: Final[int] = S3_MIN_PART_SIZE_MB * 1024 * 1024
+S3_MAX_PART_SIZE_BYTES: Final[int] = S3_MAX_PART_SIZE_MB * 1024 * 1024
+
+
 def calculate_s3_chunk_size() -> int:
     total_mem_mib = get_total_memory() or 0
     # At least 5 MiB, at most 524 MiB. Max block size used for hosts with ~210+ GB of memory
-    return max(min(int(total_mem_mib / 400), 524), 5) * 1024 * 1024
+    return max(min(int(total_mem_mib / 400), S3_MAX_PART_SIZE_MB), S3_MIN_PART_SIZE_MB) * 1024 * 1024
 
 
 # Set chunk size based on host memory. S3 supports up to 10k chunks and up to 5 TiB individual
 # files. Minimum chunk size is 5 MiB, which means max ~50 GB files can be uploaded. In order to get
 # to that 5 TiB increase the block size based on host memory; we don't want to use the max for all
 # hosts to avoid allocating too large portion of all available memory.
-S3_MULTIPART_CHUNK_SIZE: Final[int] = calculate_s3_chunk_size()
+S3_DEFAULT_MULTIPART_CHUNK_SIZE: Final[int] = calculate_s3_chunk_size()
 S3_READ_BLOCK_SIZE: Final[int] = 1024 * 1024 * 1
 
 
@@ -145,7 +152,7 @@ class S3ObjectStorageConfig(StorageModel):
     is_secure: bool = False
     is_verify_tls: bool = False
     cert_path: Optional[Path] = None
-    segment_size: int = S3_MULTIPART_CHUNK_SIZE
+    segment_size: int = S3_DEFAULT_MULTIPART_CHUNK_SIZE
     encrypted: bool = False
     proxy_info: Optional[ProxyInfo] = None
     connect_timeout: Optional[str] = None

rohmu/object_storage/s3.py

@@ -33,7 +33,9 @@
 )
 from rohmu.object_storage.config import (  # noqa: F401
     calculate_s3_chunk_size as calculate_chunk_size,
-    S3_MULTIPART_CHUNK_SIZE as MULTIPART_CHUNK_SIZE,
+    S3_DEFAULT_MULTIPART_CHUNK_SIZE as MULTIPART_CHUNK_SIZE,
+    S3_MAX_NUM_PARTS_PER_UPLOAD,
+    S3_MAX_PART_SIZE_BYTES,
     S3_READ_BLOCK_SIZE as READ_BLOCK_SIZE,
     S3AddressingStyle,
     S3ObjectStorageConfig as Config,
@@ -156,7 +158,7 @@ def __init__(
         self.read_timeout = read_timeout
         self.aws_session_token = aws_session_token
         self.use_dualstack_endpoint = use_dualstack_endpoint
-        self.multipart_chunk_size = segment_size
+        self.default_multipart_chunk_size = segment_size
         self.encrypted = encrypted
         self.s3_client: Optional[S3Client] = None
         self.location = ""
@@ -486,6 +488,29 @@ def get_file_size(self, key: str) -> int:
             else:
                 raise StorageError(f"File size lookup failed for {path}") from ex
 
+    def calculate_chunks_and_chunk_size(self, size: Optional[int]) -> tuple[int, int]:
+        """Calculate the number of chunks and chunk size for multipart upload.
+
+        If sizes provided self.default_multipart_chunk_size wil be used as first attempt,
+        if number of chunks is greater than S3_MAX_NUM_PARTS_PER_UPLOAD, chunk size will be doubled,
+        until the number of chunks is less than S3_MAX_NUM_PARTS_PER_UPLOAD.
+        """
+        if size is None:
+            return 1, self.default_multipart_chunk_size
+        chunks = math.ceil(size / self.default_multipart_chunk_size)
+        chunk_size = self.default_multipart_chunk_size
+
+        if chunks > S3_MAX_NUM_PARTS_PER_UPLOAD:
+            chunk_size = math.ceil(size / S3_MAX_NUM_PARTS_PER_UPLOAD)
+            if chunk_size > S3_MAX_PART_SIZE_BYTES:
+                raise StorageError(
+                    f"Cannot upload a file of size {size}. "
+                    f"Chunk size {chunk_size} is too big for each part of multipart upload."
+                )
+            chunks = math.ceil(size / chunk_size)
+
+        return chunks, chunk_size
+
     def multipart_upload_file_object(
         self,
         *,
@@ -500,11 +525,11 @@ def multipart_upload_file_object(
         start_of_multipart_upload = time.monotonic()
         bytes_sent = 0
 
-        chunks: int = 1
-        if size is not None:
-            chunks = math.ceil(size / self.multipart_chunk_size)
+        chunks, chunk_size = self.calculate_chunks_and_chunk_size(size)
         args, sanitized_metadata, path = self._init_args_for_multipart(key, metadata, mimetype, cache_control)
-        self.log.debug("Starting to upload multipart file: %r, size: %s, chunks: %s", path, size, chunks)
+        self.log.debug(
+            "Starting to upload multipart file: %r, size: %s, chunks: %d (chunk size: %d)", path, size, chunks, chunk_size
+        )
 
         parts: list[CompletedPartTypeDef] = []
         part_number = 1
@@ -518,7 +543,7 @@ def multipart_upload_file_object(
         mp_id = cmu_response["UploadId"]
 
         while True:
-            data = self._read_bytes(fp, self.multipart_chunk_size)
+            data = self._read_bytes(fp, chunk_size)
             if not data:
                 break
 
@@ -635,7 +660,7 @@ def store_file_object(
         upload_progress_fn: IncrementalProgressCallbackType = None,
     ) -> None:
         if not self._should_multipart(
-            fd=fd, chunk_size=self.multipart_chunk_size, default=True, metadata=metadata, multipart=multipart
+            fd=fd, chunk_size=self.default_multipart_chunk_size, default=True, metadata=metadata, multipart=multipart
         ):
             data = fd.read()
             self.store_file_from_memory(key, data, metadata, cache_control=cache_control, mimetype=mimetype)

test/object_storage/test_s3.py

@@ -8,7 +8,7 @@
 from pathlib import Path
 from pydantic.v1 import ValidationError
 from rohmu.common.models import StorageOperation
-from rohmu.errors import InvalidByteRangeError
+from rohmu.errors import InvalidByteRangeError, StorageError
 from rohmu.object_storage.base import TransferWithConcurrentUploadSupport
 from rohmu.object_storage.config import S3ObjectStorageConfig
 from rohmu.object_storage.s3 import S3Transfer
@@ -85,7 +85,7 @@ def test_store_file_object_large(infra: S3Infra) -> None:
     chunk_size = len(test_data) // 2
     file_object = BytesIO(test_data)
 
-    infra.transfer.multipart_chunk_size = chunk_size  # simulate smaller chunk size to force multiple chunks
+    infra.transfer.default_multipart_chunk_size = chunk_size  # simulate smaller chunk size to force multiple chunks
 
     metadata = {"Content-Length": len(test_data), "some-date": datetime(2022, 11, 15, 18, 30, 58, 486644)}
     infra.transfer.store_file_object(key="test_key2", fd=file_object, metadata=metadata, multipart=True)
@@ -373,3 +373,39 @@ def test_cert_path(is_verify_tls: bool, cert_path: Optional[Path], expected: Uni
         )
         mock.assert_called_once()
         assert mock.call_args[1]["verify"] == expected
+
+
+def mb_to_bytes(size: int) -> int:
+    return size * 1024 * 1024
+
+
+@pytest.mark.parametrize(
+    ("file_size", "default_multipart_chunk_size", "expected_chunks", "expected_chunk_size"),
+    [
+        # for a 100 MB file, on a system with ~4 GB of RAM, default chunk size is 9 MB, so we expect 12 chunks of 9 MB each
+        (100, 9, 12, 9),
+        (10_000, 9, 1112, 9),
+        # for a 150 GB file, on a system with ~4 GB of RAM, default chunk size is 9 MB, so we expect 10000 chunks of 15 MB
+        # each (we increase chunk size to fit the file size in 10000 chunks)
+        (150_000, 9, 10_000, 15),
+    ],
+)
+def test_calculate_chunks_and_chunk_size(
+    infra: S3Infra, file_size: int, default_multipart_chunk_size: int, expected_chunks: int, expected_chunk_size: int
+) -> None:
+    t = infra.transfer
+    t.default_multipart_chunk_size = mb_to_bytes(default_multipart_chunk_size)
+    chunks, chunk_size = t.calculate_chunks_and_chunk_size(mb_to_bytes(file_size))
+    assert chunks == expected_chunks
+    assert chunk_size == mb_to_bytes(expected_chunk_size)
+
+
+def test_calculate_chunks_and_chunk_size_error(infra: S3Infra) -> None:
+    t = infra.transfer
+    t.default_multipart_chunk_size = mb_to_bytes(9)
+    with pytest.raises(StorageError) as e:
+        t.calculate_chunks_and_chunk_size(mb_to_bytes(50000000))
+    assert (
+        str(e.value) == "Cannot upload a file of size 52428800000000. "
+        "Chunk size 5242880000 is too big for each part of multipart upload."
+    )