Coordinate shared GPU access across studio instances

scripts/comfy-host-agent.mjs

@@ -1,3 +1,5 @@
+import { createGpuReservation } from './gpu-reservation.mjs'
+import { createGpuProxy } from './gpu-proxy.mjs'
 import http from 'node:http'
 import net from 'node:net'
 import { execFile, spawn } from 'node:child_process'
@@ -141,26 +143,7 @@ let proxyTarget = 0
 
 function ensureProxyListening() {
   if (proxyServer) return
-  proxyServer = net.createServer((client) => {
-    const target = proxyTarget
-    if (!target) {
-      client.destroy()
-      return
-    }
-    const upstream = net.connect(target, '127.0.0.1')
-    const fail = () => {
-      try { client.destroy() } catch { /* ignore */ }
-      try { upstream.destroy() } catch { /* ignore */ }
-    }
-    client.on('error', fail)
-    upstream.on('error', fail)
-    client.once('data', (chunk) => {
-      if (isWorkHttp(chunk)) markWork()
-      upstream.write(chunk)
-      client.pipe(upstream)
-    })
-    upstream.pipe(client)
-  })
+  proxyServer = createGpuProxy({ target: () => proxyTarget, reservation: gpuReservation, authorized, markWork })
   proxyServer.on('error', (error) => {
     console.log(JSON.stringify({ src: 'comfy-host-agent', event: 'proxy-error', error: String(error.message || error) }))
   })
@@ -508,6 +491,7 @@ async function noteQueue(portNum) {
 }
 
 async function maybeIdleStop(healthyPort) {
+  if (gpuReservation.availability().busy) return
   if (healthyPort) {
     stoppedByAgent = false
     const queue = await noteQueue(healthyPort)
@@ -767,7 +751,7 @@ function readJson(req) {
 
 function purgeDesktopFiles(body) {
   const deleted = []
-  const allowSweep = body?.sweep !== false
+  const allowSweep = false // Never sweep another studio's files during a job cleanup.
   const names = [
     String(body?.imageName || body?.filename || ''),
     ...(Array.isArray(body?.imageNames) ? body.imageNames : [])
@@ -827,9 +811,30 @@ function purgeDesktopFiles(body) {
   return { ok: true, deleted: unique }
 }
 
-const server = http.createServer(async (req, res) => {
+const gpuReservation = createGpuReservation({ idle: async () => {
+  if ((await trainingLock()).busy) return false
+  const healthy = await syncProxy()
+  if (healthy) {
+    const queue = await fetchLocalQueue(healthy)
+    return queue.ok && queue.running === 0 && queue.pending === 0
+  }
+  // A stopped service may be reserved before waking; a silent live process may not.
+  return !(await processUp()) && !(await pythonMainUp().catch(() => true))
+} })
+
+async function handleControl(req, res) {
   if (!authorized(req)) return json(res, 401, { ok: false, error: 'unauthorized' })
   const url = new URL(req.url || '/', 'http://localhost')
+  if (req.method === 'POST' && url.pathname.startsWith('/gpu/')) {
+    const body = await readJson(req)
+    let result
+    if (url.pathname === '/gpu/acquire') result = await gpuReservation.acquire(body.ticket)
+    else if (url.pathname === '/gpu/renew') result = await gpuReservation.renew(body.token)
+    else if (url.pathname === '/gpu/release') result = await gpuReservation.release(body.token)
+    else if (url.pathname === '/gpu/cancel') result = await gpuReservation.cancel(body.ticket)
+    else return json(res, 404, { ok: false })
+    return json(res, 200, result)
+  }
   if (req.method === 'GET' && url.pathname === '/status') {
     // Answer immediately. MyMonitor allows ~500ms; probing ports + WMI here
     // made the A light stay red even while this process was running.
@@ -849,7 +854,8 @@ const server = http.createServer(async (req, res) => {
       idleMs,
       port: healthyPort || null,
       proxyPort,
-      training: lastTraining
+      gpu: gpuReservation.availability(),
+      training: { busy: lastTraining.busy }
     })
   }
   if (req.method === 'POST' && url.pathname === '/start') {
@@ -935,21 +941,21 @@ const server = http.createServer(async (req, res) => {
     const body = await readJson(req)
     return json(res, 200, purgeDesktopFiles(body))
   }
-  if (req.method === 'POST' && url.pathname === '/sweep') {
-    const body = await readJson(req)
-    const deleted = [
-      ...sweepStaleStudioInputs(body?.sweepMaxAgeMs),
-      ...sweepStaleStudioOutputs(body?.sweepPrefixes, body?.sweepSubfolders, body?.sweepMaxAgeMs)
-    ]
-    console.log(JSON.stringify({
-      src: 'comfy-host-agent',
-      event: 'sweep',
-      deleted: deleted.length,
-      files: deleted.map((path) => basename(path))
-    }))
-    return json(res, 200, { ok: true, deleted })
-  }
+  if (req.method === 'POST' && url.pathname === '/sweep') return json(res, 409, { ok: false, message: 'Global cleanup is disabled while studios share the GPU.' })
   json(res, 404, { ok: false, error: 'not found' })
+}
+
+const server = http.createServer(async (req, res) => {
+  if (!authorized(req)) return json(res, 401, { ok: false, error: 'unauthorized' })
+  try {
+    const path = new URL(req.url || '/', 'http://localhost').pathname
+    if (req.method === 'POST' && !path.startsWith('/gpu/')) {
+      await gpuReservation.permit(String(req.headers['x-aigen-gpu-lease'] || ''), () => handleControl(req, res))
+    } else await handleControl(req, res)
+  } catch (error) {
+    req.resume()
+    if (!res.headersSent) json(res, error.statusCode || 400, { ok: false, message: error.statusCode === 409 ? 'GPU is in use. Waiting for availability.' : 'GPU coordination request failed.' })
+  }
 })
 
 function logStartupSweep() {