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:
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()) {
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 {