Repository: piskvorky/smart_open

Issue:
Upload big io.Buffer to S3
#### Problem description

I am requesting a set of files, zipping them, and then upload the zipped data to S3 using `smart_open` and a `io.BytesIO()` object. The size of the compressed data exceeds the 5 Gb S3 limit, and I know that in that case a multi-parts approach should be use (just like in `boto3`). I am using `smart_open.s3.open()` for doing this, but I do not completely understand how to configure the multi-part upload to avoid the `EntityTooLarge` error. I keep getting the error when using my code. Should I divide my file before hand or specify the number of parts? Checking the source code I don't see a `num_parts` option.

```python
 (EntityTooLarge) when calling the UploadPart operation: Your proposed upload exceeds the maximum allowed size
```

My function is the following: 

```python
def stream_time_range_s3(start_date,
                         end_date,
                         aws_key,
                         aws_secret,
                         aws_bucket_name,
                         key,
                         max_workers,
                         delta):
    """
    Download individual month directory of .grd files to local directory.

    This function will download using the ftplib all the .grd files between the
    start_date and the end_date. All dates in the NOAA NARR server are
    stored following this order:
        data
        ├── year/month
            ├── year/month/day01
            ├── year/month/day02

    Here we download the monthly directory with the user-defined dates in the
    start and end dates. 

    Params:
        - start_year str: year to start download.
        - end_year str: year to stop download.
    """

    logger = logging.getLogger(__name__)

    if not isinstance(start_date, datetime):
        start_date = datetime.strptime(start_date, '%Y-%m-%d')
    else:
        ValueError(f'{start_date} is not in the correct format or not a valid type')


    session = boto3.Session(
        aws_access_key_id=aws_key,
        aws_secret_access_key=aws_secret
    )

    base_url = 'https://nomads.ncdc.noaa.gov/data/narr'
    time = ['0000', '0300', '0600', '0900', '1200', '1500', '1800', '2100']
 
    if delta is None:
        dates = datetime_range(start_date, end_date, {'days':1})
    else:
        dates = datetime_range(start_date, end_date, delta)

    urls_time_range = []
    for day, time in product(dates, time):
           file_name = f'narr-a_221_{day.strftime("%Y%m%d")}_{time}_000.grb'
           url = URL(base_url, day.strftime('%Y%m'), day.strftime('%Y%m%d'))
           urls_time_range.append(str(URL(url, file_name)))

    with multiprocessing.Pool(max_workers) as p:
        results = p.map(requests_to_s3, urls_time_range, chunksize=1)

        print('Finish download')
        buf = io.BytesIO()
        with zipfile.ZipFile(buf, mode='w', compression=zipfile.ZIP_DEFLATED, compresslevel=1) as zf:
            for content_file_name, content_file_result in results:
                try:
                    zf.writestr(content_file_name,
                                content_file_result)
                except Exception as exc:
                    print(exc)

        print('Finish zipping  - Upload Start')
        with smart_open.s3.open(aws_bucket_name, key, 'wb', session=session) as so:
            so.write(buf.getvalue())

    return None
```

You can test the function by running: 

```python
from datetime import datetime

a = stream_time_range_s3(start_date=datetime(2012, 1, 1),
end_date=datetime(2012, 2, 1),
aws_key=aws_key,
delta=None,
aws_secret=aws_secret,
aws_bucket_name=bucket_name,
key='wind_2012_test_parts.zip',
max_workers=10)

```

#### Versions

```
Darwin-18.7.0-x86_64-i386-64bit
Python 3.7.1 (default, Feb 27 2019, 18:57:54)
[Clang 10.0.0 (clang-1000.10.44.4)]
smart_open 1.8.4
```



Discussion:
Thanks for reporting this.

Could you simplify the example a bit? It's too long and requires external data (wind_2012_test_parts.zip).

From your description of the problem, it sounds like the following should reproduce your problem:

```python
import io
import smart_open
with open('some_large_file.bin', 'rb') as fin:
    with smart_open.open('s3://bucket/key.bin', 'wb') as fout:
        buf = fin.read(10e9) # read 10GiB into memory, oof
        fout.write(buf)
```

Can you confirm whether the above reproduces your problem? If not, let's look into reducing your original example, it's a bit too much for me to look at.
I can confirm that I'm encountering this error when trying to upload a file over 5GB via `fout.write(buf)` as you've stated in the simplified example.

This stack overflow article appears to explain the cause: https://stackoverflow.com/questions/26319815/entitytoolarge-error-when-uploading-a-5g-file-to-amazon-s3
@davidparks21 Thank you for confirming the problem.

I think we can resolve the issue by ensuring that a single write call never puts more than 5GB. If there is more data, then subsequent write calls should handle it.

