[Python] Why does my FastAPI application memory usage grow heavily when streaming multiple...

I'm streaming multiple files as a zip archive in a FastAPI endpoint using async_stream_zip from stream-zip (https://pypi.org/project/stream-zip/). Although this should be memory-efficient, I'm observing significant memory growth when streaming many files.

The API handles reading and converting of data (to JSON, CSV etc), and often this is CPU intense due to quirky formats and large files (context: celery is not appropriate for this API). Data is usually read via Pandas and converted using packages like orjson. I use ProcessPools to isolate the processing in async endpoints - as to avoid blocking the thread and better isolate and clean up memory. I have implemented the use of stream-zip (version 0.0.83) to yield one file at a time in a zipped folder. This should be memory efficient. But I see huge memory growth when streaming many larger files.

The application is deployed to a Kubernetes cluster via Docker with base image registry.access.redhat.com/ubi8/ubi, and ran using the following (gunicorn version 23.0.0, uvicorn-worker version 0.2.0):

CMD ["gunicorn", "-w", "3", "-k", "uvicorn.workers.UvicornWorker", "webservice.main:app", "--bind", "0.0.0.0:8000", "--max-requests", "500", "--max-requests-jitter", "100", "--access-logfile", "-"]

Now - here is a dummy example of the setup for better context, using FastAPI version 0.115.4:

from fastapi import FastAPI, Query
import asyncio
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
from datetime import datetime
from stream_zip import async_stream_zip, ZIP_32, stream_zip
from stat import S_IFREG
import orjson
from fastapi import APIRouter, FastAPI
from fastapi.responses import StreamingResponse

app = FastAPI(
title="test",
version="V1.0.0"
)


def generate_file(n=200000):
data_line = "2024-01-01 00:00:00 001 10000.00 2000.00 50000.00 51000.00\n"
return data_line * n

def parse(content):
data = []
for line in content.splitlines():
parts = line.split()
if len(parts) >= 7:
data.append({
"date": parts[0],
"time": parts[1],
"doy": int(parts[2]),
"a": float(parts[3]),
"b": float(parts[4]),
"c": float(parts[5]),
"d": float(parts[6]),
})
return orjson.dumps(data)

async def json_bytes_streamer(json_bytes: bytes, chunk_size: int = 8192):
for i in range(0, len(json_bytes), chunk_size):
yield json_bytes[i:i+chunk_size]

async def yield_data(num_files=50, pool='process'):

loop = asyncio.get_running_loop()
now = datetime.now()
content = generate_file()

executor_class = ProcessPoolExecutor if pool == 'process' else ThreadPoolExecutor
with executor_class(max_workers=1) as executor:
json_bytes = await loop.run_in_executor(executor, parse, content)
for i in range(num_files):
filename = f"data_{i+1:02d}.json"
yield (
filename,
now,
S_IFREG | 0o600,
ZIP_32,
json_bytes_streamer(json_bytes)
)

router = APIRouter(tags=["test"])

@router.get("/generate-zip-process-pool")
async def generate_zip_response_pp(num_files: int = Query(100, ge=1)):
return StreamingResponse(
async_stream_zip(yield_data(num_files=num_files)),
media_type="application/zip",
headers={"Content-Disposition": "attachment; filename=test_data.zip"},
)

app.include_router(router)




With one of my endpoints that yields hourly data in yearly files, memory exhausted per request can be as large as 50-100Mib, but local profiling shows more like 5-10Mib memory use (using psutils). For small requests (say, only 1 month of data), the issue is not really present - so are memory references being stored somewhere?

I do use asyncio tasks to handle streaming one file while another processes (I set up semaphore logic), but I don't see this being the root cause. Is this a Docker issue? The files read are from a GPFS mount if this provides any clues. I use Grafana to profile the deployment. I have tried setting strict global garbage collection settings and made use of del and call gc.collect() whenever I can. The memory growth is much more minimal using a ThreadPool, but then the memory is never released after (~10Mib per request) - again maybe a clue.

Possibly relevant - I do also check each file generated has valid data first using logic as:

async def check_data(agen):

try:
first_item = await agen.__anext__()
except StopAsyncIteration:
return None, None

async def new_gen():
yield first_item
async for item in agen:
yield item

return first_item, new_gen()


If anyone has any ideas at all please let me know - I've tried everything I can think of. I might be missing something super obvious and I am trying to learn more!

Continue reading...