fix(api): exit worker on memory allocation errors

ssube · Mar 6, 2023 · 57fed94 · 57fed94
1 parent cb460a0
commit 57fed94
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 6 deletions.
diff --git a/api/onnx_web/worker/worker.py b/api/onnx_web/worker/worker.py
@@ -39,10 +39,14 @@ def worker_main(context: WorkerContext, server: ServerContext):
             logger.info("worker got keyboard interrupt")
             exit(0)
         except ValueError as e:
-            logger.info("value error in worker: %s", e)
+            logger.info("value error in worker, exiting: %s", e)
             exit(1)
         except Exception as e:
-            logger.error(
-                "error while running job: %s",
-                format_exception(type(e), e, e.__traceback__),
-            )
+            if "Failed to allocate memory" in str(e):
+                logger.error("detected out-of-memory error, exiting: %s", e)
+                exit(2)
+            else:
+                logger.error(
+                    "error while running job: %s",
+                    format_exception(type(e), e, e.__traceback__),
+                )
diff --git a/api/scripts/test-memory-leak.sh b/api/scripts/test-memory-leak.sh
@@ -6,7 +6,7 @@ do
   curl "${test_host}/api/txt2img?"\
 'cfg=16.00&steps=30&scheduler=ddim&seed=-1&'\
 '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&'\
 'lpw=false&width=512&height=512&upscaleOrder=correction-both' \
     -X 'POST' \