Unnecessarily large memory footprint

[atzannes]

Describe the bug

I have been trying to debug a memory pressure bug which has lead me to aiohttp. I have loosely traced the issue to the fact that stream.py uses the pattern of appending byte chunks to a list, then returning b"".join(chunks). When I replace this pattern with using a buffer = io.BytesIO() that I append to and return buffer.getvalue(), the unnecessary memory pressure goes away.

To reproduce the issue, I am downloading a large file (~500MB in this example) from S3 using UPath (which uses aiohttp internally). Without my proposed change, my high-watermark memory footprint is about twice the size of the large file, and when I get rid of the reference to the data it remains at about the size of the file. With my proposed change of using io.BytesIO() and avoiding b"".join(chunks), the high-watermark memory footprint is about the size of the file and when I drop the reference to it, the memory footprint goes back down to almost nothing.

For determining the memory footprint, I am using psutil and looking at the rss.

I am happy to try to dig deeper if needed or if you think that replacing the joining of bytes by an io.BytesIO buffer is not desirable, but given that I have a simple repro and a simple fix, I thought I would report the issue to open up the discussion. I'm happy to open a RP with my proposed solution.

To Reproduce

import psutil
from upath import UPath

def mem(when: str = ""):
    print(f"Resident Memory{when}: {psutil.Process().memory_info().rss / 1024**3:.2f} GB")

def upath_download(path) -> bytes:
    up = UPath(path)
    with up.open('rb') as fp:
        return fp.read()


path = 's3://my-bucket/my-500mb-file.bin'  # ~500MB

mem(" Initially")
data = upath_download(path)
print("data size", len(data))
mem(" After download")
del data
mem(" After delete")

Below is the output with the version 3.11.13 of aiohttp

Resident Memory Initially: 0.02 GB
data size 518458228
Resident Memory After download: 1.12 GB
Resident Memory After delete: 0.64 GB

Expected behavior

When I modify aiohttp to use my proposed solution with an io.BytesIO() buffer, I see the memory footprint grow to the size of the file when I hold a reference to it, and then going back down to almost nothing when I release the reference to the data, which is what I would expect.

Resident Memory Initially: 0.02 GB
data size 518458228
Resident Memory After download: 0.53 GB
Resident Memory After delete: 0.05 GB

Logs/tracebacks

Already listed above

Python Version

$ python --version
Python 3.10.8

aiohttp Version

$ python -m pip show aiohttp
Name: aiohttp
Version: 3.11.13
Summary: Async http client/server framework (asyncio)
Home-page: https://github.com/aio-libs/aiohttp
Author: 
Author-email: 
License: Apache-2.0
Location: /home/atzannes/code/github/chunky/.venv/lib/python3.10/site-packages
Requires: aiohappyeyeballs, aiosignal, async-timeout, attrs, frozenlist, multidict, propcache, yarl
Required-by: aiobotocore, aiohttp-cors, s3fs

multidict Version

$ python -m pip show multidict
Name: multidict
Version: 6.1.0
Summary: multidict implementation
Home-page: https://github.com/aio-libs/multidict
Author: Andrew Svetlov
Author-email: andrew.svetlov@gmail.com
License: Apache 2
Location: /home/atzannes/code/github/chunky/.venv/lib/python3.10/site-packages
Requires: typing-extensions
Required-by: aiobotocore, aiohttp, yarl

propcache Version

$ python -m pip show propcache

yarl Version

$ python -m pip show yarl

OS

Unnecessarily large memory footprint

Related component

Client

Additional context

No response

Code of Conduct

• I agree to follow the aio-libs Code of Conduct

[bdraco]

It's a bit unclear from your report how you are using aiohttp as your reproduction code doesn't include aiohttp.

Can you provide a working reproducer that uses aiohttp?

[Dreamsorcerer]

When I replace this pattern with using a buffer = io.BytesIO() that I append to and return buffer.getvalue(), the unnecessary memory pressure goes away.

I'm pretty sure I recall this being something we can't do as it would break too many things, though I could be thinking of something else.. If you run the test suite with your changes (or create a draft PR and let the CI run it), then it'll likely show you why it won't work. If the tests are still passing, we can evaluate the change properly (and as mentioned above, we'd want a proper test).

[atzannes]

@bdraco admittedly, I'm not using aiohttp directly but rather indirectly through UPath -> fsspec -> s3fs -> aiohttp. Given that fsspec/s3fs are fairly common I would assume that they are using aiohttp properly but that may not be the case. I wanted to see if this solution would even be acceptable, so I'll open a provisional PR to see if any tests break and if they don't you can tell me what the next steps are. Perhaps like you said to have a repro that doesn't rely on other libraries.

[bdraco]

I'm pretty sure the issue is you are reading the whole file in at once

def upath_download(path) -> bytes:
    up = UPath(path)
    with up.open('rb') as fp:
        return fp.read()

If you read in chunks I expect the issue will go away

[atzannes]

I'm pretty sure the issue is you are reading the whole file in at once

def upath_download(path) -> bytes:
    up = UPath(path)
    with up.open('rb') as fp:
        return fp.read()

If you read in chunks I expect the issue will go away

I agree that if I were to read the file in chunks, the resulting bloat in memory use would be minimal. That being said, I am writing library code, not application code and part of the API is to hand the user a stream-like object on which they are allowed to call read(). Right now, that object comes from s3fs and uses aiohttp. Of course I could override the implementation of read() on that object to do it in chunks behind the scenes to work around this issue that comes from aiohttp but if I can help to fix the root cause in aiohttp I think could all win depending on what the solution looks like. I would think that aiohttp being that it is a library that is broadly used, would be interested in avoiding pitfalls like memory bloat whenever practical, even in cases where the user is not very sophisticated.

I see that the provisional PR I opened created a terrible performance regression. I'll see if there is a way to address the performance so that there is no regression and also not have this memory bloat. Performance here is clearly much more important.

[atzannes]

I made an iota of progress in that I have a repro that only hits aiobotocore & aiohttp. I also ran a test where I am trying the b''.join(chunks) vs. the io.BytesIO() patterns against bytes objects generated in memory and I don't see a memory bloat in either case nor a performance degradation with io.BytesIO(). I was trying to see if there is something inherent to Python with these two patterns, which apparently there isn't. Perhaps not very insightful, just trying to cover all my bases.

import asyncio
from aiobotocore.session import get_session
import psutil


def mem(when: str = ""):
    print(f"Resident Memory{when}: {psutil.Process().memory_info().rss / 1024**3:.2f} GB")


bucket = 'my-bucket'
key = 'path/to/my/500MB/object'

async def go():
    session = get_session()
    async with session.create_client( 's3', region_name='us-east-1') as client:
        # get object from s3
        response = await client.get_object(Bucket=bucket, Key=key)
        # this will ensure the connection is correctly re-used/closed
        async with response['Body'] as stream:
            return await stream.read()


mem(" Initially")
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
data = loop.run_until_complete(go())
print("data size", len(data))
mem(" After download")
del data
mem(" After delete")

With io.BytesIO():

Resident Memory Initially: 0.04 GB
data size 518458228
Resident Memory After download: 0.53 GB
Resident Memory After delete: 0.05 GB

With current aiobotocore [==2.21.1]

Resident Memory Initially: 0.04 GB
data size 518458228
Resident Memory After download: 0.98 GB
Resident Memory After delete: 0.50 GB

pip list | grep aio:

aiobotocore               2.21.1
aiohappyeyeballs          2.5.0
aiohttp                   3.11.13
aiohttp-cors              0.7.0
aioitertools              0.12.0
aiosignal                 1.3.2