Serving ML models with multiple workers linearly adds the RAM's load.

Recently, we deployed a ML model with FastAPI, and encountered an issue.

The code looks like this.

from ocr_pipeline.model.ocr_wrapper import OcrWrapper
ocr_wrapper = OcrWrapper(**config.model_load_params) # loads 1.5 GB PyTorch model

...

@api.post('/')
async def predict(file: UploadFile = File(...)):
       preds = ocr_wrapper.predict(file.file, **config.model_predict_params)
       return json.dumps({"data": preds})

The above written command consumes min. 3GB of RAM.

gunicorn --workers 2 --worker-class=uvicorn.workers.UvicornWorker app.main:api

Is there any way to scale the number of workers without consuming too much RAM?

ENVIRONMENT: Ubuntu 18.04 Python 3.6.9

fastapi==0.61.2 uvicorn==0.12.2 gunicorn==20.0.4 uvloop==0.14.0

@tiangolo

[raphaelauv]

This is not a specific fastAPI question (more a gunicorn one) , it’s about sharing memory between process

The solution would be loading the model in ram before the fork of the workers (of gunicorn)

so you need to use –preload

gunicorn --workers 2 --preload --worker-class=uvicorn.workers.UvicornWorker app.main:api

your main.py file inside folder app

def create_app():
    MY_MODEL.load("model_path")
    app = FastAPI()
    app.include_router(my_router)
    return app
api = create_app()

If you have more question about gunicorn or python or fork or copy-on-write or python reference counting or memory leak -> stackoverflow

YOU can very probably CLOSE this issue , thank you 😃

[cosimo]

Just found out that if I change my app methods from:

@app.post("/clusters", response_model=ClusteringResponse)
async def cluster(request: ClusteringRequest, model=Depends(get_model)):
    """Cluster a list of text sentences"""
    ...

to:

@app.post("/clusters", response_model=ClusteringResponse)
def cluster(request: ClusteringRequest, model=Depends(get_model)):
    """Cluster a list of text sentences"""
    ...

removing the async qualifier, the model does indeed work as expected.

@sevakharutyunyan are you able to verify if this works for you?