apextoaster
/
onnx-web
1
0
Fork 0
Code Actions

lint(api): explicitly bind the device pool to shutdown callback

This commit is contained in:
Sean Sube 2023-03-27 17:13:45 -05:00
parent c2f8fb1d31
commit 2bbc5d877f
Signed by: ssube
GPG Key ID: 3EED7B957D362AF1
2 changed files with 5 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
api/onnx_web/main.py
View File

@ -1,5 +1,6 @@
import atexit import atexit
import gc import gc
from functools import partial
from logging import getLogger from logging import getLogger
from diffusers.utils.logging import disable_progress_bar from diffusers.utils.logging import disable_progress_bar
@ -67,11 +68,11 @@ def run():
app, pool = main() app, pool = main()
pool.start() pool.start()
def quit(): def quit(p: DevicePoolExecutor):
logger.info("shutting down workers") logger.info("shutting down workers")
pool.join() p.join()
atexit.register(quit) atexit.register(partial(quit, pool))
return app return app