Coordinate shared GPU access across studio instances

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 }
+}