I am using FastAPI to perform 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's BackgroundTasks
. 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