apextoaster
/
onnx-web
1
0
Fork 0
Code Actions

fix(api): exit worker on memory allocation errors

This commit is contained in:
Sean Sube 2023-03-05 21:07:06 -06:00
parent cb460a0c59
commit 57fed94337
Signed by: ssube
GPG Key ID: 3EED7B957D362AF1
2 changed files with 10 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -39,9 +39,13 @@ def worker_main(context: WorkerContext, server: ServerContext):
logger.info("worker got keyboard interrupt") logger.info("worker got keyboard interrupt")
exit(0) exit(0)
except ValueError as e: except ValueError as e:
logger.info("value error in worker: %s", e) logger.info("value error in worker, exiting: %s", e)
exit(1) exit(1)
except Exception as e: except Exception as e:
if "Failed to allocate memory" in str(e):
logger.error("detected out-of-memory error, exiting: %s", e)
exit(2)
else:
logger.error( logger.error(
"error while running job: %s", "error while running job: %s",
format_exception(type(e), e, e.__traceback__), format_exception(type(e), e, e.__traceback__),

2
api/scripts/test-memory-leak.sh
View File

@ -6,7 +6,7 @@ do
curl "${test_host}/api/txt2img?"\ curl "${test_host}/api/txt2img?"\
'cfg=16.00&steps=30&scheduler=ddim&seed=-1&'\ 'cfg=16.00&steps=30&scheduler=ddim&seed=-1&'\
'prompt=an+astronaut+eating+a+hamburger&negativePrompt=&'\ 'prompt=an+astronaut+eating+a+hamburger&negativePrompt=&'\
'model=stable-diffusion-onnx-v1-5&platform=any&'\ 'model=stable-diffusion-onnx-v1-5&platform=cuda&'\
'upscaling=upscaling-real-esrgan-x2-plus&correction=correction-codeformer&'\ 'upscaling=upscaling-real-esrgan-x2-plus&correction=correction-codeformer&'\
'lpw=false&width=512&height=512&upscaleOrder=correction-both' \ 'lpw=false&width=512&height=512&upscaleOrder=correction-both' \
-X 'POST' \ -X 'POST' \