lint(api): explicitly bind the device pool to shutdown callback
This commit is contained in:
parent
c2f8fb1d31
commit
2bbc5d877f
|
@ -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
|
||||||
|
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue