feat: correx-managed model lifecycle slice 4 — resource telemetry

Vendor-agnostic ResourceProbe (commons): NvidiaResourceProbe reads nvidia-smi
(injectable runner) for whole-device VRAM/util and /proc/<pid>/status for the
managed llama-server RSS, both fail-soft to null; UnavailableProbe for non-GPU
hosts / the static path. DefaultModelManager.currentPid() + LlamaProcess.pid
expose the managed pid. ServerMessage.ResourceStatus (resource.status,
NonEventMessage, all-nullable) pushed every 2.5s on the global stream plus one
in the initial snapshot. Live gauge only — never event-sourced (feeds no core
decision), so invariants #8/#9 hold. Main wires the NVIDIA probe on the managed
path when nvidia-smi is present, else UnavailableProbe.

Tests: csv parsing (single/multi/malformed), gpu+rss combine, fallbacks.

Plan: docs/plans/2026-05-31-model-lifecycle-management.md (slice 4 of 5).
This commit is contained in:
2026-06-01 12:36:38 +04:00
parent 7341ab578e
commit 7b1df95627
10 changed files with 273 additions and 0 deletions
@@ -49,6 +49,9 @@ 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.NvidiaResourceProbe
import com.correx.infrastructure.inference.commons.ResourceProbe
import com.correx.infrastructure.inference.commons.UnavailableProbe
import com.correx.core.inference.InferenceProvider
import com.correx.infrastructure.inference.llama.cpp.LlamaCppInferenceProvider
import com.correx.infrastructure.tools.FileEditConfig
@@ -83,6 +86,8 @@ fun main() {
val inferenceRouter: InferenceRouter
// Non-null only on the managed path; backs manual model swap/pin from clients.
var modelSwapper: ManagedInferenceRouter? = null
// Live GPU/RAM gauge; NVIDIA-backed on the managed path when nvidia-smi is present, else unavailable.
var resourceProbe: ResourceProbe = UnavailableProbe
if (correxConfig.models.isNotEmpty()) {
val settings = correxConfig.modelsSettings
@@ -108,6 +113,11 @@ fun main() {
val managedRouter = ManagedInferenceRouter(modelManager, descriptors, targetModelConfig.id)
inferenceRouter = managedRouter
modelSwapper = managedRouter
resourceProbe = if (NvidiaResourceProbe.isAvailable()) {
NvidiaResourceProbe(pidSupplier = { modelManager.currentPid() })
} else {
UnavailableProbe
}
// Shutdown hook to kill the managed llama-server
Runtime.getRuntime().addShutdownHook(Thread {
log.info("Shutdown: unloading managed model '{}'", targetModelConfig.id)
@@ -243,6 +253,7 @@ fun main() {
toolRegistry = toolRegistry,
sessionUndoService = sessionUndoService,
modelSwapper = modelSwapper,
resourceProbe = resourceProbe,
)
module.start()
log.info("==============================")
@@ -52,6 +52,10 @@ class ServerModule(
// The managed-model router when correx owns the local model process ([[models]] configured);
// null on the static-provider path. Backs the manual SwapModel / ClearModelPin operations.
val modelSwapper: com.correx.infrastructure.inference.commons.ManagedInferenceRouter? = null,
// Live GPU/RAM gauge pushed to clients on the global stream. Defaults to the always-unavailable
// probe (static path / non-GPU host); the managed path wires an NVIDIA-backed probe.
val resourceProbe: com.correx.infrastructure.inference.commons.ResourceProbe =
com.correx.infrastructure.inference.commons.UnavailableProbe,
// Long-lived scope owned by the module — backs the event-store subscription.
// SupervisorJob so one failure doesn't kill the whole module;
// Dispatchers.Default since the work is non-blocking and CPU-light.
@@ -283,6 +283,22 @@ sealed interface ServerMessage {
override val sessionSequence: Long? = null,
) : ServerMessage, NonEventMessage
/**
* Live resource gauge pushed periodically on the global stream (not event-derived). All fields
* are nullable: GPU fields are null on a non-NVIDIA host, `processRssMb` is null when no model
* is resident.
*/
@Serializable
@SerialName("resource.status")
data class ResourceStatus(
val gpuMemoryUsedMb: Long?,
val gpuMemoryTotalMb: Long?,
val gpuUtilizationPct: Int?,
val processRssMb: Long?,
override val sequence: Long? = null,
override val sessionSequence: Long? = null,
) : ServerMessage, NonEventMessage
@Serializable
@SerialName("protocol_error")
data class ProtocolError(
@@ -25,24 +25,39 @@ import com.correx.core.events.types.SessionId
import com.correx.core.events.types.StageId
import com.correx.core.inference.ProviderHealth
import com.correx.core.utils.TypeId
import com.correx.infrastructure.inference.commons.ResourceProbe
import com.correx.infrastructure.inference.commons.ResourceSnapshot
import io.ktor.server.websocket.DefaultWebSocketServerSession
import io.ktor.websocket.Frame
import io.ktor.websocket.readText
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.channels.BufferOverflow
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.channels.ClosedReceiveChannelException
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.delay
import kotlinx.coroutines.flow.SharedFlow
import kotlinx.coroutines.flow.onStart
import kotlinx.coroutines.flow.onSubscription
import kotlinx.coroutines.launch
import kotlinx.coroutines.withContext
import kotlinx.datetime.Clock
import org.slf4j.LoggerFactory
import java.util.UUID
private val log = LoggerFactory.getLogger(GlobalStreamHandler::class.java)
private const val BUFFER_CAPACITY = 1024
private const val RESOURCE_PUSH_INTERVAL_MS = 2500L
private const val BYTES_PER_MB = 1024L * 1024L
private fun ResourceSnapshot.toResourceStatus(): ServerMessage.ResourceStatus =
ServerMessage.ResourceStatus(
gpuMemoryUsedMb = gpu?.memoryUsedMb,
gpuMemoryTotalMb = gpu?.memoryTotalMb,
gpuUtilizationPct = gpu?.utilizationPct,
processRssMb = processRssBytes?.let { it / BYTES_PER_MB },
)
class GlobalStreamHandler(private val module: ServerModule) {
@@ -72,6 +87,10 @@ class GlobalStreamHandler(private val module: ServerModule) {
launch {
streamGlobal(module.eventStore, bridge, mapper, sendFrame)
}
// Live resource gauge (GPU/RAM) pushed periodically. Cancelled with the session scope.
launch {
streamResources(module.resourceProbe, sendFrame)
}
try {
for (frame in session.incoming) {
@@ -129,6 +148,22 @@ class GlobalStreamHandler(private val module: ServerModule) {
session.send(Frame.Text(ProtocolSerializer.encodeServerMessage(
ServerMessage.WorkflowList(workflows = workflows),
)))
// Initial resource gauge so a freshly-connected client renders VRAM/RAM immediately.
val snapshot = withContext(Dispatchers.IO) { module.resourceProbe.probe() }
session.send(Frame.Text(ProtocolSerializer.encodeServerMessage(snapshot.toResourceStatus())))
}
/** Pushes a [ServerMessage.ResourceStatus] every [RESOURCE_PUSH_INTERVAL_MS]. */
private suspend fun streamResources(
probe: ResourceProbe,
sendFrame: suspend (ServerMessage) -> Unit,
) {
while (true) {
delay(RESOURCE_PUSH_INTERVAL_MS)
val snapshot = withContext(Dispatchers.IO) { probe.probe() }
sendFrame(snapshot.toResourceStatus())
}
}
private suspend fun handleClientMessage(