I am using fastapi to do predictions on a ml model. Because a single prediction is taking more than 2 minutes I have made the prediction to run as a background task using fastapi background tasks. Now when i deployed to digital ocean i was unable to run background worker. I am getting 504: Gateway Timeout Error.

startup command

gunicorn --worker-tmp-dir /dev/shm --config gunicorn.config.py main:app

gunicorn.config.py

bind = "0.0.0.0:8080"
workers = 4
worker_class = "uvicorn.workers.UvicornWorker"

main.py

response = {}
async def predictions(solute, solvent):
    mol = Chem.MolFromSmiles(solute)
    solute_graph = get_graph_from_smile(solute)
    mol = Chem.MolFromSmiles(solvent)
    solvent_graph = get_graph_from_smile(solvent)
    delta_g, interaction_map =  model([solute_graph.to(device), solvent_graph.to(device)])
    interaction_map_one = torch.trunc(interaction_map)
    response["interaction_map"] = (interaction_map_one.detach().numpy()).tolist()
    response["predictions"] = delta_g.item()

@app.post('/predict_solubility')
async def post():
    return {'result': response}

@app.post('/predict')
async def predict(background_tasks: BackgroundTasks,solute,solvent):
    background_tasks.add_task(predictions,solute,solvent)
    return {'success'}

response_two = {}
async def predictions_two(solute):
    for i in data:
        delta_g, interaction_map =  model([get_graph_from_smile(Chem.MolToSmiles(Chem.AddHs(Chem.MolFromSmiles(solute)))).to(device), get_graph_from_smile(Chem.MolToSmiles(Chem.AddHs(Chem.MolFromSmiles(i)))).to(device)])
        response_two[i] = delta_g.item()
    
@app.post('/predict_solubility_json')
async def post():
    return {'result': response_two}

@app.get('/predict_two')
async def predict_two(background_tasks: BackgroundTasks,solute):
    background_tasks.add_task(predictions_two,solute)
    return {'success'}

if __name__ == "__main__":
    uvicorn.run(app,host="0.0.0.0", port=8000, reload=True)

I have added worker in digital ocean from components even though i am getting still server timeout. I am not sure whether i have not implemented background tasks using fastapi for is their any problem with my startup command