Coordinate shared GPU access across studio instances

tests/gpu-reservation.test.mjs

@@ -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)
+})