bare-operating-system/packages/bare-os-booter/lib/bare-os-qvac-gpu-probe.mjs

/**
 * Probe host GPUs for QVAC load ordering (highest device-local memory first).
 * Uses vulkaninfo when available; nvidia-smi / sysfs as weak fallbacks.
 */

/**
 * @typedef {{
 *   index: number,
 *   name: string,
 *   memoryBytes: number,
 *   deviceType: string,
 *   source: string
 * }} BareOsQvacGpuInfo
 */

/**
 * Run a short host command; return stdout (utf8) or null.
 * @param {string} file
 * @param {string[]} args
 * @param {{ timeoutMs?: number }} [opts]
 * @returns {Promise<string | null>}
 */
export async function bareOsQvacRunHostCommand(file, args, opts = {}) {
  const timeoutMs = Math.max(500, Math.min(30000, Number(opts.timeoutMs) || 8000))
  const argv = Array.isArray(args) ? args.map((x) => String(x)) : []

  /**
   * @param {{ spawn: Function }} sp
   */
  async function viaSpawn(sp) {
    return await new Promise((resolve) => {
      let settled = false
      /** @type {string[]} */
      const chunks = []
      let child
      try {
        child = sp.spawn(file, argv, {
          stdio: ['ignore', 'pipe', 'pipe']
        })
      } catch {
        resolve(null)
        return
      }
      const finish = (out) => {
        if (settled) return
        settled = true
        try {
          if (child && typeof child.kill === 'function') child.kill()
        } catch {
          /* ignore */
        }
        resolve(out)
      }
      const timer = setTimeout(() => finish(null), timeoutMs)
      try {
        if (child.stdout && typeof child.stdout.on === 'function') {
          child.stdout.on('data', (buf) => {
            chunks.push(Buffer.isBuffer(buf) ? buf.toString('utf8') : String(buf))
          })
        }
      } catch {
        /* ignore */
      }
      const onDone = (code) => {
        clearTimeout(timer)
        if (code === 0 || chunks.length) finish(chunks.join(''))
        else finish(null)
      }
      try {
        child.on('exit', onDone)
        child.on('error', () => {
          clearTimeout(timer)
          finish(null)
        })
      } catch {
        clearTimeout(timer)
        finish(null)
      }
    })
  }

  try {
    const mod = await import('bare-subprocess')
    const sp = mod && (mod.default || mod)
    if (sp && typeof sp.spawn === 'function') {
      const out = await viaSpawn(sp)
      if (out != null) return out
    }
  } catch {
    /* try node next */
  }

  try {
    const cp = await import('child_process')
    const execFile = cp.execFile || (cp.default && cp.default.execFile)
    if (typeof execFile !== 'function') return null
    return await new Promise((resolve) => {
      execFile(
        file,
        argv,
        { timeout: timeoutMs, maxBuffer: 8 * 1024 * 1024, encoding: 'utf8' },
        (err, stdout) => {
          if (err && !stdout) resolve(null)
          else resolve(stdout != null ? String(stdout) : null)
        }
      )
    })
  } catch {
    return null
  }
}

/**
 * Parse `vulkaninfo` text for GPUs + device-local heap sizes.
 * @param {string} text
 * @returns {BareOsQvacGpuInfo[]}
 */
export function bareOsQvacParseVulkaninfo(text) {
  const src = String(text || '')
  if (!src.trim()) return []

  /** @type {BareOsQvacGpuInfo[]} */
  const out = []
  // Split on GPU N: headers (vulkaninfo --summary and full dump).
  const parts = src.split(/(?=^GPU\d+:)/m)
  for (const part of parts) {
    const mIdx = /^GPU(\d+):/m.exec(part)
    if (!mIdx) continue
    const index = Number.parseInt(mIdx[1], 10)
    if (!Number.isFinite(index)) continue

    const nameMatch =
      /deviceName\s*=\s*(.+)$/m.exec(part) ||
      /deviceName\s*=\s*(.+)$/im.exec(part)
    const name = nameMatch ? String(nameMatch[1]).trim() : 'GPU' + index

    const typeMatch = /deviceType\s*=\s*(\S+)/m.exec(part)
    const deviceType = typeMatch ? String(typeMatch[1]).trim() : ''

    let memoryBytes = 0
    // Full vulkaninfo: memoryHeaps[i] blocks with DEVICE_LOCAL
    const heapBlocks = part.split(/(?=memoryHeaps\[\d+\]:)/i)
    for (const block of heapBlocks) {
      if (!/memoryHeaps\[\d+\]:/i.test(block)) continue
      const szMatch = /\bsize\s*=\s*(\d+)/i.exec(block)
      if (!szMatch) continue
      const sz = Number(szMatch[1])
      if (!Number.isFinite(sz) || sz <= 0) continue
      if (/DEVICE_LOCAL/i.test(block) || sz > memoryBytes) {
        memoryBytes = Math.max(memoryBytes, sz)
      }
    }
    // Alternate: "heapSize = N"
    if (!memoryBytes) {
      const hs = [...part.matchAll(/heapSize\s*=\s*(\d+)/gi)]
      for (const x of hs) {
        const sz = Number(x[1])
        if (Number.isFinite(sz) && sz > memoryBytes) memoryBytes = sz
      }
    }
    // Summary-only: no heaps — score by type so discrete still ranks above integrated.
    if (!memoryBytes) {
      if (/DISCRETE/i.test(deviceType)) memoryBytes = 8 * 1024 * 1024 * 1024
      else if (/INTEGRATED/i.test(deviceType)) memoryBytes = 512 * 1024 * 1024
      else memoryBytes = 256 * 1024 * 1024
    }

    out.push({
      index,
      name,
      memoryBytes,
      deviceType,
      source: 'vulkaninfo'
    })
  }
  return out
}

