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