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.
This commit is contained in:
@@ -55,6 +55,7 @@ import com.correx.infrastructure.InfrastructureModule
|
|||||||
import com.correx.infrastructure.inference.DefaultProviderRegistry
|
import com.correx.infrastructure.inference.DefaultProviderRegistry
|
||||||
import com.correx.infrastructure.inference.FirstAvailableRoutingStrategy
|
import com.correx.infrastructure.inference.FirstAvailableRoutingStrategy
|
||||||
import com.correx.infrastructure.inference.commons.ManagedInferenceRouter
|
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.NvidiaResourceProbe
|
||||||
import com.correx.infrastructure.inference.commons.ResourceProbe
|
import com.correx.infrastructure.inference.commons.ResourceProbe
|
||||||
import com.correx.infrastructure.inference.commons.UnavailableProbe
|
import com.correx.infrastructure.inference.commons.UnavailableProbe
|
||||||
@@ -122,10 +123,12 @@ fun main() {
|
|||||||
val managedRouter = ManagedInferenceRouter(modelManager, descriptors, targetModelConfig.id)
|
val managedRouter = ManagedInferenceRouter(modelManager, descriptors, targetModelConfig.id)
|
||||||
inferenceRouter = managedRouter
|
inferenceRouter = managedRouter
|
||||||
modelSwapper = managedRouter
|
modelSwapper = managedRouter
|
||||||
resourceProbe = if (NvidiaResourceProbe.isAvailable()) {
|
resourceProbe = when {
|
||||||
NvidiaResourceProbe(pidSupplier = { modelManager.currentPid() })
|
NvidiaResourceProbe.isAvailable() ->
|
||||||
} else {
|
NvidiaResourceProbe(pidSupplier = { modelManager.currentPid() })
|
||||||
UnavailableProbe
|
AmdResourceProbe.isAvailable() ->
|
||||||
|
AmdResourceProbe(pidSupplier = { modelManager.currentPid() })
|
||||||
|
else -> UnavailableProbe
|
||||||
}
|
}
|
||||||
// Shutdown hook to kill the managed llama-server
|
// Shutdown hook to kill the managed llama-server
|
||||||
Runtime.getRuntime().addShutdownHook(Thread {
|
Runtime.getRuntime().addShutdownHook(Thread {
|
||||||
|
|||||||
+100
@@ -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()
|
||||||
|
}
|
||||||
|
}
|
||||||
+69
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user