Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

catch ^C at startup time while models are being scanned

Browse Source
This commit is contained in:
Lincoln Stein 2024-03-21 23:50:52 -04:00 committed by psychedelicious
parent 2ecbb9f720
commit 34f5259980
2 changed files with 24 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
invokeai/app/services/download/download_default.py
View File

@ -85,7 +85,7 @@ class DownloadQueueService(DownloadQueueServiceBase):
self._logger.info(f"Waiting for {len(active_jobs)} active download jobs to complete")
with self._queue.mutex:
self._queue.queue.clear()
self.join() # wait for all active jobs to finish
self.cancel_all_jobs()
self._stop_event.set()
for thread in self._worker_pool:
thread.join()

32
invokeai/app/services/model_install/model_install_default.py
View File

@ -2,6 +2,7 @@
import os
import re
import signal
import threading
import time
from hashlib import sha256
@ -112,6 +113,18 @@ class ModelInstallService(ModelInstallServiceBase):
# makes the installer harder to use outside the web app
def start(self, invoker: Optional[Invoker] = None) -> None:
"""Start the installer thread."""
# Yes, this is weird. When the installer thread is running, the
# thread masks the ^C signal. When we receive a
# sigINT, we stop the thread, reset sigINT, and send a new
# sigINT to the parent process.
def sigint_handler(signum, frame):
self.stop()
signal.signal(signal.SIGINT, signal.SIG_DFL)
signal.raise_signal(signal.SIGINT)
signal.signal(signal.SIGINT, sigint_handler)
with self._lock:
if self._running:
raise Exception("Attempt to start the installer service twice")
@ -122,15 +135,16 @@ class ModelInstallService(ModelInstallServiceBase):
def stop(self, invoker: Optional[Invoker] = None) -> None:
"""Stop the installer thread; after this the object can be deleted and garbage collected."""
with self._lock:
if not self._running:
raise Exception("Attempt to stop the install service before it was started")
self._stop_event.set()
self._clear_pending_jobs()
self._download_cache.clear()
assert self._install_thread is not None
self._install_thread.join()
self._running = False
# do manual acquire() so as to receive signals
if not self._running:
raise Exception("Attempt to stop the install service before it was started")
self._logger.debug("calling stop_event.set()")
self._stop_event.set()
self._clear_pending_jobs()
self._download_cache.clear()
assert self._install_thread is not None
self._install_thread.join()
self._running = False
def _clear_pending_jobs(self) -> None:
for job in self.list_jobs():