From 382194b49869e621887ff6194a84769cab9e9ca8 Mon Sep 17 00:00:00 2001 From: kami Date: Mon, 1 Jun 2026 11:35:51 +0400 Subject: [PATCH] =?UTF-8?q?feat:=20correx-managed=20model=20lifecycle=20sl?= =?UTF-8?q?ice=202=20=E2=80=94=20per-stage=20model=20selection?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit StageConfig.modelId; InferenceRouter gains a backward-compatible model-aware route() overload (default ignores modelId, so static routers and test fakes are unaffected). New ManagedInferenceRouter (infra commons) resolves a target model per stage (stage.modelId > capability match > default) and ensureLoaded()s it before routing, returning the live managed provider — its id changes with the resident model, so there is no stale health cache to invalidate on swap. ModelManager.ensureLoaded default delegates to the idempotent load(). Main wires the managed router on the [[models]] path; the static path keeps DefaultInferenceRouter. Plan: docs/plans/2026-05-31-model-lifecycle-management.md (slice 2 of 5). --- .../kotlin/com/correx/apps/server/Main.kt | 9 +- .../correx/core/inference/InferenceRouter.kt | 12 ++ .../orchestration/SessionOrchestrator.kt | 2 +- .../core/transitions/graph/StageConfig.kt | 1 + .../commons/ManagedInferenceRouter.kt | 70 +++++++++ .../inference/commons/ModelManager.kt | 7 + testing/contracts/build.gradle | 1 + .../inference/ManagedInferenceRouterTest.kt | 140 ++++++++++++++++++ 8 files changed, 240 insertions(+), 2 deletions(-) create mode 100644 infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/ManagedInferenceRouter.kt create mode 100644 testing/contracts/src/test/kotlin/com/correx/testing/contracts/model/inference/ManagedInferenceRouterTest.kt diff --git a/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt b/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt index 32bc9344..8498fb31 100644 --- a/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt +++ b/apps/server/src/main/kotlin/com/correx/apps/server/Main.kt @@ -16,6 +16,7 @@ import com.correx.core.inference.CapabilityScore import com.correx.core.inference.DefaultInferenceRouter import com.correx.core.inference.InferenceProjector import com.correx.core.inference.InferenceRepository +import com.correx.core.inference.InferenceRouter import com.correx.core.inference.ModelCapability import com.correx.core.kernel.orchestration.DefaultOrchestrationReducer import com.correx.core.kernel.orchestration.DefaultSessionOrchestrator @@ -47,6 +48,7 @@ import com.correx.core.toolintent.rules.PathContainmentRule 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.core.inference.InferenceProvider import com.correx.infrastructure.inference.llama.cpp.LlamaCppInferenceProvider import com.correx.infrastructure.tools.FileEditConfig @@ -78,10 +80,12 @@ fun main() { // Static path: [[models]] absent → legacy [[providers]] / env-var behavior val firstProvider: InferenceProvider val infraRegistry: DefaultProviderRegistry + val inferenceRouter: InferenceRouter if (correxConfig.models.isNotEmpty()) { val settings = correxConfig.modelsSettings val modelManager = InfrastructureModule.createModelManager(settings, eventStore, eventDispatcher) + val descriptors = correxConfig.models.map { InfrastructureModule.modelConfigToDescriptor(it) } val targetId = settings.defaultModel ?: correxConfig.models.first().id val targetModelConfig = correxConfig.models.firstOrNull { it.id == targetId } ?: correxConfig.models.first() @@ -97,6 +101,9 @@ fun main() { else -> { log.error("Unknown provider type: {}", providerConfig.type); null } } }.forEach { infraRegistry.register(it) } + // Managed router owns per-stage model selection: it ensures the resolved model is resident + // before routing (stage.modelId > capability match > default), evicting the prior model. + inferenceRouter = ManagedInferenceRouter(modelManager, descriptors, targetModelConfig.id) // Shutdown hook to kill the managed llama-server Runtime.getRuntime().addShutdownHook(Thread { log.info("Shutdown: unloading managed model '{}'", targetModelConfig.id) @@ -110,6 +117,7 @@ fun main() { require(staticProviders.isNotEmpty()) { "At least one provider must be configured" } firstProvider = staticProviders.first() infraRegistry = InfrastructureModule.createProviderRegistry(staticProviders) + inferenceRouter = DefaultInferenceRouter(infraRegistry, FirstAvailableRoutingStrategy()) } val modelId = firstProvider.id.value @@ -161,7 +169,6 @@ fun main() { ), ) - val inferenceRouter = DefaultInferenceRouter(infraRegistry, FirstAvailableRoutingStrategy()) val engines = OrchestratorEngines( transitionResolver = DefaultTransitionResolver { condition, ctx -> condition.evaluate(ctx) }, contextPackBuilder = DefaultContextPackBuilder(DefaultContextCompressor()), diff --git a/core/inference/src/main/kotlin/com/correx/core/inference/InferenceRouter.kt b/core/inference/src/main/kotlin/com/correx/core/inference/InferenceRouter.kt index 323971c3..b2436f95 100644 --- a/core/inference/src/main/kotlin/com/correx/core/inference/InferenceRouter.kt +++ b/core/inference/src/main/kotlin/com/correx/core/inference/InferenceRouter.kt @@ -27,6 +27,18 @@ interface InferenceRouter { stageId: StageId, requiredCapabilities: Set, ): InferenceProvider + + /** + * Resolves and selects a provider for a stage, optionally pinned to a specific model id. + * Implementations that own model lifecycle (e.g. a managed router) use [modelId] to select + * and load the target model before routing. The default implementation ignores [modelId] and + * falls back to capability-based selection, so static routers remain unaffected. + */ + suspend fun route( + stageId: StageId, + requiredCapabilities: Set, + modelId: String?, + ): InferenceProvider = route(stageId, requiredCapabilities) } class NoEligibleProviderException( diff --git a/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/SessionOrchestrator.kt b/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/SessionOrchestrator.kt index a5ee4cfd..f5d9fde4 100644 --- a/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/SessionOrchestrator.kt +++ b/core/kernel/src/main/kotlin/com/correx/core/kernel/orchestration/SessionOrchestrator.kt @@ -753,7 +753,7 @@ abstract class SessionOrchestrator( timeoutMs: Long, responseFormat: ResponseFormat = ResponseFormat.Text, ): InferenceResult { - val provider = inferenceRouter.route(stageId, stageConfig.requiredCapabilities) + val provider = inferenceRouter.route(stageId, stageConfig.requiredCapabilities, stageConfig.modelId) log.debug( "[Orchestrator] inference session={} stage={} provider={} timeoutMs={}", sessionId.value, stageId.value, provider.id.value, timeoutMs, diff --git a/core/transitions/src/main/kotlin/com/correx/core/transitions/graph/StageConfig.kt b/core/transitions/src/main/kotlin/com/correx/core/transitions/graph/StageConfig.kt index 5082c3b2..87f9b84d 100644 --- a/core/transitions/src/main/kotlin/com/correx/core/transitions/graph/StageConfig.kt +++ b/core/transitions/src/main/kotlin/com/correx/core/transitions/graph/StageConfig.kt @@ -7,6 +7,7 @@ import com.correx.core.inference.ModelCapability data class StageConfig( val requiredCapabilities: Set = emptySet(), + val modelId: String? = null, val tokenBudget: Int = 4096, val generationConfig: GenerationConfig = GenerationConfig( temperature = 0.7, diff --git a/infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/ManagedInferenceRouter.kt b/infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/ManagedInferenceRouter.kt new file mode 100644 index 00000000..fe375147 --- /dev/null +++ b/infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/ManagedInferenceRouter.kt @@ -0,0 +1,70 @@ +package com.correx.infrastructure.inference.commons + +import com.correx.core.events.types.StageId +import com.correx.core.inference.InferenceProvider +import com.correx.core.inference.InferenceRouter +import com.correx.core.inference.ModelCapability +import com.correx.core.inference.NoEligibleProviderException + +/** + * An [InferenceRouter] that owns the local model lifecycle: it resolves a single target model + * per stage and ensures it is resident (via [ModelManager.ensureLoaded]) before returning the + * managed provider to run inference against. + * + * Resolution precedence (per the model-lifecycle spec): + * ``` + * stage.modelId > capability match over configured models > defaultModelId + * ``` + * The manual-swap pin (a higher-precedence override) layers in on top of this in a later slice. + * + * Because the returned provider is the live managed provider (its [InferenceProvider.id] changes + * with the resident model), there is no stale router-level health cache to invalidate on swap — + * selection always reflects the model the manager just made resident. + */ +class ManagedInferenceRouter( + private val manager: ModelManager, + descriptors: List, + private val defaultModelId: String, +) : InferenceRouter { + + private val byId: Map = descriptors.associateBy { it.modelId } + private val ordered: List = descriptors + + override suspend fun route( + stageId: StageId, + requiredCapabilities: Set, + ): InferenceProvider = route(stageId, requiredCapabilities, null) + + override suspend fun route( + stageId: StageId, + requiredCapabilities: Set, + modelId: String?, + ): InferenceProvider { + val targetId = modelId + ?: capabilityMatch(requiredCapabilities) + ?: defaultModelId + val descriptor = byId[targetId] + ?: throw NoEligibleProviderException(stageId, requiredCapabilities) + return manager.ensureLoaded(descriptor) + } + + /** + * Selects the configured model whose declared capabilities cover every required capability, + * preferring the highest summed score over those required capabilities. Returns null when no + * capabilities are required or none of the models satisfy them (callers fall back to default). + */ + private fun capabilityMatch(required: Set): String? { + if (required.isEmpty()) return null + return ordered + .filter { descriptor -> + val declared = descriptor.capabilities.map { it.capability }.toSet() + declared.containsAll(required) + } + .maxByOrNull { descriptor -> + descriptor.capabilities + .filter { it.capability in required } + .sumOf { it.score } + } + ?.modelId + } +} diff --git a/infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/ModelManager.kt b/infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/ModelManager.kt index 03020324..8e2ca129 100644 --- a/infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/ModelManager.kt +++ b/infrastructure/inference/commons/src/main/kotlin/com/correx/infrastructure/inference/commons/ModelManager.kt @@ -16,6 +16,13 @@ interface ModelManager { */ suspend fun load(descriptor: ModelDescriptor): InferenceProvider + /** + * Ensures the model described by [descriptor] is the resident model, returning its provider. + * A no-op load if that model is already current; otherwise loads it (evicting the previous + * model). [load] is idempotent for the already-loaded model, so the default delegates to it. + */ + suspend fun ensureLoaded(descriptor: ModelDescriptor): InferenceProvider = load(descriptor) + /** * Unloads the currently loaded model. * diff --git a/testing/contracts/build.gradle b/testing/contracts/build.gradle index 55459ea1..cda5a4a9 100644 --- a/testing/contracts/build.gradle +++ b/testing/contracts/build.gradle @@ -15,6 +15,7 @@ dependencies { testImplementation(project(":testing:fixtures")) testImplementation(project(":core:inference")) testImplementation(project(":core:context")) + testImplementation(project(":infrastructure:inference:commons")) testFixturesImplementation(project(":core:events")) testFixturesImplementation(project(":core:sessions")) testFixturesImplementation(project(":core:artifacts-store")) diff --git a/testing/contracts/src/test/kotlin/com/correx/testing/contracts/model/inference/ManagedInferenceRouterTest.kt b/testing/contracts/src/test/kotlin/com/correx/testing/contracts/model/inference/ManagedInferenceRouterTest.kt new file mode 100644 index 00000000..ba83f9f6 --- /dev/null +++ b/testing/contracts/src/test/kotlin/com/correx/testing/contracts/model/inference/ManagedInferenceRouterTest.kt @@ -0,0 +1,140 @@ +package com.correx.testing.contracts.model.inference + +import com.correx.core.events.types.ProviderId +import com.correx.core.events.types.StageId +import com.correx.core.inference.CapabilityScore +import com.correx.core.inference.InferenceProvider +import com.correx.core.inference.ModelCapability +import com.correx.core.inference.NoEligibleProviderException +import com.correx.core.inference.ProviderHealth +import com.correx.infrastructure.inference.commons.ManagedInferenceRouter +import com.correx.infrastructure.inference.commons.ModelDescriptor +import com.correx.infrastructure.inference.commons.ModelManager +import com.correx.infrastructure.inference.commons.ResidencyMode +import com.correx.testing.fixtures.inference.MockInferenceProvider +import kotlinx.coroutines.runBlocking +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Test +import org.junit.jupiter.api.assertThrows + +class ManagedInferenceRouterTest { + + private val stage = StageId("s1") + + /** Records every (idempotent) load/ensureLoaded so we can assert which model was made resident. */ + private class FakeModelManager : ModelManager { + val loads = mutableListOf() + private var current: ModelDescriptor? = null + + override suspend fun load(descriptor: ModelDescriptor): InferenceProvider { + loads.add(descriptor.modelId) + current = descriptor + return MockInferenceProvider(id = ProviderId("managed:${descriptor.modelId}")) + } + + override suspend fun unload(modelId: String) { current = null } + override fun currentModel(): ModelDescriptor? = current + override fun healthCheck(): ProviderHealth = ProviderHealth.Healthy + } + + private fun desc(id: String, vararg caps: Pair): ModelDescriptor = + ModelDescriptor( + modelId = id, + modelPath = "/models/$id", + residencyMode = ResidencyMode.PERSISTENT, + capabilities = caps.map { CapabilityScore(it.first, it.second) }.toSet(), + ) + + @Test + fun `stage modelId takes precedence over capability and default`(): Unit = runBlocking { + val manager = FakeModelManager() + val router = ManagedInferenceRouter( + manager = manager, + descriptors = listOf(desc("default-model"), desc("coder", ModelCapability.Coding to 1.0)), + defaultModelId = "default-model", + ) + + val provider = router.route(stage, setOf(ModelCapability.Coding), modelId = "coder") + + assertEquals("managed:coder", provider.id.value) + assertEquals(listOf("coder"), manager.loads) + } + + @Test + fun `capability match selects a model that covers the required capabilities`(): Unit = runBlocking { + val manager = FakeModelManager() + val router = ManagedInferenceRouter( + manager = manager, + descriptors = listOf( + desc("generalist", ModelCapability.General to 1.0), + desc("coder", ModelCapability.Coding to 1.0), + ), + defaultModelId = "generalist", + ) + + val provider = router.route(stage, setOf(ModelCapability.Coding), modelId = null) + + assertEquals("managed:coder", provider.id.value) + assertEquals(listOf("coder"), manager.loads) + } + + @Test + fun `capability match prefers the highest summed score`(): Unit = runBlocking { + val manager = FakeModelManager() + val router = ManagedInferenceRouter( + manager = manager, + descriptors = listOf( + desc("weak", ModelCapability.Coding to 0.3), + desc("strong", ModelCapability.Coding to 0.9), + ), + defaultModelId = "weak", + ) + + val provider = router.route(stage, setOf(ModelCapability.Coding), modelId = null) + + assertEquals("managed:strong", provider.id.value) + } + + @Test + fun `falls back to default when no modelId and no required capabilities`(): Unit = runBlocking { + val manager = FakeModelManager() + val router = ManagedInferenceRouter( + manager = manager, + descriptors = listOf(desc("default-model"), desc("coder", ModelCapability.Coding to 1.0)), + defaultModelId = "default-model", + ) + + val provider = router.route(stage, emptySet(), modelId = null) + + assertEquals("managed:default-model", provider.id.value) + assertEquals(listOf("default-model"), manager.loads) + } + + @Test + fun `falls back to default when no model satisfies the required capabilities`(): Unit = runBlocking { + val manager = FakeModelManager() + val router = ManagedInferenceRouter( + manager = manager, + descriptors = listOf(desc("default-model")), + defaultModelId = "default-model", + ) + + val provider = router.route(stage, setOf(ModelCapability.Reasoning), modelId = null) + + assertEquals("managed:default-model", provider.id.value) + } + + @Test + fun `unknown explicit modelId throws NoEligibleProviderException`() { + val manager = FakeModelManager() + val router = ManagedInferenceRouter( + manager = manager, + descriptors = listOf(desc("default-model")), + defaultModelId = "default-model", + ) + + assertThrows { + runBlocking { router.route(stage, emptySet(), modelId = "ghost") } + } + } +}