apextoaster
/
onnx-web
1
0
Fork 0
Code Actions

fix(api): handle more out-of-memory errors in the workers

This commit is contained in:
Sean Sube 2023-03-16 18:34:28 -05:00
parent dfba5f0231
commit c8c5e9f42e
Signed by: ssube
GPG Key ID: 3EED7B957D362AF1
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
api/onnx_web/worker/worker.py
View File

@ -60,7 +60,8 @@ def worker_main(context: WorkerContext, server: ServerContext):
) )
exit(EXIT_ERROR) exit(EXIT_ERROR)
except Exception as e: except Exception as e:
if "Failed to allocate memory" in str(e) or "CUDA out of memory" in str(e): e_str = str(e)
if "Failed to allocate memory" in e_str or "out of memory" in e_str:
logger.error("detected out-of-memory error, exiting: %s", e) logger.error("detected out-of-memory error, exiting: %s", e)
exit(EXIT_MEMORY) exit(EXIT_MEMORY)
else: else: