82 lines
4.6 KiB
JavaScript
82 lines
4.6 KiB
JavaScript
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})
|
|
})
|