Are you able to make a PR?
Oh, so just raise an exception when one `fout.write(buff)` is called with a `buff > 5GB`?

That would be an easy solution to deal with for me. I think I could do a PR for that. There was one other small thing I wanted to do a PR for too, so this would probably get me off my butt to do both.
> so just raise an exception when one fout.write(buff) is called with a buff > 5GB?

smart_open's promise is to handle large uploads (and downloads) transparently. So instead of raising an exception, isn't it better to split the chunk into multipart pieces, each smaller than 5GB?

IIRC smart_open is already handling multipart uploads transparently under the hood, so this should be no different.
I have a similar issue with trying to stream/write large files to S3 via `pickle.dump(obj, fout, protocol=4)` as smart_open is trying to upload parts that are each greater than 5GB.

Is this still a 'needs-info' or is the problem understood? 
I think we understand the problem, now we "just" need to fix it.
@JamalRahman @pythric @ivanhigueram could you help with a fix, prepare the PR?
@mpenkov is it still an open issue? I found this when considering smart_open to upload 5TB file to S3. 
@davidparks21 how did you resolve the problem in your case? Could you share some code snippet?
I started working on a solution, but it wasn't a very trivial change the way it's currently coded. I ran out of time and abandoned the effort back when I posted. I'm not sure about the current status, but my solution was to simply chunk the calls to `file.write(n)` which was pretty trivial in my use case.
I've been chugging away at this, and finally hit upon a solution where it wouldn't need to make any extra copies to the buffer of the data - which would be a significant improvement when dealing with the size of files we're talking about, but unfortunately https://github.com/boto/boto3/issues/3423 stopped me from such a perfect solution.

I'll be opening a PR soon with a compromise solution though, but if my PR to botocore is accepted and it's released it'll open up not needing to buffer the data at all before sending (unless smaller than min-size-upload writes are involved).



Patch:
diff --git a/smart_open/s3.py b/smart_open/s3.py
index 1be5ece6..d07498f9 100644
--- a/smart_open/s3.py
+++ b/smart_open/s3.py
@@ -6,6 +6,7 @@
 # from the MIT License (MIT).
 #
 """Implements file-like objects for reading and writing from/to AWS S3."""
+from __future__ import annotations
 
 import http
 import io
@@ -14,6 +15,12 @@
 import time
 import warnings
 
+from typing import (
+    Callable,
+    List,
+    TYPE_CHECKING,
+)
+
 try:
     import boto3
     import botocore.client
@@ -28,18 +35,34 @@
 
 from smart_open import constants
 
-from typing import (
-    Callable,
-    List,
-)
+
+if TYPE_CHECKING:
+    from mypy_boto3_s3.client import S3Client
+    from typing_extensions import Buffer
 
 logger = logging.getLogger(__name__)
 
-DEFAULT_MIN_PART_SIZE = 50 * 1024**2
-"""Default minimum part size for S3 multipart uploads"""
-MIN_MIN_PART_SIZE = 5 * 1024 ** 2
+#
+# AWS puts restrictions on the part size for multipart uploads.
+# Each part must be more than 5MB, and less than 5GB.
+#
+# On top of that, our MultipartWriter has a min_part_size option.
+# In retrospect, it's an unfortunate name, because it conflicts with the
+# minimum allowable part size (5MB), but it's too late to change it, because
+# people are using that parameter (unlike the MIN, DEFAULT, MAX constants).
+# It really just means "part size": as soon as you have this many bytes,
+# write a part to S3 (see the MultipartWriter.write method).
+#
+
+MIN_PART_SIZE = 5 * 1024 ** 2
 """The absolute minimum permitted by Amazon."""
 
+DEFAULT_PART_SIZE = 50 * 1024**2
+"""The default part size for S3 multipart uploads, chosen carefully by smart_open"""
+
+MAX_PART_SIZE = 5 * 1024 ** 3
+"""The absolute maximum permitted by Amazon."""
+
 SCHEMES = ("s3", "s3n", 's3u', "s3a")
 DEFAULT_PORT = 443
 DEFAULT_HOST = 's3.amazonaws.com'
