put_async resulting in OOM

[aabed-aa]

I am working on a project which is using Obstore in combination with FastAPI. I have implemented an endpoint to perform file uploads, what I noticed is that bigger files cause OOM errors due to put_async. Here are some details:

I have an endpoint (redacted for brevity) that looks like this:

@router.post(
    "/upload",
)
async def upload(
    request: Request,
):
     # Stream the file to S3
     store = S3Store(...)
     await obstore.put_async(store, some_path, request.stream())

Since I am using request.stream() as the file I am expecting the put operation to directly stream the data to the S3 bucket. What is actually happening is that my application is crashing with an OOM error. The application is running in a pod with a limit of 2gb memory., the file I am uploading is 3.5gb. Increasing the limit solves the issue but this means multiple uploads to this service will still cause an OOM. For sanity I iterated over the request.stream() without put_async and this did not affect memory usage.

Any advice how to perform a streaming upload using Obstore?

Obstore version: 0.8.0
Python version: 3.13.5

[kylebarron]

Hmm.

put_async does use streaming. In particular, it'll create a Multipart upload with each iteration yielded from the stream.

---

When you pass an async iterable of bytes in, it becomes an AsyncPushSource,

obstore/obstore/src/put.rs

Lines 147 to 154 in 1253f8d

	/// Sources to `put` that are push-based and asynchronous.
	///
	/// I.e. we don't know how large each chunk will be before we receive it.
	pub(crate) enum AsyncPushSource {
	/// A Python Async Iterator: An object with an __anext__ method that returns a buffer protocol
	/// object (anything that can be extracted into `PyBytes`)
	AsyncIterator(Py<PyAny>),
	}

Then in push_async we always infer a multipart upload:

obstore/obstore/src/put.rs

Lines 366 to 370 in 1253f8d

	let mut use_multipart = if let Some(use_multipart) = use_multipart {
	use_multipart
	} else {
	file.use_multipart(chunk_size)?
	};

obstore/obstore/src/put.rs

Lines 214 to 216 in 1253f8d

	// We always use multipart uploads for push-based sources because we have no way of
	// knowing how large they'll be and we don't want to buffer them into memory.
	_ => Ok(true),

Then that passes to put_multipart_inner

obstore/obstore/src/put.rs

Lines 381 to 390 in 1253f8d

	put_multipart_inner(
	store.into_inner(),
	path.as_ref(),
	file,
	chunk_size,
	max_concurrency,
	attributes,
	tags,
	)
	.await?

That iterates over __anext__ one at a time

obstore/obstore/src/put.rs

Lines 485 to 491 in 1253f8d

	PutInput::AsyncPush(mut push_reader) => {
	// Note: I believe that only one __anext__ call can happen at a time
	while let Some(buf) = push_reader.next_chunk().await? {
	writer.wait_for_capacity(max_concurrency).await?;
	writer.put(buf);
	}
	}

---

One thing you can try is reducing the max_concurrency and chunk_size.

Another thing I wonder about is whether there's an issue here when download speed is much higher than upload speed. Because the download is push-based, perhaps it buffers more in memory than desired while waiting for the upload to happen?

If you have a single upload concurrency, you should be able to verify easily through network tools that the upload has started before the download has finished.

[aabed-aa]

Thank you for the clear explanation!

Regarding the download vs upload speed I am not able to reproduce this I performed the following test:

@router.post(
    "/upload",
)
async def upload(
    request: Request,
):
     # "slowly" iterate over the stream
     async for x in request.stream():
          print(len(x))
          await asyncio.sleep(0.05)

Running this did not increase the memory usage but it did take some time to finish the request.

I performed another test:

@router.post(
    "/upload",
)
async def upload(
    request: Request,
):
     # "slowly" iterate over the stream
    async def stream():
        async for x in request.stream():
            await asyncio.sleep(0.05)
            yield x

     store = S3Store(...)
     await obstore.put_async(store, some_path, stream())

This caused the memory usage to slowly increase.. I am not sure about obstore internals but for some reason consuming the stream from put_async is causing memory usage to grow.

[kylebarron]

I don't think putting an asyncio.sleep in there would do much of anything.

Would you be able to run your code with memray and share the flamegraph?

[kylebarron]

I implemented a quick little bench script and couldn't see any sort of memory leak with memray:

memray-flamegraph-upload.py.68787.html

import asyncio
from pathlib import Path

from boto3.session import Session

from obstore.auth.boto3 import Boto3CredentialProvider
from obstore.store import LocalStore, S3Store

path = Path("~/Downloads").expanduser()
local_store = LocalStore(path)
session = Session(profile_name="ds")
credential_provider = Boto3CredentialProvider(session=session)


async def upload():
    resp = await local_store.get_async("population_areas.fgb")

    remote_store = S3Store(
        "ds-wheels",
        credential_provider=credential_provider,
        region="us-east-1",
    )
    await remote_store.put_async("population_areas.fgb", resp)


asyncio.run(upload())

That script uploads this 11GB file to S3.

[aabed-aa]

I have found the issue:

I was initially using:

 await obstore.put_async(..., mode="create") 

It seems that if the file exists and I use "create" this issue arrises.. Obstore is not raising an exception if the file already exists, and the memory usage is going up (I guess because of retries in the background of the parts).

 await obstore.put_async(..., mode="overwrite") 

This solves the issue..

Is there any way to enable logging of obstore? I have not found anything related to this in the docs, the logging might show this happening.

[kylebarron]

Is there any way to enable logging of obstore? I have not found anything related to this in the docs, the logging might show this happening.

No, that's tracked in #93

[kylebarron]

I prototyped logging in #544 if you'd like to test it out

[aabed-aa]

I prototyped logging in #544 if you'd like to test it out

Sure I'll give this a go!

[kylebarron]

It seems that if the file exists and I use "create" this issue arrises.. Obstore is not raising an exception if the file already exists, and the memory usage is going up (I guess because of retries in the background of the parts).

It looks like it should raise an error if the file already exists.

This test verifies that it raises an AlreadyExistsError when using mode="create" when a file already exists there:

obstore/tests/test_put.py

Lines 58 to 64 in aaa9261

	store = MemoryStore()
	store.put("file1.txt", b"foo")
	store.put("file1.txt", b"bar", mode="overwrite")
	with pytest.raises(AlreadyExistsError):
	store.put("file1.txt", b"foo", mode="create")