diff --git a/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt b/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt index 7ebdcdcc..832647ab 100644 --- a/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt +++ b/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt @@ -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()) { - NvidiaResourceProbe(pidSupplier = { modelManager.currentPid() }) - } else { - UnavailableProbe + resourceProbe = when { + NvidiaResourceProbe.isAvailable() -> + NvidiaResourceProbe(pidSupplier = { modelManager.currentPid() }) + AmdResourceProbe.isAvailable() -> + AmdResourceProbe(pidSupplier = { modelManager.currentPid() }) + else -> UnavailableProbe } // Shutdown hook to kill the managed llama-server Runtime.getRuntime().addShutdownHook(Thread { diff --git a/infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/AmdResourceProbe.kt b/infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/AmdResourceProbe.kt new file mode 100644 index 00000000..2d1ac812 --- /dev/null +++ b/infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/AmdResourceProbe.kt @@ -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//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() + } +} diff --git a/testing/contracts/src/test/kotlin/com/correx/testing/contracts/model/inference/AmdResourceProbeTest.kt b/testing/contracts/src/test/kotlin/com/correx/testing/contracts/model/inference/AmdResourceProbeTest.kt new file mode 100644 index 00000000..8ccdc853 --- /dev/null +++ b/testing/contracts/src/test/kotlin/com/correx/testing/contracts/model/inference/AmdResourceProbeTest.kt @@ -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) + } +}