Coordinate shared GPU access across studio instances
This commit is contained in:
Towsty
@@ -0,0 +1,75 @@
|
||||
import test from 'node:test'
|
||||
import assert from 'node:assert/strict'
|
||||
import { createGpuReservation, GPU_WAIT_MESSAGE } from '../scripts/gpu-reservation.mjs'
|
||||
const a = 'aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa'
|
||||
const b = 'bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb'
|
||||
|
||||
test('simultaneous sites receive exactly one reservation', async () => {
|
||||
const gpu = createGpuReservation({ idle: async () => true })
|
||||
const replies = await Promise.all([gpu.acquire(a), gpu.acquire(b)])
|
||||
assert.equal(replies.filter(reply => reply.acquired).length, 1)
|
||||
assert.deepEqual(replies[1], { acquired: false, message: GPU_WAIT_MESSAGE, retryAfterMs: 2500 })
|
||||
assert.deepEqual(gpu.availability(), { busy: true })
|
||||
})
|
||||
test('another site cannot renew, release, or execute with an invalid token', async () => {
|
||||
const gpu = createGpuReservation({ idle: async () => true })
|
||||
const first = await gpu.acquire(a)
|
||||
assert.deepEqual(await gpu.renew(b), { renewed: false })
|
||||
assert.deepEqual(await gpu.release(b), { released: false })
|
||||
await assert.rejects(gpu.permit(b, () => assert.fail('must not execute')), { statusCode: 409 })
|
||||
assert.equal((await gpu.acquire(b)).acquired, false)
|
||||
assert.equal(await gpu.permit(first.token, () => 'allowed'), 'allowed')
|
||||
})
|
||||
test('release never grants another site while Comfy still has work', async () => {
|
||||
let idle = true
|
||||
const gpu = createGpuReservation({ idle: async () => idle })
|
||||
const first = await gpu.acquire(a)
|
||||
idle = false
|
||||
await gpu.release(first.token)
|
||||
assert.equal((await gpu.acquire(b)).acquired, false)
|
||||
idle = true
|
||||
assert.equal((await gpu.acquire(b)).acquired, true)
|
||||
await assert.rejects(gpu.permit(first.token, () => {}), { statusCode: 409 })
|
||||
})
|
||||
test('expired reservations recover only after confirmed GPU idle', async () => {
|
||||
let time = 0, idle = true
|
||||
const gpu = createGpuReservation({ idle: async () => idle, now: () => time, ttlMs: 100 })
|
||||
const first = await gpu.acquire(a)
|
||||
time = 101; idle = false
|
||||
assert.equal((await gpu.acquire(b)).acquired, false)
|
||||
assert.equal((await gpu.renew(first.token)).renewed, false)
|
||||
idle = true
|
||||
assert.equal((await gpu.acquire(b)).acquired, true)
|
||||
})
|
||||
test('an unanswered idle probe fails closed', async () => {
|
||||
const gpu = createGpuReservation({ idle: async () => null })
|
||||
assert.equal((await gpu.acquire(a)).acquired, false)
|
||||
})
|
||||
test('a waiting site gets the next reservation before the previous owner can reacquire', async () => {
|
||||
const gpu = createGpuReservation({ idle: async () => true })
|
||||
const first = await gpu.acquire(a)
|
||||
await gpu.acquire(b)
|
||||
await gpu.release(first.token)
|
||||
assert.equal((await gpu.acquire(a)).acquired, false)
|
||||
assert.equal((await gpu.acquire(b)).acquired, true)
|
||||
})
|
||||
test('an expired reservation cannot transfer while submission is in flight', async () => {
|
||||
let time = 0, finish
|
||||
const gpu = createGpuReservation({ idle: async () => true, now: () => time, ttlMs: 100 })
|
||||
const first = await gpu.acquire(a)
|
||||
const work = gpu.permit(first.token, () => new Promise(resolve => { finish = resolve }))
|
||||
await new Promise(resolve => setImmediate(resolve))
|
||||
time = 101
|
||||
assert.equal((await gpu.acquire(b)).acquired, false)
|
||||
finish(); await work
|
||||
assert.equal((await gpu.acquire(b)).acquired, true)
|
||||
})
|
||||
test('disconnected waiters expire so they cannot block the device forever', async () => {
|
||||
let time = 0
|
||||
const gpu = createGpuReservation({ idle: async () => true, now: () => time, ttlMs: 100, ticketTtlMs: 50 })
|
||||
const first = await gpu.acquire(a)
|
||||
await gpu.acquire(b)
|
||||
time = 51
|
||||
await gpu.release(first.token)
|
||||
assert.equal((await gpu.acquire(a)).acquired, true)
|
||||
})
|
||||
@@ -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})
|
||||
})
|
||||
@@ -27,6 +27,7 @@ function fixture({ rows = [], lives = [], queue = { running: 0, pending: 0 }, hi
|
||||
const timers = []
|
||||
const started = []
|
||||
const modules = {
|
||||
'~/server/utils/sharedGpu': { withSharedGpuStart: async fn => fn(), maintainSharedGpu: async () => {}, acquireSharedGpu: async () => false, sharedGpuWaitReason: () => 'GPU is in use. Waiting for availability.' },
|
||||
'~/server/utils/jobs': {
|
||||
listJobs: () => lives,
|
||||
getJob: id => lives.find(job => job.id === id),
|
||||
|
||||
Reference in New Issue
Block a user