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

33
extensions_built_in/captioner/BaseCaptioner.py
View File

@@ -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:

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)

41
extensions_built_in/sd_trainer/UITrainer.py
View File

@@ -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())

7
run.py
View File

@@ -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

58
toolkit/win_console.py Normal file
View File

@@ -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

232
ui/cron/actions/startJob.ts
View File

@@ -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 <this>` 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<RelayResult> => {
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<void>(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 },

15
ui/cron/pythonPath.ts
View File

@@ -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;
};

46
ui/src/app/api/jobs/[jobID]/stop/route.ts
View File

@@ -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');
}