fix(inference): tolerate a briefly-absent provider instead of killing the session (#299)
A transient provider connection drop (e.g. OOM-killed llama.cpp mid-request) collapsed into a hard NoEligibleProviderException that escaped runInference's try/catch (the route() call sat outside it), propagated past the stage retry loop, and landed in ServerModule's session-level catch — failing the WHOLE session even though the failure was retryable and the provider recovered seconds later. - SessionOrchestrator.runInference: wrap router.route() in try/catch so routing failures become InferenceResult.Failed and flow through the normal retryable/backoff/exhaustion machinery instead of escaping. - DefaultInferenceRouter.route: distinguish "capability never configured on any provider" (fail fast, waiting can't help) from "configured but currently unhealthy" (bounded wait/backoff — default 3 attempts x 2s — re-checking health before declaring NoEligibleProvider terminal). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01GeyGFXczJb8RUWGBKmkm6G
This commit is contained in:
@@ -2,6 +2,7 @@ package com.correx.core.inference
|
||||
|
||||
import com.correx.core.events.types.ProviderId
|
||||
import com.correx.core.events.types.StageId
|
||||
import kotlinx.coroutines.delay
|
||||
import kotlinx.coroutines.sync.Mutex
|
||||
import kotlinx.coroutines.sync.withLock
|
||||
import kotlin.time.Duration
|
||||
@@ -21,6 +22,12 @@ class DefaultInferenceRouter(
|
||||
private val strategy: RoutingStrategy,
|
||||
private val cacheTtl: Duration = 5.seconds,
|
||||
private val timeSource: TimeSource = TimeSource.Monotonic,
|
||||
// A provider that briefly drops (crash + qa-stack restart) shouldn't collapse into a hard
|
||||
// NoEligibleProvider abort — give it bounded time to come back before giving up. This only
|
||||
// applies when the capability IS configured on some provider but that provider is currently
|
||||
// unhealthy; a capability nobody ever declared fails immediately (see routeCapabilityCandidates).
|
||||
private val unavailableRetryAttempts: Int = 3,
|
||||
private val unavailableRetryDelay: Duration = 2.seconds,
|
||||
) : InferenceRouter {
|
||||
|
||||
private val cache = mutableMapOf<ProviderId, HealthEntry>()
|
||||
@@ -45,12 +52,39 @@ class DefaultInferenceRouter(
|
||||
}
|
||||
}
|
||||
|
||||
private suspend fun refreshedHealth(provider: InferenceProvider): ProviderHealth =
|
||||
lockFor(provider.id).withLock {
|
||||
val fresh = provider.healthCheck()
|
||||
mapMutex.withLock { cache[provider.id] = HealthEntry(fresh, timeSource.markNow()) }
|
||||
fresh
|
||||
}
|
||||
|
||||
override suspend fun route(stageId: StageId, requiredCapabilities: Set<ModelCapability>): InferenceProvider {
|
||||
val candidates = requiredCapabilities
|
||||
.flatMap { registry.resolve(it) }
|
||||
.distinctBy { it.id }
|
||||
.ifEmpty { registry.listAll() }
|
||||
val healthy = candidates.filter { cachedHealth(it) !is ProviderHealth.Unavailable }
|
||||
|
||||
// Nobody was ever configured with this capability set — no amount of waiting fixes that,
|
||||
// fail fast instead of burning the bounded-wait budget below.
|
||||
if (requiredCapabilities.isNotEmpty() &&
|
||||
candidates.none { it.capabilities().map { c -> c.capability }.toSet().containsAll(requiredCapabilities) }
|
||||
) {
|
||||
throw NoEligibleProviderException(stageId, requiredCapabilities)
|
||||
}
|
||||
|
||||
var healthy = candidates.filter { cachedHealth(it) !is ProviderHealth.Unavailable }
|
||||
var attempt = 0
|
||||
while (healthy.isEmpty() && attempt < unavailableRetryAttempts) {
|
||||
attempt++
|
||||
log.warn(
|
||||
"route: capability {} configured but all candidates unhealthy for stage={}" +
|
||||
" — waiting {} (attempt {}/{}) before declaring NoEligibleProvider",
|
||||
requiredCapabilities, stageId.value, unavailableRetryDelay, attempt, unavailableRetryAttempts,
|
||||
)
|
||||
delay(unavailableRetryDelay)
|
||||
healthy = candidates.filter { refreshedHealth(it) !is ProviderHealth.Unavailable }
|
||||
}
|
||||
val selected = strategy.select(healthy, requiredCapabilities)
|
||||
// Post-selection re-check closes the TOCTOU window between initial filter and dispatch.
|
||||
when (val postHealth = selected.healthCheck()) {
|
||||
|
||||
+9
-1
@@ -300,7 +300,15 @@ abstract class SessionOrchestrator(
|
||||
// ToolCalling score and routes the stage to the best tool-caller.
|
||||
val requiredCapabilities = stageConfig.requiredCapabilities +
|
||||
if (withTools && stageConfig.allowedTools.isNotEmpty()) setOf(ModelCapability.ToolCalling) else emptySet()
|
||||
val provider = inferenceRouter.route(stageId, requiredCapabilities, stageConfig.modelId)
|
||||
// Routing itself can fail transiently (provider mid-crash-recovery) — it must be retryable
|
||||
// like any other inference failure, not escape and kill the whole session (see #299).
|
||||
val provider = try {
|
||||
inferenceRouter.route(stageId, requiredCapabilities, stageConfig.modelId)
|
||||
} catch (e: CancellationException) {
|
||||
throw e
|
||||
} catch (e: Exception) {
|
||||
return InferenceResult.Failed(e.message ?: "routing failed")
|
||||
}
|
||||
log.debug(
|
||||
"[Orchestrator] inference session={} stage={} provider={} timeoutMs={}",
|
||||
sessionId.value, stageId.value, provider.id.value, timeoutMs,
|
||||
|
||||
+56
@@ -13,8 +13,10 @@ import com.correx.core.inference.RoutingStrategy
|
||||
import com.correx.testing.fixtures.inference.MockInferenceProvider
|
||||
import kotlinx.coroutines.runBlocking
|
||||
import org.junit.jupiter.api.Assertions.assertSame
|
||||
import org.junit.jupiter.api.Assertions.assertTrue
|
||||
import org.junit.jupiter.api.Test
|
||||
import org.junit.jupiter.api.assertThrows
|
||||
import kotlin.time.Duration.Companion.milliseconds
|
||||
|
||||
class DefaultInferenceRouterTest {
|
||||
|
||||
@@ -157,4 +159,58 @@ class DefaultInferenceRouterTest {
|
||||
val result = router.route(stage, setOf(ModelCapability.General), "llama-cpp:sick")
|
||||
assertSame(healthy, result)
|
||||
}
|
||||
|
||||
// ── bounded wait for a briefly-absent provider (#299) ─────────────────────
|
||||
|
||||
@Test
|
||||
fun `waits for the sole capable provider to recover instead of failing immediately`(): Unit = runBlocking {
|
||||
var healthChecks = 0
|
||||
val recovering = object : InferenceProvider by provider("a", ModelCapability.ToolCalling) {
|
||||
override suspend fun healthCheck(): ProviderHealth {
|
||||
healthChecks++
|
||||
return if (healthChecks < 3) ProviderHealth.Unavailable("connection dropped") else ProviderHealth.Healthy
|
||||
}
|
||||
}
|
||||
val router = DefaultInferenceRouter(
|
||||
registryOf(recovering),
|
||||
firstStrategy(),
|
||||
unavailableRetryAttempts = 5,
|
||||
unavailableRetryDelay = 5.milliseconds,
|
||||
)
|
||||
val result = router.route(stage, setOf(ModelCapability.ToolCalling))
|
||||
assertSame(recovering, result)
|
||||
assertTrue(healthChecks >= 3) { "expected at least 3 health checks, got $healthChecks" }
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `still throws NoEligibleProviderException if the sole provider never recovers within the bound`() {
|
||||
val neverRecovers = MockInferenceProvider(
|
||||
id = ProviderId("a"),
|
||||
declaredCapabilities = setOf(CapabilityScore(ModelCapability.ToolCalling, 1.0)),
|
||||
health = ProviderHealth.Unavailable("still down"),
|
||||
)
|
||||
val router = DefaultInferenceRouter(
|
||||
registryOf(neverRecovers),
|
||||
throwingStrategy(),
|
||||
unavailableRetryAttempts = 2,
|
||||
unavailableRetryDelay = 5.milliseconds,
|
||||
)
|
||||
assertThrows<NoEligibleProviderException> {
|
||||
runBlocking { router.route(stage, setOf(ModelCapability.ToolCalling)) }
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `fails fast without waiting when the capability was never configured on any provider`(): Unit = runBlocking {
|
||||
val p = provider("a", ModelCapability.General) // does not declare ToolCalling
|
||||
val router = DefaultInferenceRouter(
|
||||
registryOf(p),
|
||||
throwingStrategy(),
|
||||
unavailableRetryAttempts = 5,
|
||||
unavailableRetryDelay = 10_000.milliseconds, // would time the test out if the wait loop ran
|
||||
)
|
||||
assertThrows<NoEligibleProviderException> {
|
||||
router.route(stage, setOf(ModelCapability.ToolCalling))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -46,6 +46,7 @@ import com.correx.core.inference.InferenceResponse
|
||||
import com.correx.core.inference.InferenceRouter
|
||||
import com.correx.core.inference.InferenceState
|
||||
import com.correx.core.inference.ModelCapability
|
||||
import com.correx.core.inference.NoEligibleProviderException
|
||||
import com.correx.core.inference.ProviderHealth
|
||||
import com.correx.core.inference.Token
|
||||
import com.correx.core.inference.TokenUsage
|
||||
@@ -216,6 +217,42 @@ class SessionOrchestratorIntegrationTest {
|
||||
assertTrue(failed.retryExhausted)
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `transient NoEligibleProviderException on a retry is tolerated, not a session-killing crash (#299)`(): Unit =
|
||||
runBlocking {
|
||||
val sessionId = SessionId("s2b")
|
||||
val config = OrchestrationConfig(
|
||||
retryPolicy = RetryPolicy(maxAttempts = 3, backoffMs = 0),
|
||||
)
|
||||
val graph = threeStageGraph()
|
||||
var routeCalls = 0
|
||||
|
||||
val recoveringOrchestrator = DefaultSessionOrchestrator(
|
||||
repositories = repositories,
|
||||
engines = engines.copy(
|
||||
inferenceRouter = object : InferenceRouter {
|
||||
override suspend fun route(stageId: StageId, requiredCapabilities: Set<ModelCapability>): InferenceProvider {
|
||||
routeCalls++
|
||||
// First attempt hits the provider mid-crash — subsequent retries find it back.
|
||||
if (routeCalls == 1) throw NoEligibleProviderException(stageId, requiredCapabilities)
|
||||
return MockInferenceProvider()
|
||||
}
|
||||
},
|
||||
),
|
||||
retryCoordinator = retryCoordinator,
|
||||
artifactStore = artifactStore,
|
||||
decisionJournalRepository = decisionJournalRepository,
|
||||
)
|
||||
recoveringOrchestrator.run(sessionId, graph, config)
|
||||
|
||||
val events = eventStore.read(sessionId)
|
||||
// A crash would have skipped straight to an unhandled failure with no WorkflowCompletedEvent
|
||||
// and no orderly retry — assert the session survived and made it through via the retry path.
|
||||
assertTrue(routeCalls > 1) { "expected routing to be retried, only called $routeCalls time(s)" }
|
||||
assertNotNull(events.find { it.payload is WorkflowCompletedEvent })
|
||||
assertTrue(events.none { it.payload is WorkflowFailedEvent })
|
||||
}
|
||||
|
||||
@Test
|
||||
fun `stage that matches no transition retries through the budget before failing`(): Unit = runBlocking {
|
||||
val sessionId = SessionId("s-no-transition-retry")
|
||||
|
||||
Reference in New Issue
Block a user