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.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 {
|
||||
|
||||
Reference in New Issue
Block a user