83 lines
4.2 KiB
JavaScript
83 lines
4.2 KiB
JavaScript
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 = () => {}, externalBusy = () => false }) {
|
|
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 (externalBusy()) { req.resume(); return reply(res, 409, { error: { message: 'A music or upscale job is using the GPU.' }, code: 'GPU_BUSY' }) }
|
|
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
|
|
}
|