apextoaster
/
onnx-web
1
0
Fork 0
Code Actions

fix(api): handle CUDA memory errors in workers

This commit is contained in:
Sean Sube 2023-03-15 08:51:29 -05:00
parent a2e21f427f
commit b2eb406197
Signed by: ssube
GPG Key ID: 3EED7B957D362AF1
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -60,7 +60,7 @@ 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): if "Failed to allocate memory" in str(e) or "CUDA out of memory" in str(e):
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: