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,81 @@
+import http from 'node:http'
+import net from 'node:net'
+import { GPU_WAIT_MESSAGE } from './gpu-reservation.mjs'
+
+function reply(res, status, body) {
+  if (res.headersSent) return res.destroy()
+  res.writeHead(status, { 'Content-Type': 'application/json', 'Cache-Control': 'no-store' })
+  res.end(JSON.stringify(body))
+}
+
+/** Streams Comfy HTTP/WebSocket traffic; all mutations require the current reservation. */
+export function createGpuProxy({ target, reservation, authorized = () => true, markWork = () => {} }) {
+  const server = http.createServer(async (req, res) => {
+    const path = new URL(req.url || '/', 'http://localhost').pathname.replace(/\/+$/, '') || '/'
+    // A global history list can expose another studio's prompts and outputs.
+    if (req.method === 'GET' && (path === '/history' || path === '/history/')) return reply(res, 200, {})
+    const forward = () => new Promise(resolve => {
+      const port = target()
+      if (!port) { reply(res, 503, { error: 'GPU service is unavailable.' }); return resolve() }
+      const headers = { ...req.headers }
+      delete headers['x-aigen-gpu-lease']
+      delete headers.authorization
+      const upstream = http.request({ hostname: '127.0.0.1', port, path: req.url, method: req.method, headers }, response => {
+        if (req.method === 'GET' && path === '/queue') {
+          let size = 0
+          const chunks = []
+          response.on('data', chunk => {
+            size += chunk.length
+            if (size > 16 * 1024 * 1024) { response.destroy(); reply(res, 502, { error: 'GPU status unavailable.' }); resolve(); return }
+            chunks.push(chunk)
+          })
+          response.on('end', () => {
+            try {
+              if (response.statusCode !== 200) throw new Error('status unavailable')
+              const queue = JSON.parse(Buffer.concat(chunks).toString('utf8'))
+              if (!Array.isArray(queue.queue_running) || !Array.isArray(queue.queue_pending)) throw new Error('invalid queue')
+              // Preserve count compatibility without revealing IDs, graphs, or prompts.
+              reply(res, 200, { queue_running: queue.queue_running.map(() => null), queue_pending: queue.queue_pending.map(() => null) })
+            } catch { reply(res, 502, { error: 'GPU status unavailable.' }) }
+            resolve()
+          })
+        } else {
+          res.writeHead(response.statusCode || 502, response.headers)
+          response.pipe(res)
+          response.on('end', resolve)
+        }
+        response.on('error', () => { res.destroy(); resolve() })
+      })
+      upstream.on('error', () => { reply(res, 502, { error: 'GPU service is unavailable.' }); resolve() })
+      upstream.setTimeout(5 * 60 * 1000, () => upstream.destroy())
+      req.on('aborted', () => upstream.destroy())
+      res.on('close', () => { upstream.destroy(); resolve() })
+      req.pipe(upstream)
+    })
+    try {
+      if (req.method === 'GET' || req.method === 'HEAD') await forward()
+      else {
+        if (!authorized(req)) { req.resume(); return reply(res, 401, { error: 'Unauthorized' }) }
+        await reservation.permit(String(req.headers['x-aigen-gpu-lease'] || ''), async () => { markWork(); await forward() })
+      }
+    } catch {
+      req.resume()
+      reply(res, 409, { error: { message: GPU_WAIT_MESSAGE }, code: 'GPU_BUSY' })
+    }
+  })
+  server.on('upgrade', (req, socket, head) => {
+    if (new URL(req.url || '/', 'http://localhost').pathname !== '/ws' || !target()) return socket.destroy()
+    const upstream = net.connect(target(), '127.0.0.1', () => {
+      const headers = Object.entries(req.headers).filter(([key]) => key !== 'authorization' && key !== 'x-aigen-gpu-lease')
+        .map(([key, value]) => `${key}: ${value}`).join('\r\n')
+      upstream.write(`${req.method} ${req.url} HTTP/${req.httpVersion}\r\n${headers}\r\n\r\n`)
+      if (head.length) upstream.write(head)
+      socket.pipe(upstream); upstream.pipe(socket)
+    })
+    socket.on('error', () => upstream.destroy())
+    upstream.on('error', () => socket.destroy())
+    socket.on('close', () => upstream.destroy())
+    upstream.on('close', () => socket.destroy())
+  })
+  return server
+}