github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Rework windows start and stop methods so that command windows dont appear. Stop with signal since we cannot signint.

Browse Source
This commit is contained in:
Jaret Burkett
2026-07-27 17:43:01 -06:00
parent 6b0449c326
commit 461e798708
8 changed files with 389 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
extensions_built_in/sd_trainer/DiffusionTrainer.py
View File

@@ -39,7 +39,12 @@ class DiffusionTrainer(SDTrainer):
# Initialize the status
self._run_async_operation(self._update_status("running", "Starting"))
self._stop_watcher_started = False
# self.start_stop_watcher(interval_sec=2.0)
if os.name == "nt":
# On Windows the stop route cannot send us SIGINT from outside
# (no console to deliver a Ctrl+C to), so watch the stop flag
# and raise the interrupt from inside. On Linux the route
# sends a real SIGINT to the pid and this is unnecessary.
self.start_stop_watcher(interval_sec=2.0)
def start_stop_watcher(self, interval_sec: float = 5.0):
"""
@@ -60,26 +65,21 @@ class DiffusionTrainer(SDTrainer):
while True:
try:
if self.should_stop():
# Mark and update status (non-blocking; uses existing infra)
self.is_stopping = True
self._run_async_operation(
self._update_status("stopped", "Job stopped (remote)")
)
# Best-effort flush pending async ops
try:
asyncio.run(self.wait_for_all_async())
except RuntimeError:
pass
# Try to stop DB thread pool quickly
try:
self.thread_pool.shutdown(wait=False, cancel_futures=True)
except TypeError:
self.thread_pool.shutdown(wait=False)
if self.is_stopping:
# maybe_stop() already started the graceful shutdown;
# a second interrupt would only break its cleanup.
return
print("")
print("****************************************************")
print(" Stop signal received; terminating process. ")
print("****************************************************")
os.kill(os.getpid(), signal.SIGINT)
# Deliver a real KeyboardInterrupt to the main thread so
# on_error runs the normal shutdown (final DB write, last
# log). os.kill(pid, SIGINT) must not be used here: on
# Windows it is TerminateProcess and kills us instantly.
# Leave the thread pool alone -- on_error still needs it.
signal.raise_signal(signal.SIGINT)
return
time.sleep(interval_sec)
except Exception:
time.sleep(interval_sec)