/**
 * Parse nvidia-smi CSV: name, memory.total (MiB).
 * Indices are NVIDIA order, NOT Vulkan — used only as a presence / size hint.
 * @param {string} text
 * @returns {{ name: string, memoryBytes: number }[]}
 */
export function bareOsQvacParseNvidiaSmi(text) {
  /** @type {{ name: string, memoryBytes: number }[]} */
  const out = []
  for (const line of String(text || '').split(/\r?\n/)) {
    const t = line.trim()
    if (!t) continue
    const parts = t.split(',').map((x) => x.trim())
    if (parts.length < 2) continue
    const name = parts[0]
    const mib = Number(parts[1])
    if (!name || !Number.isFinite(mib) || mib <= 0) continue
    out.push({ name, memoryBytes: Math.floor(mib * 1024 * 1024) })
  }
  return out
}

/**
 * Merge nvidia memory hints onto vulkan devices by fuzzy name match.
 * @param {BareOsQvacGpuInfo[]} vulkan
 * @param {{ name: string, memoryBytes: number }[]} nvidia
 */
export function bareOsQvacMergeNvidiaHints(vulkan, nvidia) {
  if (!vulkan.length || !nvidia.length) return vulkan
  return vulkan.map((g) => {
    const gName = g.name.toLowerCase()
    let best = null
    let bestScore = 0
    for (const n of nvidia) {
      const nName = n.name.toLowerCase()
      if (gName.includes(nName) || nName.includes(gName) || gName.includes('nvidia')) {
        const score = n.memoryBytes
        if (score > bestScore) {
          bestScore = score
          best = n
        }
      }
    }
    if (best && best.memoryBytes > g.memoryBytes) {
      return {
        ...g,
        memoryBytes: best.memoryBytes,
        source: g.source + '+nvidia-smi'
      }
    }
    return g
  })
}

/**
 * Sort GPUs for load attempts: highest memory first; discrete before integrated on ties.
 * @param {BareOsQvacGpuInfo[]} gpus
 * @returns {BareOsQvacGpuInfo[]}
 */
export function bareOsQvacRankGpus(gpus) {
  return [...gpus].sort((a, b) => {
    if (b.memoryBytes !== a.memoryBytes) return b.memoryBytes - a.memoryBytes
    const ad = /DISCRETE/i.test(a.deviceType) ? 1 : 0
    const bd = /DISCRETE/i.test(b.deviceType) ? 1 : 0
    if (bd !== ad) return bd - ad
    return a.index - b.index
  })
}

/**
 * Build ordered main-gpu candidates for auto mode.
 * @param {BareOsQvacGpuInfo[]} ranked
 * @returns {Array<string | number>}
 */
export function bareOsQvacMainGpuCandidates(ranked) {
  /** @type {Array<string | number>} */
  const out = []
  const seen = new Set()
  const push = (v) => {
    const k = String(v)
    if (seen.has(k)) return
    seen.add(k)
    out.push(v)
  }

  const vulkanBacked = ranked.filter((g) => String(g.source || '').includes('vulkaninfo'))
  if (vulkanBacked.length) {
    for (const g of vulkanBacked) push(g.index)
    push('dedicated')
    return out
  }

  // nvidia-smi indices ≠ Vulkan indices on hybrid laptops — prefer class + common slots.
  if (ranked.length) {
    push('dedicated')
    push(1)
    push(0)
    push(2)
    push(3)
    return out
  }

  push('dedicated')
  push(1)
  push(0)
  push(2)
  push(3)
  return out
}

/**
 * Probe host GPUs (best-effort). Never throws.
 * @param {{ runCommand?: typeof bareOsQvacRunHostCommand }} [opts]
 * @returns {Promise<BareOsQvacGpuInfo[]>}
 */
export async function bareOsQvacProbeGpus(opts = {}) {
  const run = opts.runCommand || bareOsQvacRunHostCommand
  /** @type {BareOsQvacGpuInfo[]} */
  let vulkan = []

  const full =
    (await run('vulkaninfo', [], { timeoutMs: 12000 })) ||
    (await run('vulkaninfo', ['--summary'], { timeoutMs: 8000 }))
  if (full) vulkan = bareOsQvacParseVulkaninfo(full)

  const smi = await run(
    'nvidia-smi',
    ['--query-gpu=name,memory.total', '--format=csv,noheader,nounits'],
    { timeoutMs: 5000 }
  )
  const nvidia = smi ? bareOsQvacParseNvidiaSmi(smi) : []
  if (vulkan.length && nvidia.length) {
    vulkan = bareOsQvacMergeNvidiaHints(vulkan, nvidia)
  } else if (!vulkan.length && nvidia.length) {
    // No Vulkan enumeration — synthesize placeholders; load path still tries dedicated + indices.
    vulkan = nvidia.map((n, i) => ({
      index: i,
      name: n.name,
      memoryBytes: n.memoryBytes,
      deviceType: 'PHYSICAL_DEVICE_TYPE_DISCRETE_GPU',
      source: 'nvidia-smi'
    }))
  }

  return bareOsQvacRankGpus(vulkan)
}