From a705610721e86f98d740511f87759022624240e8 Mon Sep 17 00:00:00 2001 From: Towsty Date: Tue, 1 Sep 2026 20:53:52 -0500 Subject: [PATCH] Refuse to start Comfy while AITraining has the GPU. Co-authored-by: Cursor --- scripts/comfy-host-agent.mjs | 39 +++++++++++++++++++++++++++++++++- server/utils/comfyLifecycle.ts | 19 ++++++++++++++++- 2 files changed, 56 insertions(+), 2 deletions(-) diff --git a/scripts/comfy-host-agent.mjs b/scripts/comfy-host-agent.mjs index 88b9634..267033c 100644 --- a/scripts/comfy-host-agent.mjs +++ b/scripts/comfy-host-agent.mjs @@ -12,18 +12,43 @@ const proxyPort = Number(process.env.COMFY_PROXY_PORT || 8198) const token = process.env.COMFY_CONTROL_TOKEN || '' const defaultHttp = (process.env.COMFY_HOST || 'http://127.0.0.1:8188').replace(/\/$/, '') const idleMs = Math.max(60_000, Number(process.env.COMFY_IDLE_MS || 30 * 60 * 1000) || 30 * 60 * 1000) +const trainingUrl = String(process.env.TRAINING_CONTROL_URL || 'http://127.0.0.1:8200').replace(/\/$/, '') let lastWorkAt = Date.now() let lastQueueRunning = 0 let stoppedByAgent = false let lastHealthyPort = 0 let lastProcessUp = false +let lastTraining = { busy: false, jobId: null, status: null, name: null, message: '' } function markWork() { lastWorkAt = Date.now() stoppedByAgent = false } +async function trainingLock() { + try { + const res = await fetch(`${trainingUrl}/status`, { signal: AbortSignal.timeout(2500) }) + const body = await res.json().catch(() => null) + const job = body?.job + const status = String(job?.status || '') + const busy = Boolean(job && ['preparing', 'queued', 'running', 'stopping'].includes(status)) + lastTraining = { + busy, + jobId: job?.id || null, + status: busy ? status : null, + name: busy ? (job?.outputName || null) : null, + message: busy + ? `AITraining is ${status}${job?.outputName ? ` (${job.outputName})` : ''}. Stop that job before poking Comfy.` + : '' + } + return lastTraining + } catch { + lastTraining = { busy: false, jobId: null, status: null, name: null, message: '' } + return lastTraining + } +} + function json(res, status, body) { const payload = JSON.stringify(body) res.writeHead(status, { 'Content-Type': 'application/json', 'Content-Length': Buffer.byteLength(payload) }) @@ -603,10 +628,20 @@ const server = http.createServer(async (req, res) => { lastActivityAt: new Date(lastWorkAt).toISOString(), idleMs, port: healthyPort || null, - proxyPort: healthyPort ? proxyPort : null + proxyPort: healthyPort ? proxyPort : null, + training: lastTraining }) } if (req.method === 'POST' && url.pathname === '/start') { + const training = await trainingLock() + if (training.busy) { + return json(res, 409, { + ok: false, + error: 'train-busy', + message: training.message, + training + }) + } const healthyPort = await syncProxy() if (healthyPort) { markWork() @@ -663,6 +698,8 @@ server.listen(port, '0.0.0.0', async () => { idleMs, candidates: candidatePorts() })) + void trainingLock() + setInterval(() => { void trainingLock() }, 3000) let lastIdleCheck = 0 setInterval(() => { syncProxy().then(async (nextPort) => { diff --git a/server/utils/comfyLifecycle.ts b/server/utils/comfyLifecycle.ts index 8f67b33..9cb8b83 100644 --- a/server/utils/comfyLifecycle.ts +++ b/server/utils/comfyLifecycle.ts @@ -109,6 +109,8 @@ async function controlRequest(path: string, method = 'GET', timeoutMs = 5000) { stopped?: boolean killed?: unknown error?: string + message?: string + training?: { busy?: boolean; jobId?: string | null; status?: string | null; name?: string | null; message?: string } running?: number pending?: number status?: number @@ -118,6 +120,8 @@ async function controlRequest(path: string, method = 'GET', timeoutMs = 5000) { ok: false as const, status: res.status, error: body?.error, + message: body?.message, + training: body?.training, running: body?.running, pending: body?.pending, killed: body?.killed @@ -133,6 +137,15 @@ export async function fetchComfyControlStatus() { return await controlRequest('/status') } +function throwIfTrainingBusy(remote: { training?: { busy?: boolean; message?: string }; error?: string; message?: string } | null) { + const busy = Boolean(remote?.training?.busy || remote?.error === 'train-busy') + if (!busy) return + throw createError({ + statusCode: 409, + statusMessage: remote?.training?.message || remote?.message || 'AITraining is using the GPU. Stop that job before poking Comfy.' + }) +} + export async function requestComfyWake() { const { controlUrl } = settings() if (!controlUrl) { @@ -148,10 +161,11 @@ export async function requestComfyWake() { statusMessage: 'Beast host agent is unreachable. Confirm it is running on the desktop.' }) } + throwIfTrainingBusy(remote) if (remote.ok === false) { throw createError({ statusCode: Number(remote.status) || 502, - statusMessage: 'Beast host agent refused to start Comfy.' + statusMessage: String(remote.message || remote.error || 'Beast host agent refused to start Comfy.') }) } return { @@ -250,6 +264,7 @@ async function startComfy() { lastStartAt = now const remote = await controlRequest('/start', 'POST') + throwIfTrainingBusy(remote) if (remote?.ok || remote?.started) { log('start-remote', { pid: remote.pid || null }) return { started: true, reason: 'remote' as const } @@ -322,6 +337,8 @@ function hostWithPort(host: string, port: number) { async function ensureUnlocked(onStatus: StatusFn, skipBusyWait = false) { const cfg = settings() const host = cfg.host + const lock = await controlRequest('/status') + throwIfTrainingBusy(lock) onStatus({ state: 'offline', message: 'Checking ComfyUI...', host }) const health = await checkComfyHttp(cfg.healthTimeoutMs)