Coordinate shared GPU access across studio instances

tests/shared-gpu.test.mjs

@@ -0,0 +1,81 @@
+import test from 'node:test'
+import assert from 'node:assert/strict'
+import http from 'node:http'
+import { createRequire } from 'node:module'
+import { readFileSync } from 'node:fs'
+import ts from 'typescript'
+import { createGpuReservation } from '../scripts/gpu-reservation.mjs'
+import { createGpuProxy } from '../scripts/gpu-proxy.mjs'
+const require = createRequire(import.meta.url)
+const compiled = ts.transpileModule(readFileSync(new URL('../server/utils/sharedGpu.ts', import.meta.url), 'utf8'), {
+  compilerOptions: { module: ts.ModuleKind.CommonJS, target: ts.ScriptTarget.ES2022 }
+}).outputText
+const listen = server => new Promise(resolve => server.listen(0, '127.0.0.1', () => resolve(server.address().port)))
+const close = server => new Promise(resolve => { server.closeAllConnections(); server.close(resolve) })
+const json = (res, body) => { res.setHeader('Content-Type', 'application/json'); res.end(JSON.stringify(body)) }
+function client(url) {
+  const exports = {}
+  new Function('require', 'exports', 'useRuntimeConfig', 'createError', compiled)(require, exports,
+    () => ({ comfyControlUrl: url, comfyControlToken: 'test-secret' }),
+    info => Object.assign(new Error(info.statusMessage), info))
+  let active = false
+  exports.registerSharedGpuWork(() => active)
+  return { api: exports, active: value => { active = value } }
+}
+
+test('two independent studio clients share the real HTTP gateway without sharing job details', async t => {
+  let busy = false
+  const mutations = []
+  const comfy = http.createServer((req,res) => {
+    if (req.method === 'GET') return json(res, { queue_running: busy ? [[1, 'private-id', { prompt: 'private prompt from other site' }]] : [], queue_pending: [] })
+    mutations.push(req.url)
+    req.resume(); req.on('end', () => json(res, { prompt_id: 'opaque-result' }))
+  })
+  const comfyPort = await listen(comfy)
+  t.after(() => close(comfy))
+  const reservation = createGpuReservation({ idle: async () => !busy })
+  const control = http.createServer(async (req,res) => {
+    let text = ''; for await (const chunk of req) text += chunk
+    const body = JSON.parse(text || '{}')
+    assert.equal(req.headers.authorization, 'Bearer test-secret')
+    const op = req.url.split('/').at(-1)
+    json(res, await reservation[op](body.ticket || body.token))
+  })
+  const controlPort = await listen(control)
+  t.after(() => close(control))
+  const proxy = createGpuProxy({ target: () => comfyPort, reservation, authorized: req => req.headers.authorization === 'Bearer test-secret' })
+  const proxyPort = await listen(proxy)
+  t.after(() => close(proxy))
+  const a = client(`http://127.0.0.1:${controlPort}`)
+  const b = client(`http://127.0.0.1:${controlPort}`)
+  const grants = await Promise.all([a.api.acquireSharedGpu(), b.api.acquireSharedGpu()])
+  assert.deepEqual(grants, [true, false])
+  a.active(true)
+  await a.api.maintainSharedGpu()
+  assert.equal((await fetch(`http://127.0.0.1:${proxyPort}/prompt`, { method:'POST', headers:a.api.sharedGpuHeaders(), body:'{}' })).status, 200)
+  assert.equal((await fetch(`http://127.0.0.1:${proxyPort}/interrupt`, { method:'POST', headers:{Authorization:'Bearer test-secret'}, body:'{}' })).status, 409)
+  assert.deepEqual(mutations, ['/prompt'])
+  busy = true
+  const queue = await (await fetch(`http://127.0.0.1:${proxyPort}/queue`)).json()
+  assert.deepEqual(queue, { queue_running: [null], queue_pending: [] })
+  assert.equal(JSON.stringify(queue).includes('private'), false)
+  assert.deepEqual(await (await fetch(`http://127.0.0.1:${proxyPort}/history`)).json(), {})
+  assert.equal(await b.api.acquireSharedGpu(), false)
+  // Saving still owns the reservation even after Comfy reports idle.
+  busy = false
+  assert.equal(await b.api.acquireSharedGpu(), false)
+  a.active(false); await a.api.maintainSharedGpu()
+  assert.equal(await a.api.acquireSharedGpu(), false, 'waiting B must run before another A job')
+  assert.equal(await b.api.acquireSharedGpu(), true)
+  assert.equal((await fetch(`http://127.0.0.1:${proxyPort}/prompt`, { method:'POST', headers:b.api.sharedGpuHeaders(), body:'{}' })).status, 200)
+  // A reset cannot even enter its local clear routine while B owns the device.
+  await assert.rejects(a.api.withSharedGpuReset(async () => assert.fail('must not clear jobs')), {statusCode:409})
+  b.active(false); await b.api.maintainSharedGpu()
+})
+
+test('missing coordinator cannot silently fall back to uncoordinated generation', async () => {
+  const c = client('http://127.0.0.1:1')
+  assert.equal(await c.api.acquireSharedGpu(), false)
+  await assert.rejects(c.api.assertSharedGpu(), {statusCode:409})
+  assert.throws(() => c.api.sharedGpuHeaders(), {statusCode:409})
+})