diff --git a/extensions_built_in/captioner/BaseCaptioner.py b/extensions_built_in/captioner/BaseCaptioner.py index 6100fbb..3c8f298 100644 --- a/extensions_built_in/captioner/BaseCaptioner.py +++ b/extensions_built_in/captioner/BaseCaptioner.py @@ -274,26 +274,21 @@ class BaseCaptioner(BaseExtensionProcess): 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) @@ -455,7 +450,11 @@ class BaseCaptioner(BaseExtensionProcess): super(BaseCaptioner, self).on_error(e) if self.is_ui_captioner: try: - if not self.is_stopping: + if isinstance(e, KeyboardInterrupt): + # SIGINT (UI stop button or ctrl+c) is a stop, not an error + self.is_stopping = True + self.update_status("stopped", "Job stopped") + elif not self.is_stopping: self.update_status("error", str(e)) asyncio.run(self.wait_for_all_async()) except Exception as db_err: diff --git a/extensions_built_in/sd_trainer/DiffusionTrainer.py b/extensions_built_in/sd_trainer/DiffusionTrainer.py index f7d2ad6..ef9d58f 100644 --- a/extensions_built_in/sd_trainer/DiffusionTrainer.py +++ b/extensions_built_in/sd_trainer/DiffusionTrainer.py @@ -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) diff --git a/extensions_built_in/sd_trainer/UITrainer.py b/extensions_built_in/sd_trainer/UITrainer.py index 8b5fa79..f290fc9 100644 --- a/extensions_built_in/sd_trainer/UITrainer.py +++ b/extensions_built_in/sd_trainer/UITrainer.py @@ -33,7 +33,12 @@ class UITrainer(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): """ @@ -52,26 +57,21 @@ class UITrainer(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) @@ -217,7 +217,12 @@ class UITrainer(SDTrainer): def on_error(self, e: Exception): super(UITrainer, self).on_error(e) - if self.accelerator.is_main_process and not self.is_stopping: + if isinstance(e, KeyboardInterrupt): + # SIGINT (UI stop button or ctrl+c) is a stop, not an error + self.is_stopping = True + if self.accelerator.is_main_process: + self.update_status("stopped", "Job stopped") + elif self.accelerator.is_main_process and not self.is_stopping: self.update_status("error", str(e)) self.update_db_key("step", self.last_save_step) asyncio.run(self.wait_for_all_async()) diff --git a/run.py b/run.py index a77a792..3ee5708 100644 --- a/run.py +++ b/run.py @@ -14,6 +14,13 @@ if "SEED" in os.environ: print(f"Invalid SEED value: {os.environ['SEED']}. SEED must be an integer.") sys.path.insert(0, os.getcwd()) + +# The UI launches jobs with no console; keep anything we shell out to (torch +# compiles, HF git downloads) from flashing a console window. Must come before +# any import that might spawn a subprocess. +from toolkit.win_console import suppress_child_consoles +suppress_child_consoles() + # must come before ANY torch or fastai imports # import toolkit.cuda_malloc diff --git a/toolkit/win_console.py b/toolkit/win_console.py new file mode 100644 index 0000000..d7eca0f --- /dev/null +++ b/toolkit/win_console.py @@ -0,0 +1,58 @@ +"""Keep child processes from popping console windows on Windows. + +The UI worker launches training jobs with pythonw.exe and DETACHED_PROCESS so +they survive the UI shutting down (see ui/cron/actions/startJob.ts). That +leaves the job with no console at all, and Windows hands a brand new console +-- with a visible window -- to any console program launched from a process +that has none. MSVC during a torch/triton compile, git during a HF download +and ffmpeg would each flash a window on the user's desktop. + +Defaulting those spawns to CREATE_NO_WINDOW suppresses the flash. This is a +no-op unless we are on Windows *and* have no console, so running run.py from a +terminal behaves exactly as it did before. +""" + +import subprocess +import sys + +CREATE_NEW_CONSOLE = 0x00000010 +DETACHED_PROCESS = 0x00000008 +CREATE_NO_WINDOW = 0x08000000 + +# creationflags is the 14th positional parameter of Popen.__init__ after self. +_CREATIONFLAGS_POSITION = 14 + +_patched = False + + +def _has_console(): + import ctypes + + return bool(ctypes.windll.kernel32.GetConsoleWindow()) + + +def suppress_child_consoles(): + """Make CREATE_NO_WINDOW the default for subprocesses, where it matters.""" + global _patched + if _patched or sys.platform != "win32": + return + try: + if _has_console(): + return + except Exception: + # Never let a console tweak take down a training run. + return + + original_init = subprocess.Popen.__init__ + + def patched_init(self, *args, **kwargs): + if len(args) >= _CREATIONFLAGS_POSITION: + # Passed positionally; leave the caller's choice alone. + return original_init(self, *args, **kwargs) + flags = kwargs.get("creationflags", 0) + if not flags & (CREATE_NEW_CONSOLE | DETACHED_PROCESS | CREATE_NO_WINDOW): + kwargs["creationflags"] = flags | CREATE_NO_WINDOW + return original_init(self, *args, **kwargs) + + subprocess.Popen.__init__ = patched_init + _patched = True diff --git a/ui/cron/actions/startJob.ts b/ui/cron/actions/startJob.ts index 8beb186..0312b03 100644 --- a/ui/cron/actions/startJob.ts +++ b/ui/cron/actions/startJob.ts @@ -1,10 +1,10 @@ import prisma from '../prisma'; import { Job } from '@prisma/client'; -import { spawn } from 'child_process'; +import { spawn, ChildProcess } from 'child_process'; import path from 'path'; import fs from 'fs'; import { TOOLKIT_ROOT, getTrainingFolder, getHFToken } from '../paths'; -import { resolvePythonPath } from '../pythonPath'; +import { resolveDetachedPythonPath } from '../pythonPath'; const isWindows = process.platform === 'win32'; const appendJobLog = (logPath: string, message: string) => { @@ -13,6 +13,168 @@ const appendJobLog = (logPath: string, message: string) => { }); }; +// Windows only. Launched as `node -e ` so the job ends up outside the +// worker's process tree: `taskkill /T` (the dev script's `concurrently -k`, +// or any shutdown that kills the tree) walks parent/child links and would take +// a direct child down with the UI. This relay exits immediately, orphaning the +// job, and `detached` keeps the job alive once its parent is gone. Its own +// stdout/stderr are the job log, so the job inherits them as fds 1 and 2. +// Python failing to launch at all (broken venv, missing interpreter) happens +// inside the relay, so the relay -- not the worker -- is what sees that error. +// It reports it two ways: on stderr, which is the job log, and through the pid +// file, so the worker can put the real reason in the database. +const RELAY_ERROR_PREFIX = 'error:'; +const WINDOWS_RELAY_SCRIPT = ` +const { spawn } = require('child_process'); +const fs = require('fs'); +const [pidFile, command, ...args] = process.argv.slice(1); +const child = spawn(command, args, { + detached: true, + windowsHide: true, + stdio: ['ignore', 1, 2], +}); +child.once('error', error => { + process.stderr.write('Error launching job process: ' + error.message + '\\n'); + try { + fs.writeFileSync(pidFile, '${RELAY_ERROR_PREFIX}' + error.message); + } catch (e) { + process.stderr.write('Could not write job pid file: ' + e.message + '\\n'); + } + process.exit(1); +}); +if (child.pid) { + fs.writeFileSync(pidFile, String(child.pid)); + child.unref(); +} +`; + +const RELAY_PID_TIMEOUT_MS = 30000; + +type RelayResult = { pid: number | null; error?: string }; + +// The relay exits as soon as it has launched the job, leaving the real pid in +// pidPath. Without this we would only ever know the (already dead) relay's pid. +const readRelayPid = (relay: ChildProcess, pidPath: string): Promise => { + return new Promise(resolve => { + let settled = false; + const finish = (value: RelayResult) => { + if (settled) return; + settled = true; + clearTimeout(timer); + resolve(value); + }; + + const timer = setTimeout( + () => finish({ pid: null, error: 'Timed out waiting for the job process to start' }), + RELAY_PID_TIMEOUT_MS, + ); + + relay.once('exit', () => { + let contents: string; + try { + contents = fs.readFileSync(pidPath, 'utf8').trim(); + } catch { + finish({ pid: null, error: 'Job process did not report a pid' }); + return; + } + + if (contents.startsWith(RELAY_ERROR_PREFIX)) { + finish({ pid: null, error: contents.slice(RELAY_ERROR_PREFIX.length) }); + return; + } + + const pid = Number(contents); + finish( + Number.isInteger(pid) && pid > 0 + ? { pid } + : { pid: null, error: 'Job process did not report a usable pid' }, + ); + }); + + relay.once('error', error => finish({ pid: null, error: error.message })); + }); +}; + +const isProcessAlive = (pid: number): boolean => { + try { + process.kill(pid, 0); + return true; + } catch (e: any) { + // EPERM means it exists but belongs to someone else, which still counts. + return e?.code === 'EPERM'; + } +}; + +// We cannot read an exit code off a process that is not our child, so pull the +// last thing it said instead -- for a job that dies on startup (bad venv, +// missing CUDA libs) that traceback line is the whole diagnosis. +const LOG_TAIL_BYTES = 4096; +const LOG_TAIL_MAX_CHARS = 300; + +const readLogTail = (logPath: string): string | null => { + let fd: number | null = null; + try { + const size = fs.statSync(logPath).size; + const length = Math.min(size, LOG_TAIL_BYTES); + if (length === 0) return null; + + const buffer = Buffer.alloc(length); + fd = fs.openSync(logPath, 'r'); + fs.readSync(fd, buffer, 0, length, size - length); + + const lines = buffer.toString('utf8').split(/\r?\n/).filter(line => line.trim() !== ''); + const lastLine = lines[lines.length - 1]; + if (!lastLine) return null; + return lastLine.length > LOG_TAIL_MAX_CHARS ? `${lastLine.slice(-LOG_TAIL_MAX_CHARS)}` : lastLine; + } catch { + return null; + } finally { + if (fd !== null) { + try { + fs.closeSync(fd); + } catch { + // nothing useful to do if the log handle will not close + } + } + } +}; + +// The job is not our child anymore, so there is no 'exit' event to listen for. +// Poll instead, so a job that dies without updating its own row (OOM kill, +// hard crash) still gets marked as an error rather than sitting on 'running'. +const JOB_POLL_INTERVAL_MS = 2000; + +const watchDetachedJob = (pid: number, jobID: string, logPath: string) => { + const timer = setInterval(() => { + if (isProcessAlive(pid)) return; + clearInterval(timer); + + // A stopped or completed job writes its own ending (status row + final + // log lines) via its KeyboardInterrupt/done handlers -- stay out of the + // way. Only a job that vanished while still marked 'running' died without + // getting to say anything; record that. There is no exit code to read off + // a process that is not our child, so report the last thing it logged. + const tail = readLogTail(logPath); + const message = tail + ? `Job process exited unexpectedly. Last log line: ${tail}` + : 'Job process exited unexpectedly.'; + void prisma.job + .updateMany({ + where: { id: jobID, status: 'running' }, + data: { status: 'error', info: message, pid: null }, + }) + .then(result => { + if (result.count > 0) appendJobLog(logPath, `\n${message}\n`); + }) + .catch(updateError => { + console.error('Error updating job after process disappeared:', updateError); + }); + }, JOB_POLL_INTERVAL_MS); + + // Never hold the worker open on account of this poll. + if (timer.unref) timer.unref(); +}; + const startAndWatchJob = (job: Job) => { // starts and watches the job asynchronously return new Promise(async (resolve, reject) => { @@ -59,7 +221,7 @@ const startAndWatchJob = (job: Job) => { // write the config file fs.writeFileSync(configPath, JSON.stringify(jobConfig, null, 2)); - const pythonPath = resolvePythonPath(); + const pythonPath = resolveDetachedPythonPath(); const runFilePath = path.join(TOOLKIT_ROOT, 'run.py'); if (!fs.existsSync(runFilePath)) { @@ -90,6 +252,9 @@ const startAndWatchJob = (job: Job) => { const args = [runFilePath, configPath]; + // Where the Windows relay reports the job's real pid back to us. + const relayPidPath = path.join(trainingFolder, '.job_pid'); + let logFd: number | null = null; try { // Capture errors that occur before run.py can initialize file logging. @@ -97,14 +262,23 @@ const startAndWatchJob = (job: Job) => { let subprocess; if (isWindows) { - // Spawn Python directly on Windows so the process can survive parent exit - subprocess = spawn(pythonPath, args, { + // Launch through the relay (see WINDOWS_RELAY_SCRIPT) so the job is not + // a descendant of this worker and survives the UI being shut down or + // tree-killed. The relay spawns the job `detached`, which is what keeps + // it alive once the relay exits; that in turn means DETACHED_PROCESS, + // so pythonPath is pythonw.exe to avoid Windows handing the job a + // console window of its own. + try { + fs.unlinkSync(relayPidPath); + } catch { + // no stale pid file to clear + } + subprocess = spawn(process.execPath, ['-e', WINDOWS_RELAY_SCRIPT, relayPidPath, pythonPath, ...args], { env: { ...process.env, ...additionalEnv, }, cwd: TOOLKIT_ROOT, - detached: true, windowsHide: true, stdio: ['ignore', logFd, logFd], // don't tie stdio to parent; log fd passed as stdout and stderr }); @@ -136,25 +310,39 @@ const startAndWatchJob = (job: Job) => { }); }); - // Record abnormal termination and repair jobs Python could not update itself. - subprocess.once('exit', (code, signal) => { - if (code === 0) return; + let pid: number | null; - const result = signal ? `signal ${signal}` : `exit code ${code}`; - const message = `Job process terminated with ${result}.`; - appendJobLog(logPath, `\n${message}\n`); - void prisma.job - .updateMany({ - where: { id: jobID, status: 'running' }, - data: { status: 'error', info: message, pid: null }, - }) - .catch(updateError => { - console.error('Error updating job after abnormal process exit:', updateError); - }); - }); + if (isWindows) { + // The relay is gone within a few hundred ms; the pid it leaves behind is + // the job's. Poll that pid for liveness since we get no 'exit' event. + const relayResult = await readRelayPid(subprocess, relayPidPath); + if (relayResult.pid == null) { + throw new Error(relayResult.error ?? 'Job process did not report a pid'); + } + pid = relayResult.pid; + watchDetachedJob(pid, jobID, logPath); + } else { + pid = subprocess.pid ?? null; + + // Record abnormal termination and repair jobs Python could not update itself. + subprocess.once('exit', (code, signal) => { + if (code === 0) return; + + const result = signal ? `signal ${signal}` : `exit code ${code}`; + const message = `Job process terminated with ${result}.`; + appendJobLog(logPath, `\n${message}\n`); + void prisma.job + .updateMany({ + where: { id: jobID, status: 'running' }, + data: { status: 'error', info: message, pid: null }, + }) + .catch(updateError => { + console.error('Error updating job after abnormal process exit:', updateError); + }); + }); + } // Save the PID to the database and a file for future management (stop/inspect) - const pid = subprocess.pid ?? null; if (pid != null) { await prisma.job.update({ where: { id: jobID }, diff --git a/ui/cron/pythonPath.ts b/ui/cron/pythonPath.ts index 64ac231..05187e5 100644 --- a/ui/cron/pythonPath.ts +++ b/ui/cron/pythonPath.ts @@ -25,3 +25,18 @@ export const resolvePythonPath = (): string => { return isWindows ? 'python.exe' : 'python3'; }; + +// Interpreter for jobs we detach so they outlive the UI. On Windows detaching +// means DETACHED_PROCESS, and Windows gives any *console* app started that way +// a fresh console with a visible window. pythonw.exe is the GUI-subsystem +// build, so it never gets one; stdout/stderr still go to the handles we pass. +export const resolveDetachedPythonPath = (): string => { + const pythonPath = resolvePythonPath(); + if (!isWindows) return pythonPath; + + // Always take the pythonw next to the interpreter we already resolved, so + // both come from the same environment. Falling back to python.exe still + // works, it just shows a console window. + const pythonwPath = path.join(path.dirname(pythonPath), 'pythonw.exe'); + return fs.existsSync(pythonwPath) ? pythonwPath : pythonPath; +}; diff --git a/ui/src/app/api/jobs/[jobID]/stop/route.ts b/ui/src/app/api/jobs/[jobID]/stop/route.ts index cb14ad6..abfd11b 100644 --- a/ui/src/app/api/jobs/[jobID]/stop/route.ts +++ b/ui/src/app/api/jobs/[jobID]/stop/route.ts @@ -6,6 +6,43 @@ import { promisify } from 'util'; const execAsync = promisify(exec); const isWindows = process.platform === 'win32'; +// How long a stopping job gets to run its graceful shutdown (KeyboardInterrupt +// -> on_error -> final DB write) before we assume it is hung and force-kill it. +const GRACEFUL_STOP_TIMEOUT_MS = 60_000; +const BACKSTOP_POLL_MS = 2_000; + +const isProcessAlive = (pid: number): boolean => { + try { + process.kill(pid, 0); + return true; + } catch (e: any) { + // EPERM means it exists but belongs to someone else, which still counts. + return e?.code === 'EPERM'; + } +}; + +// Windows: the trainer's stop watcher sees the `stop` flag and raises SIGINT +// inside the process, so the graceful path runs without us sending anything. +// This backstop only exists for a job that is too hung to notice the flag. +// Poll (rather than one long timer) so we stop watching the moment the pid +// dies and never touch a recycled pid. +const scheduleForceKillBackstop = (pid: number, jobID: string) => { + const startedAt = Date.now(); + const timer = setInterval(() => { + if (!isProcessAlive(pid)) { + clearInterval(timer); + return; + } + if (Date.now() - startedAt < GRACEFUL_STOP_TIMEOUT_MS) return; + clearInterval(timer); + console.warn(`Job ${jobID} (pid ${pid}) still alive ${GRACEFUL_STOP_TIMEOUT_MS / 1000}s after stop request, force killing`); + execAsync(`taskkill /PID ${pid} /T /F`, { windowsHide: true }).catch(() => { + // already gone + }); + }, BACKSTOP_POLL_MS); + timer.unref?.(); +}; + export async function GET(request: NextRequest, { params }: { params: { jobID: string } }) { const { jobID } = await params; @@ -30,9 +67,12 @@ export async function GET(request: NextRequest, { params }: { params: { jobID: s console.log(`Attempting to stop job ${jobID} with PID ${job.pid}`); try { if (isWindows) { - // Windows doesn't support SIGINT for arbitrary processes. - // Use taskkill with /T (tree) to send a CTRL+C-like termination. - await execAsync(`taskkill /PID ${job.pid} /T /F`, { windowsHide: true }); + // No external SIGINT possible on Windows (the job runs under pythonw + // with no console), and none is needed: the `stop` flag written above + // is the signal. The trainer's stop watcher polls it and raises + // SIGINT in-process, which runs the same KeyboardInterrupt shutdown + // as Linux -- final DB write and the closing log lines. + scheduleForceKillBackstop(job.pid, jobID); } else { process.kill(job.pid, 'SIGINT'); }