@@ -286,7 +309,7 @@ def open(
     mode,
     version_id=None,
     buffer_size=DEFAULT_BUFFER_SIZE,
-    min_part_size=DEFAULT_MIN_PART_SIZE,
+    min_part_size=DEFAULT_PART_SIZE,
     multipart_upload=True,
     defer_seek=False,
     client=None,
@@ -306,12 +329,29 @@ def open(
     buffer_size: int, optional
         The buffer size to use when performing I/O.
     min_part_size: int, optional
-        The minimum part size for multipart uploads.  For writing only.
+        The minimum part size for multipart uploads, in bytes.
+
+        When the writebuffer contains this many bytes, smart_open will upload
+        the bytes to S3 as a single part of a multi-part upload, freeing the
+        buffer either partially or entirely.  When you close the writer, it
+        will assemble the parts together.
+
+        The value determines the upper limit for the writebuffer.  If buffer
+        space is short (e.g. you are buffering to memory), then use a smaller
+        value for min_part_size, or consider buffering to disk instead (see
+        the writebuffer option).
+
+        The value must be between 5MB and 5GB.  If you specify a value outside
+        of this range, smart_open will adjust it for you, because otherwise the
+        upload _will_ fail.
+
+        For writing only.  Does not apply if you set multipart_upload=False.
     multipart_upload: bool, optional
         Default: `True`
         If set to `True`, will use multipart upload for writing to S3. If set
         to `False`, S3 upload will use the S3 Single-Part Upload API, which
         is more ideal for small file sizes.
+
         For writing only.
     version_id: str, optional
         Version of the object, used when reading object.
@@ -358,10 +398,10 @@ def open(
             fileobj = MultipartWriter(
                 bucket_id,
                 key_id,
-                min_part_size=min_part_size,
                 client=client,
                 client_kwargs=client_kwargs,
                 writebuffer=writebuffer,
+                part_size=min_part_size,
             )
         else:
             fileobj = SinglepartWriter(
@@ -829,17 +869,21 @@ def __init__(
         self,
         bucket,
         key,
-        min_part_size=DEFAULT_MIN_PART_SIZE,
+        part_size=DEFAULT_PART_SIZE,
         client=None,
         client_kwargs=None,
-        writebuffer=None,
+        writebuffer: io.BytesIO | None = None,
     ):
-        if min_part_size < MIN_MIN_PART_SIZE:
-            logger.warning("S3 requires minimum part size >= 5MB; \
-multipart upload may fail")
-        self._min_part_size = min_part_size
+        adjusted_ps = smart_open.utils.clamp(part_size, MIN_PART_SIZE, MAX_PART_SIZE)
+        if part_size != adjusted_ps:
+            logger.warning(f"adjusting part_size from {part_size} to {adjusted_ps}")
+            part_size = adjusted_ps
+        self._part_size = part_size
 
         _initialize_boto3(self, client, client_kwargs, bucket, key)
+        self._client: S3Client
+        self._bucket: str
+        self._key: str
 
         try:
             partial = functools.partial(
@@ -862,12 +906,12 @@ def __init__(
 
         self._total_bytes = 0
         self._total_parts = 0
-        self._parts = []
+        self._parts: list[dict[str, object]] = []
 
         #
         # This member is part of the io.BufferedIOBase interface.
         #
-        self.raw = None
+        self.raw = None  # type: ignore[assignment]
 
     def flush(self):
         pass
@@ -943,7 +987,7 @@ def tell(self):
     def detach(self):
         raise io.UnsupportedOperation("detach() not supported")
 
-    def write(self, b):
+    def write(self, b: Buffer) -> int:
         """Write the given buffer (bytes, bytearray, memoryview or any buffer
         interface implementation) to the S3 file.
 
@@ -951,14 +995,30 @@ def write(self, b):
 
         There's buffering happening under the covers, so this may not actually
         do any HTTP transfer right away."""
+        offset = 0
+        mv = memoryview(b)
+        self._total_bytes += len(mv)
 
-        length = self._buf.write(b)
-        self._total_bytes += length
+        #
+        # botocore does not accept memoryview, otherwise we could've gotten
+        # away with not needing to write a copy to the buffer aside from cases
+        # where b is smaller than part_size
+        #
+        while offset < len(mv):
+            start = offset
+            end = offset + self._part_size - self._buf.tell()
+            self._buf.write(mv[start:end])
+            if self._buf.tell() < self._part_size:
+                #
+                # Not enough data to write a new part just yet. The assert
+                # ensures that we've consumed all of the input buffer.
+                #
+                assert end >= len(mv)
+                return len(mv)
 
-        if self._buf.tell() >= self._min_part_size:
             self._upload_next_part()
-
-        return length
+            offset = end
+        return len(mv)
 
     def terminate(self):
         """Cancel the underlying multipart upload."""
@@ -984,7 +1044,7 @@ def to_boto3(self, resource):
     #
     # Internal methods.
     #
-    def _upload_next_part(self):
+    def _upload_next_part(self) -> None:
         part_num = self._total_parts + 1
         logger.info(
             "%s: uploading part_num: %i, %i bytes (total %.3fGB)",
@@ -1033,10 +1093,10 @@ def __str__(self):
         return "smart_open.s3.MultipartWriter(%r, %r)" % (self._bucket, self._key)
 
     def __repr__(self):
-        return "smart_open.s3.MultipartWriter(bucket=%r, key=%r, min_part_size=%r)" % (
+        return "smart_open.s3.MultipartWriter(bucket=%r, key=%r, part_size=%r)" % (
             self._bucket,
             self._key,
-            self._min_part_size,
+            self._part_size,
         )
 
 


Test patch:
diff --git a/smart_open/tests/test_s3.py b/smart_open/tests/test_s3.py
index 430c29a6..6f12dbdd 100644
--- a/smart_open/tests/test_s3.py
+++ b/smart_open/tests/test_s3.py
@@ -200,8 +200,8 @@ def test_every_second_read_fails(self):
 class ReaderTest(BaseTest):
     def setUp(self):
         # lower the multipart upload size, to speed up these tests
-        self.old_min_part_size = smart_open.s3.DEFAULT_MIN_PART_SIZE
-        smart_open.s3.DEFAULT_MIN_PART_SIZE = 5 * 1024**2
+        self.old_min_part_size = smart_open.s3.DEFAULT_PART_SIZE
+        smart_open.s3.DEFAULT_PART_SIZE = 5 * 1024**2
 
         ignore_resource_warnings()
 
@@ -214,7 +214,7 @@ def setUp(self):
         s3.Object(BUCKET_NAME, KEY_NAME).put(Body=self.body)
 
     def tearDown(self):
-        smart_open.s3.DEFAULT_MIN_PART_SIZE = self.old_min_part_size
+        smart_open.s3.DEFAULT_PART_SIZE = self.old_min_part_size
 
     def test_iter(self):
         """Are S3 files iterated over correctly?"""
@@ -461,27 +461,64 @@ def test_write_02(self):
             fout.write(u"testžížáč".encode("utf-8"))
             self.assertEqual(fout.tell(), 14)
 
+    #
+    # Nb. Under Windows, the byte offsets are different for some reason
+    #
+    @pytest.mark.skipif(condition=sys.platform == 'win32', reason="does not run on windows")
     def test_write_03(self):
         """Does s3 multipart chunking work correctly?"""
-        # write
-        smart_open_write = smart_open.s3.MultipartWriter(
-            BUCKET_NAME, WRITE_KEY_NAME, min_part_size=10
-        )
-        with smart_open_write as fout:
-            fout.write(b"test")
-            self.assertEqual(fout._buf.tell(), 4)
 
-            fout.write(b"test\n")
-            self.assertEqual(fout._buf.tell(), 9)
-            self.assertEqual(fout._total_parts, 0)
+        #
+        # generate enough test data for a single multipart upload part.
+        # We need this because moto behaves like S3; it refuses to upload
+        # parts smaller than 5MB.
+        #
+        data_dir = os.path.join(os.path.dirname(__file__), "test_data")
+        with open(os.path.join(data_dir, "crime-and-punishment.txt"), "rb") as fin:
+            crime = fin.read()
+        data = b''
+        ps = 5 * 1024 * 1024
+        while len(data) < ps:
+            data += crime
+
+        title = "Преступление и наказание\n\n".encode()
+        to_be_continued = "\n\n... продолжение следует ...\n\n".encode()
+
+        with smart_open.s3.MultipartWriter(BUCKET_NAME, WRITE_KEY_NAME, part_size=ps) as fout:
+            #
+            # Write some data without triggering an upload
+            #
+            fout.write(title)
+            assert fout._total_parts == 0
+            assert fout._buf.tell() == 48
 
-            fout.write(b"test")
-            self.assertEqual(fout._buf.tell(), 0)
-            self.assertEqual(fout._total_parts, 1)
+            #
+            # Trigger a part upload
+            #
+            fout.write(data)
+            assert fout._total_parts == 1
+            assert fout._buf.tell() == 661
 
+            #
+            # Write _without_ triggering a part upload
+            #
+            fout.write(to_be_continued)
+            assert fout._total_parts == 1
+            assert fout._buf.tell() == 710
+
+        #
+        # We closed the writer, so the final part must have been uploaded
+        #
+        assert fout._buf.tell() == 0
+        assert fout._total_parts == 2
+
+        #
         # read back the same key and check its content
-        output = list(smart_open.s3.open(BUCKET_NAME, WRITE_KEY_NAME, 'rb'))
-        self.assertEqual(output, [b"testtest\n", b"test"])
+        #
+        with smart_open.s3.open(BUCKET_NAME, WRITE_KEY_NAME, 'rb') as fin:
+            got = fin.read()
+        want = title + data + to_be_continued
+        assert want == got
 
     def test_write_04(self):
         """Does writing no data cause key with an empty value to be created?"""
@@ -586,7 +623,6 @@ def test_write_gz_with_error(self):
             ) as fout:
                 fout.write(b"test12345678test12345678")
                 fout.write(b"test\n")
-
                 raise ValueError("some error")
 
         # no multipart upload was committed: