70 lines
2.8 KiB
JavaScript
70 lines
2.8 KiB
JavaScript
import { randomUUID } from 'node:crypto'
|
|
|
|
export const GPU_WAIT_MESSAGE = 'GPU is in use. Waiting for availability.'
|
|
|
|
/** Device-local arbitration. Tickets/tokens are opaque; no studio or job metadata is stored. */
|
|
export function createGpuReservation({ idle, now = Date.now, ttlMs = 60_000, ticketTtlMs = 30_000 } = {}) {
|
|
if (typeof idle !== 'function') throw new Error('An authoritative GPU idle probe is required')
|
|
let owner = null
|
|
let gate = Promise.resolve()
|
|
let inFlight = 0
|
|
const waiting = new Map()
|
|
const serialize = fn => {
|
|
const run = gate.then(fn)
|
|
gate = run.catch(() => {})
|
|
return run
|
|
}
|
|
const valid = token => Boolean(owner && !owner.released && owner.token === token && owner.until > now())
|
|
const busy = () => ({ acquired: false, message: GPU_WAIT_MESSAGE, retryAfterMs: 2500 })
|
|
const prune = () => {
|
|
for (const [ticket, seen] of waiting) if (now() - seen >= ticketTtlMs) waiting.delete(ticket)
|
|
}
|
|
return {
|
|
isOwner: valid,
|
|
acquire(ticket) {
|
|
return serialize(async () => {
|
|
if (typeof ticket !== 'string' || !/^[a-zA-Z0-9-]{16,80}$/.test(ticket)) throw new Error('Invalid reservation ticket')
|
|
prune()
|
|
if (owner?.ticket === ticket && valid(owner.token)) {
|
|
owner.until = now() + ttlMs
|
|
return { acquired: true, token: owner.token, ttlMs }
|
|
}
|
|
waiting.set(ticket, now())
|
|
if (owner && (valid(owner.token) || inFlight > 0)) return busy()
|
|
if (waiting.keys().next().value !== ticket) return busy()
|
|
// Unknown/offline is not idle; the adapter must explicitly account for a stopped GPU.
|
|
if (await idle() !== true) return busy()
|
|
owner = { ticket, token: randomUUID(), until: now() + ttlMs, released: false }
|
|
waiting.delete(ticket)
|
|
return { acquired: true, token: owner.token, ttlMs }
|
|
})
|
|
},
|
|
renew(token) {
|
|
return serialize(() => {
|
|
if (!valid(token)) return { renewed: false }
|
|
owner.until = now() + ttlMs
|
|
return { renewed: true, ttlMs }
|
|
})
|
|
},
|
|
release(token) {
|
|
return serialize(() => {
|
|
if (!owner || owner.token !== token) return { released: false }
|
|
owner.released = true
|
|
// The next acquire still verifies that Comfy has drained before granting access.
|
|
return { released: true }
|
|
})
|
|
},
|
|
cancel(ticket) { return serialize(() => { waiting.delete(ticket); return { cancelled: true } }) },
|
|
async permit(token, operation) {
|
|
const accepted = await serialize(() => {
|
|
if (!valid(token)) return false
|
|
inFlight += 1
|
|
return true
|
|
})
|
|
if (!accepted) throw Object.assign(new Error(GPU_WAIT_MESSAGE), { statusCode: 409 })
|
|
try { return await operation() } finally { inFlight -= 1 }
|
|
},
|
|
availability() { return { busy: Boolean(owner && (valid(owner.token) || inFlight > 0)) } }
|
|
}
|
|
}
|