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
@@ -55,6 +55,7 @@ import com.correx.infrastructure.InfrastructureModule
 import com.correx.infrastructure.inference.DefaultProviderRegistry
 import com.correx.infrastructure.inference.FirstAvailableRoutingStrategy
 import com.correx.infrastructure.inference.commons.ManagedInferenceRouter
+import com.correx.infrastructure.inference.commons.AmdResourceProbe
 import com.correx.infrastructure.inference.commons.NvidiaResourceProbe
 import com.correx.infrastructure.inference.commons.ResourceProbe
 import com.correx.infrastructure.inference.commons.UnavailableProbe
@@ -122,10 +123,12 @@ fun main() {
        val managedRouter = ManagedInferenceRouter(modelManager, descriptors, targetModelConfig.id)
        inferenceRouter = managedRouter
        modelSwapper = managedRouter
-        resourceProbe = if (NvidiaResourceProbe.isAvailable()) {
+        resourceProbe = when {
+            NvidiaResourceProbe.isAvailable() ->
                NvidiaResourceProbe(pidSupplier = { modelManager.currentPid() })
-        } else {
-            UnavailableProbe
+            AmdResourceProbe.isAvailable() ->
+                AmdResourceProbe(pidSupplier = { modelManager.currentPid() })
+            else -> UnavailableProbe
        }
        // Shutdown hook to kill the managed llama-server
        Runtime.getRuntime().addShutdownHook(Thread {
@@ -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()
+    }
+}
@@ -0,0 +1,69 @@
+package com.correx.testing.contracts.model.inference
+
+import com.correx.infrastructure.inference.commons.AmdResourceProbe
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Assertions.assertNull
+import org.junit.jupiter.api.Test
+
+class AmdResourceProbeTest {
+
+    // Real `rocm-smi --showmeminfo vram --showuse --csv` output (ROCm 4.0.0 / lib 7.8.0),
+    // including the low-power warning + blank line that precede the CSV.
+    private val realCsv = """
+        WARNING: AMD GPU device(s) is/are in a low-power state. Check power control/runtime_status
+
+        device,GPU use (%),VRAM Total Memory (B),VRAM Total Used Memory (B)
+        card0,2,17163091968,12186726400
+    """.trimIndent()
+
+    @Test
+    fun `parses real rocm-smi csv, converting bytes to MiB`() {
+        val gpu = AmdResourceProbe.parseGpu(realCsv)
+        assertEquals(16368L, gpu?.memoryTotalMb) // 17163091968 B / 1 MiB
+        assertEquals(11622L, gpu?.memoryUsedMb) // 12186726400 B / 1 MiB
+        assertEquals(2, gpu?.utilizationPct)
+    }
+
+    @Test
+    fun `tolerates reordered columns by matching header names`() {
+        val reordered = """
+            device,VRAM Total Used Memory (B),VRAM Total Memory (B),GPU use (%)
+            card0,12186726400,17163091968,2
+        """.trimIndent()
+        val gpu = AmdResourceProbe.parseGpu(reordered)
+        assertEquals(16368L, gpu?.memoryTotalMb)
+        assertEquals(11622L, gpu?.memoryUsedMb)
+        assertEquals(2, gpu?.utilizationPct)
+    }
+
+    @Test
+    fun `returns null when header or data is missing`() {
+        assertNull(AmdResourceProbe.parseGpu(""))
+        assertNull(AmdResourceProbe.parseGpu("WARNING: something\n\n"))
+        assertNull(AmdResourceProbe.parseGpu("device,GPU use (%),VRAM Total Memory (B),VRAM Total Used Memory (B)"))
+    }
+
+    @Test
+    fun `probe combines gpu csv and rss reader`() {
+        val probe = AmdResourceProbe(
+            pidSupplier = { 4242L },
+            rocmSmiRunner = { realCsv },
+            rssReader = { pid -> if (pid == 4242L) 2_097_152L else null },
+        )
+        val snapshot = probe.probe()
+        assertEquals(11622L, snapshot.gpu?.memoryUsedMb)
+        assertEquals(2_097_152L, snapshot.processRssBytes)
+    }
+
+    @Test
+    fun `probe reports gpu unavailable when rocm-smi yields null`() {
+        val probe = AmdResourceProbe(
+            pidSupplier = { null },
+            rocmSmiRunner = { null },
+            rssReader = { error("must not be called when pid is null") },
+        )
+        val snapshot = probe.probe()
+        assertNull(snapshot.gpu)
+        assertNull(snapshot.processRssBytes)
+    }
+}