107 lines
4.0 KiB
TypeScript
107 lines
4.0 KiB
TypeScript
import { randomUUID } from 'node:crypto'
|
|
|
|
const ticket = randomUUID()
|
|
let lease = ''
|
|
let expiresAt = 0
|
|
let lost = false
|
|
let starting = 0
|
|
let resetting = false
|
|
let gate: Promise<unknown> = Promise.resolve()
|
|
let localWork: () => boolean = () => true
|
|
let reason = 'GPU is in use. Waiting for availability.'
|
|
|
|
function settings() {
|
|
const config = useRuntimeConfig()
|
|
return {
|
|
url: String(config.comfyControlUrl || process.env.COMFY_CONTROL_URL || '').replace(/\/$/, ''),
|
|
token: String(config.comfyControlToken || process.env.COMFY_CONTROL_TOKEN || '')
|
|
}
|
|
}
|
|
export function sharedGpuConfigured() { return Boolean(settings().url) }
|
|
export function sharedGpuWaitReason() { return reason }
|
|
export function registerSharedGpuWork(check: () => boolean) { localWork = check }
|
|
|
|
async function request(path: string, body: Record<string, string>) {
|
|
const { url, token } = settings()
|
|
const response = await fetch(`${url}/gpu/${path}`, {
|
|
method: 'POST', headers: { 'Content-Type': 'application/json', ...(token ? { Authorization: `Bearer ${token}` } : {}) },
|
|
body: JSON.stringify(body), signal: AbortSignal.timeout(15_000)
|
|
})
|
|
if (!response.ok) throw new Error('GPU coordinator unavailable')
|
|
return await response.json() as { acquired?: boolean; token?: string; ttlMs?: number; renewed?: boolean; released?: boolean }
|
|
}
|
|
function serialized<T>(fn: () => Promise<T>): Promise<T> {
|
|
const run = gate.then(fn)
|
|
gate = run.catch(() => {})
|
|
return run
|
|
}
|
|
|
|
export function acquireSharedGpu() {
|
|
return serialized(async () => {
|
|
if (!sharedGpuConfigured()) return true
|
|
if (lost) return false
|
|
if (lease && Date.now() < expiresAt) return true
|
|
if (lease) { lost = true; return false }
|
|
try {
|
|
const result = await request('acquire', { ticket })
|
|
if (!result.acquired || !result.token || !Number.isFinite(result.ttlMs) || Number(result.ttlMs) < 10_000) {
|
|
reason = 'GPU is in use. Waiting for availability.'
|
|
return false
|
|
}
|
|
lease = result.token
|
|
expiresAt = Date.now() + Math.max(1000, Number(result.ttlMs) - 5000)
|
|
return true
|
|
} catch {
|
|
reason = 'GPU coordinator unavailable. Waiting to reconnect.'
|
|
return false
|
|
}
|
|
})
|
|
}
|
|
|
|
export async function withSharedGpuStart<T>(fn: () => Promise<T>, blocked: () => Promise<T>): Promise<T> {
|
|
starting += 1
|
|
try { return await (!resetting && await acquireSharedGpu() ? fn() : blocked()) }
|
|
finally { starting -= 1 }
|
|
}
|
|
|
|
export async function assertSharedGpu() {
|
|
if (!await acquireSharedGpu()) throw createError({ statusCode: 409, statusMessage: reason, data: { code: 'GPU_BUSY' } })
|
|
}
|
|
|
|
export function sharedGpuHeaders() {
|
|
if (!sharedGpuConfigured()) return {} as Record<string, string>
|
|
if (!lease || lost || Date.now() >= expiresAt) {
|
|
throw createError({ statusCode: 409, statusMessage: 'GPU reservation is unavailable. Waiting to reconnect.', data: { code: 'GPU_BUSY' } })
|
|
}
|
|
const { token } = settings()
|
|
return { 'x-aigen-gpu-lease': lease, ...(token ? { Authorization: `Bearer ${token}` } : {}) }
|
|
}
|
|
|
|
export function maintainSharedGpu() {
|
|
return serialized(async () => {
|
|
if (!sharedGpuConfigured()) return
|
|
if (!starting && !resetting && !localWork()) {
|
|
if (lease) {
|
|
try { await request('release', { token: lease }) } catch { /* expires if the host is unreachable */ }
|
|
}
|
|
lease = ''; expiresAt = 0; lost = false
|
|
return
|
|
}
|
|
if (!lease || lost) return
|
|
try {
|
|
const result = await request('renew', { token: lease })
|
|
if (!result.renewed || !Number.isFinite(result.ttlMs) || Number(result.ttlMs) < 10_000) { lost = true; return }
|
|
expiresAt = Date.now() + Math.max(1000, Number(result.ttlMs) - 5000)
|
|
} catch {
|
|
if (Date.now() >= expiresAt) lost = true
|
|
}
|
|
})
|
|
}
|
|
|
|
export async function withSharedGpuReset<T>(fn: () => Promise<T>) {
|
|
if (resetting) throw createError({ statusCode: 409, statusMessage: 'GPU reset is already in progress.' })
|
|
resetting = true
|
|
try { await assertSharedGpu(); return await fn() }
|
|
finally { resetting = false }
|
|
}
|