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:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user