Coordinate shared GPU access across studio instances

2026-09-06 16:32:18 -05:00
parent 911ddeab55
commit 0c6a035d04
17 changed files with 582 additions and 64 deletions
@@ -0,0 +1,68 @@
+import { randomUUID } from 'node:crypto'
+
+export const GPU_WAIT_MESSAGE = 'GPU is in use. Waiting for availability.'
+
+/** Device-local arbitration. Tickets/tokens are opaque; no studio or job metadata is stored. */
+export function createGpuReservation({ idle, now = Date.now, ttlMs = 60_000, ticketTtlMs = 30_000 } = {}) {
+  if (typeof idle !== 'function') throw new Error('An authoritative GPU idle probe is required')
+  let owner = null
+  let gate = Promise.resolve()
+  let inFlight = 0
+  const waiting = new Map()
+  const serialize = fn => {
+    const run = gate.then(fn)
+    gate = run.catch(() => {})
+    return run
+  }
+  const valid = token => Boolean(owner && !owner.released && owner.token === token && owner.until > now())
+  const busy = () => ({ acquired: false, message: GPU_WAIT_MESSAGE, retryAfterMs: 2500 })
+  const prune = () => {
+    for (const [ticket, seen] of waiting) if (now() - seen >= ticketTtlMs) waiting.delete(ticket)
+  }
+  return {
+    acquire(ticket) {
+      return serialize(async () => {
+        if (typeof ticket !== 'string' || !/^[a-zA-Z0-9-]{16,80}$/.test(ticket)) throw new Error('Invalid reservation ticket')
+        prune()
+        if (owner?.ticket === ticket && valid(owner.token)) {
+          owner.until = now() + ttlMs
+          return { acquired: true, token: owner.token, ttlMs }
+        }
+        waiting.set(ticket, now())
+        if (owner && (valid(owner.token) || inFlight > 0)) return busy()
+        if (waiting.keys().next().value !== ticket) return busy()
+        // Unknown/offline is not idle; the adapter must explicitly account for a stopped GPU.
+        if (await idle() !== true) return busy()
+        owner = { ticket, token: randomUUID(), until: now() + ttlMs, released: false }
+        waiting.delete(ticket)
+        return { acquired: true, token: owner.token, ttlMs }
+      })
+    },
+    renew(token) {
+      return serialize(() => {
+        if (!valid(token)) return { renewed: false }
+        owner.until = now() + ttlMs
+        return { renewed: true, ttlMs }
+      })
+    },
+    release(token) {
+      return serialize(() => {
+        if (!owner || owner.token !== token) return { released: false }
+        owner.released = true
+        // The next acquire still verifies that Comfy has drained before granting access.
+        return { released: true }
+      })
+    },
+    cancel(ticket) { return serialize(() => { waiting.delete(ticket); return { cancelled: true } }) },
+    async permit(token, operation) {
+      const accepted = await serialize(() => {
+        if (!valid(token)) return false
+        inFlight += 1
+        return true
+      })
+      if (!accepted) throw Object.assign(new Error(GPU_WAIT_MESSAGE), { statusCode: 409 })
+      try { return await operation() } finally { inFlight -= 1 }
+    },
+    availability() { return { busy: Boolean(owner && (valid(owner.token) || inFlight > 0)) } }
+  }
+}