Coordinate shared GPU access across studio instances

server/utils/comfy.ts

@@ -1,3 +1,4 @@
+import { assertSharedGpu, sharedGpuHeaders } from '~/server/utils/sharedGpu'
 import { LTX_NEGATIVE } from '~/utils/videoModels'
 import { createHash } from 'node:crypto'
 import { comfyJobPrefix } from '~/utils/outputNames'
@@ -73,6 +74,9 @@ function configuredComfyBase() {
 }
 
 export function getComfyHost() {
+  // Shared studios must never bypass the reservation-enforcing proxy.
+  const sharedHost = agentFrontDoorOrigin()
+  if (sharedHost) return sharedHost
   if (comfyHostOverride) return viaAgentFrontDoor(comfyHostOverride) || comfyHostOverride
   return viaAgentFrontDoor(configuredComfyBase()) || configuredComfyBase()
 }
@@ -83,6 +87,12 @@ export function comfyWsUrl(clientId: string) {
 
 export async function comfyFetch(path: string, init?: RequestInit) {
   const url = `${getComfyHost()}${path}`
+  if (init?.method && !['GET', 'HEAD'].includes(init.method.toUpperCase())) {
+    await assertSharedGpu()
+    const headers = new Headers(init.headers)
+    for (const [key, value] of Object.entries(sharedGpuHeaders())) headers.set(key, value)
+    init = { ...init, headers }
+  }
   try {
     return await fetch(url, init)
   } catch (error) {
@@ -623,7 +633,8 @@ async function purgeOnDesktop(opts: {
       headers: {
         'Content-Type': 'application/json',
         Accept: 'application/json',
-        ...(token ? { Authorization: `Bearer ${token}` } : {})
+        ...(token ? { Authorization: `Bearer ${token}` } : {}),
+        ...sharedGpuHeaders()
       },
       body: JSON.stringify({
         imageName: imageNames[0] || '',

server/utils/comfyLifecycle.ts

@@ -1,3 +1,4 @@
+import { assertSharedGpu, sharedGpuHeaders, withSharedGpuStart } from '~/server/utils/sharedGpu'
 import { execFile, spawn } from 'node:child_process'
 import { promisify } from 'node:util'
 import { getComfyHost, setComfyHostOverride, comfyConfigured } from '~/server/utils/comfy'
@@ -85,13 +86,15 @@ export async function fetchQueue() {
 async function controlRequest(path: string, method = 'GET', timeoutMs = 5000, body?: Record<string, unknown>) {
   const { controlUrl, controlToken } = settings()
   if (!controlUrl) return null
+  if (method !== 'GET') await assertSharedGpu()
   try {
     const res = await fetch(`${controlUrl}${path}`, {
       method,
       headers: {
         Accept: 'application/json',
         ...(body ? { 'Content-Type': 'application/json' } : {}),
-        ...(controlToken ? { Authorization: `Bearer ${controlToken}` } : {})
+        ...(controlToken ? { Authorization: `Bearer ${controlToken}` } : {}),
+        ...(method !== 'GET' ? sharedGpuHeaders() : {})
       },
       ...(body ? { body: JSON.stringify(body) } : {}),
       signal: AbortSignal.timeout(timeoutMs)
@@ -344,7 +347,7 @@ async function waitWhileBusy(onStatus: StatusFn) {
   const started = Date.now()
   while (Date.now() - started < busyWaitMs) {
     const queue = await fetchQueue()
-    if (queue.running === 0) return
+    if (queue.ok && queue.running === 0 && queue.pending === 0) return
     onStatus({
       state: 'busy',
       message: `ComfyUI is busy with another prompt (${queue.running} running, ${queue.pending} queued). Waiting up to 3 minutes, then this job will be saved.`,
@@ -488,7 +491,10 @@ async function ensureUnlocked(onStatus: StatusFn, skipBusyWait = false) {
 
 export function ensureComfyReady(onStatus: StatusFn = () => undefined, options?: { skipBusyWait?: boolean }) {
   const skipBusyWait = options?.skipBusyWait === true
-  const run = gate.then(() => ensureUnlocked(onStatus, skipBusyWait))
+  const run = gate.then(() => withSharedGpuStart(
+    () => ensureUnlocked(onStatus, skipBusyWait),
+    async () => { await assertSharedGpu() }
+  ))
   gate = run.then(() => undefined, () => undefined)
   return run
 }

server/utils/imageComfy.ts

@@ -1,3 +1,4 @@
+import { assertSharedGpu, sharedGpuHeaders, sharedGpuConfigured } from '~/server/utils/sharedGpu'
 import { createHash } from 'node:crypto'
 import { AsyncLocalStorage } from 'node:async_hooks'
 import { getComfyHost, viaAgentFrontDoor } from '~/server/utils/comfy'
@@ -61,6 +62,7 @@ export function getSidecarImageHost() {
 }
 
 export function getImageComfyHost() {
+  if (sharedGpuConfigured()) return getComfyHost()
   const fromAls = imageHostAls.getStore()
   if (fromAls) return viaAgentFrontDoor(fromAls) || fromAls
   if (imageComfyHostOverride) return viaAgentFrontDoor(imageComfyHostOverride) || imageComfyHostOverride
@@ -80,6 +82,12 @@ export async function imageComfyFetch(path: string, init?: RequestInit) {
   if (!host) {
     throw createError({ statusCode: 500, statusMessage: 'No image ComfyUI host is configured' })
   }
+  if (init?.method && !['GET', 'HEAD'].includes(init.method.toUpperCase())) {
+    await assertSharedGpu()
+    const headers = new Headers(init.headers)
+    for (const [key, value] of Object.entries(sharedGpuHeaders())) headers.set(key, value)
+    init = { ...init, headers }
+  }
   try {
     return await fetch(`${host}${path}`, init)
   } catch (error) {
@@ -390,7 +398,8 @@ async function purgeImageOnDesktop(opts: {
       headers: {
         'Content-Type': 'application/json',
         Accept: 'application/json',
-        ...(token ? { Authorization: `Bearer ${token}` } : {})
+        ...(token ? { Authorization: `Bearer ${token}` } : {}),
+        ...sharedGpuHeaders()
       },
       body: JSON.stringify({
         imageName: imageNames[0] || '',

server/utils/imageComfyLifecycle.ts

@@ -1,3 +1,4 @@
+import { assertSharedGpu, sharedGpuHeaders, withSharedGpuStart } from '~/server/utils/sharedGpu'
 import { getImageComfyHost, setImageComfyHostOverride, imageComfyConfigured, getBeastImageHost, getSidecarImageHost, withImageComfyHost, type ImageEditBox } from '~/server/utils/imageComfy'
 
 export type ImageComfyState = 'online' | 'booting' | 'busy' | 'starting' | 'offline'
@@ -92,12 +93,14 @@ export async function fetchImageQueue() {
 async function controlRequest(path: string, method = 'GET', timeoutMs = 5000) {
   const { controlUrl, controlToken } = settings()
   if (!controlUrl) return null
+  if (method !== 'GET') await assertSharedGpu()
   try {
     const res = await fetch(`${controlUrl}${path}`, {
       method,
       headers: {
         Accept: 'application/json',
-        ...(controlToken ? { Authorization: `Bearer ${controlToken}` } : {})
+        ...(controlToken ? { Authorization: `Bearer ${controlToken}` } : {}),
+        ...(method !== 'GET' ? sharedGpuHeaders() : {})
       },
       signal: AbortSignal.timeout(timeoutMs)
     })
@@ -238,7 +241,7 @@ async function ensureUnlocked(onStatus: StatusFn) {
 }
 
 export function ensureImageComfyReady(onStatus: StatusFn = () => undefined) {
-  const run = gate.then(() => ensureUnlocked(onStatus))
+  const run = gate.then(() => withSharedGpuStart(() => ensureUnlocked(onStatus), async () => { await assertSharedGpu() }))
   gate = run.then(() => undefined, () => undefined)
   return run
 }

server/utils/promptComfy.ts

@@ -1,4 +1,5 @@
-import { comfyInputFilename } from '~/server/utils/comfy'
+import { assertSharedGpu, sharedGpuHeaders, sharedGpuConfigured } from '~/server/utils/sharedGpu'
+import { comfyInputFilename, getComfyHost } from '~/server/utils/comfy'
 import { getSidecarImageHost } from '~/server/utils/imageComfy'
 import { ensureSidecarReady } from '~/server/utils/imageComfyLifecycle'
 import { buildVisionPromptWorkflow } from '~/server/utils/promptWorkflow'
@@ -19,6 +20,7 @@ const jobs = new Map<string, PromptJob>()
 const MAX_JOBS = 20
 
 export function getPromptComfyHost() {
+  if (sharedGpuConfigured()) return getComfyHost()
   return getSidecarImageHost()
 }
 
@@ -31,6 +33,12 @@ async function promptComfyFetch(path: string, init?: RequestInit) {
   if (!host) {
     throw createError({ statusCode: 503, statusMessage: 'Qwen VL is not configured. Set COMFY_HOST or IMAGE_COMFY_HOST to Beast.' })
   }
+  if (init?.method && !['GET', 'HEAD'].includes(init.method.toUpperCase())) {
+    await assertSharedGpu()
+    const headers = new Headers(init.headers)
+    for (const [key, value] of Object.entries(sharedGpuHeaders())) headers.set(key, value)
+    init = { ...init, headers }
+  }
   try {
     return await fetch(`${host}${path}`, init)
   } catch (error) {
@@ -260,3 +268,7 @@ export function failPromptJob(job: PromptJob, error: unknown) {
   job.error = err.statusMessage || err.message || 'Prompt recommend failed'
   job.message = job.error
 }
+
+export function hasActivePromptJobs() {
+  return [...jobs.values()].some(job => job.status === 'queued' || job.status === 'running')
+}

server/utils/sharedGpu.ts

@@ -0,0 +1,106 @@
+import { randomUUID } from 'node:crypto'
+
+const ticket = randomUUID()
+let lease = ''
+let expiresAt = 0
+let lost = false
+let starting = 0
+let resetting = false
+let gate: Promise<unknown> = Promise.resolve()
+let localWork: () => boolean = () => true
+let reason = 'GPU is in use. Waiting for availability.'
+
+function settings() {
+  const config = useRuntimeConfig()
+  return {
+    url: String(config.comfyControlUrl || process.env.COMFY_CONTROL_URL || '').replace(/\/$/, ''),
+    token: String(config.comfyControlToken || process.env.COMFY_CONTROL_TOKEN || '')
+  }
+}
+export function sharedGpuConfigured() { return Boolean(settings().url) }
+export function sharedGpuWaitReason() { return reason }
+export function registerSharedGpuWork(check: () => boolean) { localWork = check }
+
+async function request(path: string, body: Record<string, string>) {
+  const { url, token } = settings()
+  const response = await fetch(`${url}/gpu/${path}`, {
+    method: 'POST', headers: { 'Content-Type': 'application/json', ...(token ? { Authorization: `Bearer ${token}` } : {}) },
+    body: JSON.stringify(body), signal: AbortSignal.timeout(15_000)
+  })
+  if (!response.ok) throw new Error('GPU coordinator unavailable')
+  return await response.json() as { acquired?: boolean; token?: string; ttlMs?: number; renewed?: boolean; released?: boolean }
+}
+function serialized<T>(fn: () => Promise<T>): Promise<T> {
+  const run = gate.then(fn)
+  gate = run.catch(() => {})
+  return run
+}
+
+export function acquireSharedGpu() {
+  return serialized(async () => {
+    if (!sharedGpuConfigured()) return true
+    if (lost) return false
+    if (lease && Date.now() < expiresAt) return true
+    if (lease) { lost = true; return false }
+    try {
+      const result = await request('acquire', { ticket })
+      if (!result.acquired || !result.token || !Number.isFinite(result.ttlMs) || Number(result.ttlMs) < 10_000) {
+        reason = 'GPU is in use. Waiting for availability.'
+        return false
+      }
+      lease = result.token
+      expiresAt = Date.now() + Math.max(1000, Number(result.ttlMs) - 5000)
+      return true
+    } catch {
+      reason = 'GPU coordinator unavailable. Waiting to reconnect.'
+      return false
+    }
+  })
+}
+
+export async function withSharedGpuStart<T>(fn: () => Promise<T>, blocked: () => Promise<T>): Promise<T> {
+  starting += 1
+  try { return await (!resetting && await acquireSharedGpu() ? fn() : blocked()) }
+  finally { starting -= 1 }
+}
+
+export async function assertSharedGpu() {
+  if (!await acquireSharedGpu()) throw createError({ statusCode: 409, statusMessage: reason, data: { code: 'GPU_BUSY' } })
+}
+
+export function sharedGpuHeaders() {
+  if (!sharedGpuConfigured()) return {} as Record<string, string>
+  if (!lease || lost || Date.now() >= expiresAt) {
+    throw createError({ statusCode: 409, statusMessage: 'GPU reservation is unavailable. Waiting to reconnect.', data: { code: 'GPU_BUSY' } })
+  }
+  const { token } = settings()
+  return { 'x-aigen-gpu-lease': lease, ...(token ? { Authorization: `Bearer ${token}` } : {}) }
+}
+
+export function maintainSharedGpu() {
+  return serialized(async () => {
+    if (!sharedGpuConfigured()) return
+    if (!starting && !resetting && !localWork()) {
+      if (lease) {
+        try { await request('release', { token: lease }) } catch { /* expires if the host is unreachable */ }
+      }
+      lease = ''; expiresAt = 0; lost = false
+      return
+    }
+    if (!lease || lost) return
+    try {
+      const result = await request('renew', { token: lease })
+      if (!result.renewed || !Number.isFinite(result.ttlMs) || Number(result.ttlMs) < 10_000) { lost = true; return }
+      expiresAt = Date.now() + Math.max(1000, Number(result.ttlMs) - 5000)
+    } catch {
+      if (Date.now() >= expiresAt) lost = true
+    }
+  })
+}
+
+export async function withSharedGpuReset<T>(fn: () => Promise<T>) {
+  if (resetting) throw createError({ statusCode: 409, statusMessage: 'GPU reset is already in progress.' })
+  resetting = true
+  try { await assertSharedGpu(); return await fn() }
+  finally { resetting = false }
+}

server/utils/studioQueue.ts

@@ -1,3 +1,4 @@
+import { withSharedGpuStart, sharedGpuWaitReason, maintainSharedGpu, acquireSharedGpu } from '~/server/utils/sharedGpu'
 import { existsSync, mkdirSync, readdirSync, readFileSync, renameSync, writeFileSync } from 'node:fs'
 import { join } from 'node:path'
 import { getJob, listJobs, emitJob, type Job } from '~/server/utils/jobs'
@@ -98,6 +99,7 @@ export interface StudioJob {
   pausedByUser?: boolean
   resumeAutoRun?: boolean
   lastError?: string
+  waitReason?: string
 }
 
 type StudioQueueStore = {
@@ -246,7 +248,8 @@ export function summarizeStudioJob(job: StudioJob) {
     holdForCutIn: job.holdForCutIn === true,
     pauseAfterCurrent: job.pauseAfterCurrent === true,
     pausedByUser: job.pausedByUser === true,
-    lastError: job.lastError
+    lastError: job.lastError,
+    waitReason: job.status === 'waiting' ? job.waitReason : undefined
   }
 }
 
@@ -415,14 +418,11 @@ export async function videoJobsBusy() {
   }))) return true
   const queue = await fetchLiveQueue()
   if (queue) {
-    if (queue.running > 0) return true
-    if (queue.pending > 0) {
-      if (listJobs().some(job => liveJobOwnsGpu(job) || jobIsLocallySubmitting(job))) return true
-      if (listPendingJobs().some((pending) => {
-        if (!pending.promptId || pending.stopAfterCurrent) return false
-        return Date.now() - pending.startedAt < PENDING_ORPHAN_MS
-      })) return true
-      return false
+    if (queue.running > 0 || queue.pending > 0) {
+      // Join the device queue while another instance is rendering, so it cannot
+      // reclaim every gap ahead of this waiting instance.
+      await acquireSharedGpu()
+      return true
     }
     return false
   }
@@ -865,7 +865,7 @@ function pendingAlive(job: StudioJob) {
 }
 
 function isTransientComfyError(error?: string) {
-  return /still busy|COMFY_BUSY|offline|never answered|host agent|unreachable|asleep|Starting Comfy|ComfyUI \(1\)|instance picker|refused to start|API never answered|cold start|ECONNREFUSED|ETIMEDOUT|fetch failed|502|503/i.test(String(error || ''))
+  return /GPU is in use|GPU reservation|GPU coordinator|GPU_BUSY|still busy|COMFY_BUSY|offline|never answered|host agent|unreachable|asleep|Starting Comfy|ComfyUI \(1\)|instance picker|refused to start|API never answered|cold start|ECONNREFUSED|ETIMEDOUT|fetch failed|502|503/i.test(String(error || ''))
 }
 
 async function parkStudioOnStartFailure(owner: string, id: string, message: string) {
@@ -1010,11 +1010,12 @@ async function dispatchStudioQueue() {
   }
   await settleAttachedTerminalLives()
   if (await videoJobsBusy()) {
-    if (owners.some(owner => readJobs(owner).some(job => job.status === 'waiting'))) {
+    if (owners.some(owner => readJobs(owner).some(job => job.status === 'waiting' || (job.status === 'held' && !job.pausedByUser)))) {
       scheduleKickRetry()
     }
     return
   }
+  await maintainSharedGpu()
   // Re-read after videoJobsBusy — reap/settle may have rewritten disk rows.
   for (const owner of owners) {
     const store = readStore(owner)
@@ -1039,6 +1040,13 @@ async function dispatchStudioQueue() {
 }
 
 async function resumeHeldStudioJob(item: StudioJob) {
+  return withSharedGpuStart(() => resumeHeldStudioJobReserved(item), async () => {
+    scheduleKickRetry()
+    return false
+  })
+}
+
+async function resumeHeldStudioJobReserved(item: StudioJob) {
   const latest = readJobs(item.ownerKey).find(job => job.id === item.id)
   if (!latest || latest.status === 'cancelled') return false
   if (!latest.shotQueueId) {
@@ -1386,6 +1394,13 @@ async function startStudioMusicJob(item: StudioJob) {
 }
 
 export async function startStudioJob(item: StudioJob) {
+  return withSharedGpuStart(() => startStudioJobReserved(item), async () => {
+    await patchStudioJob(item.ownerKey, item.id, row => { row.waitReason = sharedGpuWaitReason() })
+    scheduleKickRetry()
+  })
+}
+
+async function startStudioJobReserved(item: StudioJob) {
   if (studioJobKind(item) === 'music') {
     await startStudioMusicJob(item)
     return