feat(infra): AMD/ROCm resource probe for the VRAM gauge

The resource gauge was NVIDIA-only — on an AMD/ROCm box Main fell back to UnavailableProbe, so the TUI showed no VRAM. Add AmdResourceProbe backed by `rocm-smi --showmeminfo vram --showuse --csv`: it skips the warning preamble, locates the header by column name (order/version tolerant), and converts the byte VRAM figures to MiB. Process RSS reuses /proc/<pid>/status (only populated for a managed-model pid). Main now selects NVIDIA → AMD → Unavailable. Tests cover the real ROCm 4.0 CSV + reordered columns.
2026-06-03 01:52:51 +04:00
parent c63929d79f
commit 00b08660bd
3 changed files with 176 additions and 4 deletions
@@ -0,0 +1,100 @@
+package com.correx.infrastructure.inference.commons
+
+import org.slf4j.LoggerFactory
+import java.io.File
+
+private val log = LoggerFactory.getLogger("AmdResourceProbe")
+
+/**
+ * [ResourceProbe] backed by `rocm-smi --showmeminfo vram --showuse --csv` for AMD GPU stats and
+ * `/proc/<pid>/status` for the managed model process RSS. Both reads fail soft to null, so a
+ * missing `rocm-smi`, a non-AMD host, or a dead pid never throws.
+ *
+ * `rocm-smi` prefixes its CSV with warning/blank lines, and the column order/names vary between
+ * ROCm versions — so the parser locates the header by its column names rather than by position.
+ * VRAM figures are reported in bytes and converted to MiB to match [GpuStatus].
+ *
+ * The command runner and RSS read are injectable so the parsing is unit-testable without a GPU.
+ *
+ * @param pidSupplier returns the managed `llama-server` pid, or null when no model is resident.
+ * @param rocmSmiRunner returns raw `rocm-smi --csv` stdout, or null when unavailable.
+ * @param rssReader returns the RSS in bytes for a pid, or null when unreadable.
+ */
+class AmdResourceProbe(
+    private val pidSupplier: () -> Long?,
+    private val rocmSmiRunner: () -> String? = ::runRocmSmi,
+    private val rssReader: (Long) -> Long? = ::readRssBytes,
+) : ResourceProbe {
+
+    override fun probe(): ResourceSnapshot {
+        val gpu = rocmSmiRunner()?.let { parseGpu(it) }
+        val rss = pidSupplier()?.let { rssReader(it) }
+        return ResourceSnapshot(gpu = gpu, processRssBytes = rss)
+    }
+
+    companion object {
+        private val QUERY = listOf(
+            "rocm-smi", "--showmeminfo", "vram", "--showuse", "--csv",
+        )
+        private const val BYTES_PER_MB = 1024L * 1024L
+        private const val BYTES_PER_KB = 1024L
+
+        /**
+         * Parses `rocm-smi --csv` output. Skips the warning/blank preamble, finds the header by its
+         * column names (so column order and version differences don't matter), reads the first card
+         * row, and converts the byte VRAM figures to MiB. Returns null on any shape mismatch so the
+         * caller surfaces "unavailable" rather than partial data.
+         *
+         * Example input (AMD GPU in low-power state warning elided):
+         * ```
+         * device,GPU use (%),VRAM Total Memory (B),VRAM Total Used Memory (B)
+         * card0,2,17163091968,12186726400
+         * ```
+         */
+        fun parseGpu(csv: String): GpuStatus? = runCatching {
+            val lines = csv.lineSequence().map { it.trim() }.filter { it.isNotEmpty() }.toList()
+            val headerIdx = lines.indexOfFirst {
+                it.startsWith("device,") && it.contains("VRAM Total", ignoreCase = true)
+            }
+            require(headerIdx >= 0 && headerIdx + 1 < lines.size)
+
+            val header = lines[headerIdx].split(',').map { it.trim() }
+            val row = lines[headerIdx + 1].split(',').map { it.trim() }
+
+            val usedCol = header.indexOfFirst { it.contains("VRAM", true) && it.contains("Used", true) }
+            val totalCol = header.indexOfFirst { it.contains("VRAM", true) && it.contains("Total Memory", true) }
+            val useCol = header.indexOfFirst { it.contains("use (%)", true) }
+            require(usedCol >= 0 && totalCol >= 0)
+
+            val util = useCol.takeIf { it >= 0 }?.let { row.getOrNull(it)?.toIntOrNull() } ?: 0
+            GpuStatus(
+                memoryUsedMb = row[usedCol].toLong() / BYTES_PER_MB,
+                memoryTotalMb = row[totalCol].toLong() / BYTES_PER_MB,
+                utilizationPct = util,
+            )
+        }.getOrNull()
+
+        /** True if `rocm-smi` runs successfully — used to pick the probe at boot. */
+        fun isAvailable(): Boolean = runCatching {
+            val proc = ProcessBuilder("rocm-smi", "--showid")
+                .redirectErrorStream(true)
+                .start()
+            proc.waitFor() == 0
+        }.getOrElse { false }
+
+        private fun runRocmSmi(): String? = runCatching {
+            val proc = ProcessBuilder(QUERY).redirectErrorStream(true).start()
+            val out = proc.inputStream.bufferedReader().readText()
+            if (proc.waitFor() == 0) out else null
+        }.onFailure { log.debug("rocm-smi query failed: {}", it.message) }.getOrNull()
+
+        private fun readRssBytes(pid: Long): Long? = runCatching {
+            val statusFile = File("/proc/$pid/status")
+            if (!statusFile.exists()) return null
+            val vmRss = statusFile.readLines().firstOrNull { it.startsWith("VmRSS:") } ?: return null
+            // Format: "VmRSS:\t   12345 kB"
+            val kb = vmRss.removePrefix("VmRSS:").trim().substringBefore(' ').toLongOrNull() ?: return null
+            kb * BYTES_PER_KB
+        }.onFailure { log.debug("RSS read failed for pid {}: {}", pid, it.message) }.getOrNull()
+    }
+}