mirror of
https://github.com/diegosouzapw/OmniRoute.git
synced 2026-08-05 14:52:09 +03:00
fix(combo): parallel pre-screen + circuit-breaker fast-exit for priority combos (#3169)
Integrated into release/v3.8.18
This commit is contained in:
@@ -13,6 +13,7 @@ import {
|
||||
recordProviderFailure,
|
||||
isProviderFailureCode,
|
||||
isProviderExhaustedReason,
|
||||
type ProviderProfile,
|
||||
} from "./accountFallback.ts";
|
||||
import { FETCH_TIMEOUT_MS, RateLimitReason } from "../config/constants.ts";
|
||||
import { errorResponse, unavailableResponse } from "../utils/error.ts";
|
||||
@@ -27,6 +28,7 @@ import {
|
||||
resolveComboConfig,
|
||||
getDefaultComboConfig,
|
||||
resolveComboTargetTimeoutMs,
|
||||
PRE_SCREEN_CONCURRENCY,
|
||||
} from "./comboConfig.ts";
|
||||
import {
|
||||
maybeGenerateHandoff,
|
||||
@@ -1833,6 +1835,42 @@ async function fetchResetAwareQuotaWithCache({
|
||||
return refresh();
|
||||
}
|
||||
|
||||
type PreScreenResult = { profile: ProviderProfile | null; available: boolean };
|
||||
|
||||
export async function preScreenTargets(
|
||||
targets: ResolvedComboTarget[],
|
||||
isModelAvailable?: ((model: string, target: ResolvedComboTarget) => Promise<boolean>) | null
|
||||
): Promise<Map<string, PreScreenResult>> {
|
||||
if (targets.length === 0) {
|
||||
return new Map();
|
||||
}
|
||||
|
||||
const results = await mapWithConcurrency(
|
||||
targets,
|
||||
PRE_SCREEN_CONCURRENCY,
|
||||
async (target): Promise<{ key: string; result: PreScreenResult }> => {
|
||||
const profile = await getRuntimeProviderProfile(target.provider).catch(() => null);
|
||||
|
||||
const breaker = getCircuitBreaker(target.provider);
|
||||
if (breaker.getStatus().state === "OPEN") {
|
||||
return { key: target.executionKey, result: { profile, available: false } };
|
||||
}
|
||||
|
||||
let available = true;
|
||||
if (isModelAvailable) {
|
||||
available = await isModelAvailable(target.modelStr, target).catch(() => true);
|
||||
}
|
||||
return { key: target.executionKey, result: { profile, available } };
|
||||
}
|
||||
);
|
||||
|
||||
const map = new Map<string, PreScreenResult>();
|
||||
for (const { key, result } of results) {
|
||||
map.set(key, result);
|
||||
}
|
||||
return map;
|
||||
}
|
||||
|
||||
async function orderTargetsByResetAwareQuota(
|
||||
targets: ResolvedComboTarget[],
|
||||
comboName: string,
|
||||
@@ -3217,6 +3255,15 @@ export async function handleComboChat({
|
||||
orderedTargets = orderTargetsByEvalScores(orderedTargets, config.evalRouting, log);
|
||||
orderedTargets = filterTargetsByRequestCompatibility(orderedTargets, body, log);
|
||||
|
||||
// Parallel pre-screen: check provider profiles and model availability for all targets
|
||||
// Only runs for priority strategy where sequential checking causes latency
|
||||
const preScreenMap =
|
||||
strategy === "priority"
|
||||
? await preScreenTargets(orderedTargets, isModelAvailable).catch(
|
||||
() => new Map<string, PreScreenResult>()
|
||||
)
|
||||
: new Map<string, PreScreenResult>();
|
||||
|
||||
if (orderedTargets.length === 0) {
|
||||
return comboModelNotFoundResponse("Combo has no executable targets");
|
||||
}
|
||||
@@ -3285,7 +3332,18 @@ export async function handleComboChat({
|
||||
const target = orderedTargets[i];
|
||||
const modelStr = target.modelStr;
|
||||
const provider = target.provider;
|
||||
const profile = await getRuntimeProviderProfile(provider);
|
||||
|
||||
const cb = getCircuitBreaker(provider);
|
||||
if (cb.getStatus().state === "OPEN") {
|
||||
log.info("COMBO", `Skipping ${modelStr} — circuit breaker OPEN for ${provider}`);
|
||||
if (i > 0) fallbackCount++;
|
||||
return null;
|
||||
}
|
||||
|
||||
// Use pre-screened profile if available, otherwise fetch on demand
|
||||
const preScreenEntry = preScreenMap.get(target.executionKey);
|
||||
const profile = preScreenEntry?.profile ?? (await getRuntimeProviderProfile(provider));
|
||||
|
||||
const allowRateLimitedConnection =
|
||||
Boolean(provider && provider !== "unknown") &&
|
||||
transientRateLimitedProviders.has(provider);
|
||||
@@ -3307,7 +3365,18 @@ export async function handleComboChat({
|
||||
return null;
|
||||
}
|
||||
|
||||
// Pre-check: skip models where no credentials are available (excluded, rate-limited, or unavailable)
|
||||
// Pre-screen may have already determined this target unavailable (e.g.
|
||||
// circuit-breaker OPEN at resolve time). Skip immediately in that case.
|
||||
// For targets pre-screened as "available" we still call isModelAvailable
|
||||
// below because connection cooldowns (rateLimitedUntil) can change
|
||||
// mid-request after a same-provider failure — the pre-screen snapshot is
|
||||
// stale by the time we reach the 2nd/3rd same-provider target.
|
||||
const preCheckedAvailable = preScreenEntry?.available ?? null;
|
||||
if (preCheckedAvailable === false) {
|
||||
log.info("COMBO", `Skipping ${modelStr} — pre-screen marked unavailable`);
|
||||
if (i > 0) fallbackCount++;
|
||||
return null;
|
||||
}
|
||||
if (isModelAvailable) {
|
||||
const available = await isModelAvailable(modelStr, targetForAttempt);
|
||||
if (!available) {
|
||||
|
||||
@@ -7,6 +7,12 @@
|
||||
|
||||
import { MAX_TIMER_TIMEOUT_MS } from "../../src/shared/utils/runtimeTimeouts.ts";
|
||||
|
||||
/**
|
||||
* Maximum number of concurrent pre-screen checks (provider profile + availability)
|
||||
* when running parallel pre-screening for priority strategy combos.
|
||||
*/
|
||||
export const PRE_SCREEN_CONCURRENCY = 5;
|
||||
|
||||
const DEFAULT_COMBO_CONFIG = {
|
||||
strategy: "priority",
|
||||
maxRetries: 1,
|
||||
|
||||
205
tests/unit/combo-prescreen.test.ts
Normal file
205
tests/unit/combo-prescreen.test.ts
Normal file
@@ -0,0 +1,205 @@
|
||||
import test, { after } from "node:test";
|
||||
import assert from "node:assert/strict";
|
||||
import fs from "node:fs";
|
||||
import os from "node:os";
|
||||
import path from "node:path";
|
||||
|
||||
const TEST_DATA_DIR = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-combo-prescreen-"));
|
||||
const ORIGINAL_DATA_DIR = process.env.DATA_DIR;
|
||||
process.env.DATA_DIR = TEST_DATA_DIR;
|
||||
|
||||
const dbCore = await import("../../src/lib/db/core.ts");
|
||||
const { handleComboChat } = await import("../../open-sse/services/combo.ts");
|
||||
const combosDb = await import("../../src/lib/db/combos.ts");
|
||||
|
||||
after(() => {
|
||||
dbCore.resetDbInstance();
|
||||
fs.rmSync(TEST_DATA_DIR, { recursive: true, force: true });
|
||||
if (ORIGINAL_DATA_DIR === undefined) {
|
||||
delete process.env.DATA_DIR;
|
||||
} else {
|
||||
process.env.DATA_DIR = ORIGINAL_DATA_DIR;
|
||||
}
|
||||
});
|
||||
|
||||
function okResponse(model: string) {
|
||||
return Response.json({ choices: [{ message: { role: "assistant", content: model } }] });
|
||||
}
|
||||
|
||||
function makeLog() {
|
||||
return {
|
||||
info() {},
|
||||
warn() {},
|
||||
debug() {},
|
||||
error() {},
|
||||
};
|
||||
}
|
||||
|
||||
const reqBody = {
|
||||
model: "prescreen-test",
|
||||
messages: [{ role: "user", content: "hi" }],
|
||||
stream: false,
|
||||
};
|
||||
|
||||
test("pre-screen: all targets checked in parallel", async () => {
|
||||
const checkOrder: string[] = [];
|
||||
|
||||
const combo = await combosDb.createCombo({
|
||||
name: "prescreen-parallel",
|
||||
strategy: "priority",
|
||||
models: ["p1/m1", "p2/m2", "p3/m3"],
|
||||
});
|
||||
|
||||
const response = await handleComboChat({
|
||||
body: { ...reqBody, model: combo.name },
|
||||
combo,
|
||||
allCombos: [combo],
|
||||
isModelAvailable: async (modelStr: string) => {
|
||||
checkOrder.push(modelStr);
|
||||
return true;
|
||||
},
|
||||
relayOptions: undefined,
|
||||
signal: undefined,
|
||||
settings: {},
|
||||
log: makeLog(),
|
||||
handleSingleModel: async (_body: unknown, modelStr: string) => {
|
||||
return okResponse(modelStr);
|
||||
},
|
||||
});
|
||||
|
||||
assert.equal(response.status, 200);
|
||||
assert.ok(checkOrder.length >= 3, "pre-screen must check all 3 targets concurrently");
|
||||
});
|
||||
|
||||
test("pre-screen: unavailable targets are skipped", async () => {
|
||||
const calls: string[] = [];
|
||||
const availability = new Map([
|
||||
["p1/m1", false],
|
||||
["p2/m2", true],
|
||||
["p3/m3", true],
|
||||
]);
|
||||
|
||||
const combo = await combosDb.createCombo({
|
||||
name: "prescreen-skip",
|
||||
strategy: "priority",
|
||||
models: ["p1/m1", "p2/m2", "p3/m3"],
|
||||
});
|
||||
|
||||
const response = await handleComboChat({
|
||||
body: { ...reqBody, model: combo.name },
|
||||
combo,
|
||||
allCombos: [combo],
|
||||
isModelAvailable: async (modelStr: string) => {
|
||||
return availability.get(modelStr) ?? true;
|
||||
},
|
||||
relayOptions: undefined,
|
||||
signal: undefined,
|
||||
settings: {},
|
||||
log: makeLog(),
|
||||
handleSingleModel: async (_body: unknown, modelStr: string) => {
|
||||
calls.push(modelStr);
|
||||
return okResponse(modelStr);
|
||||
},
|
||||
});
|
||||
|
||||
assert.equal(response.status, 200);
|
||||
assert.equal(calls.length, 1);
|
||||
assert.equal(calls[0], "p2/m2");
|
||||
});
|
||||
|
||||
test("pre-screen: failure treated as unknown availability", async () => {
|
||||
const calls: string[] = [];
|
||||
let checkCount = 0;
|
||||
|
||||
const combo = await combosDb.createCombo({
|
||||
name: "prescreen-failure",
|
||||
strategy: "priority",
|
||||
models: ["p1/m1", "p2/m2"],
|
||||
});
|
||||
|
||||
const response = await handleComboChat({
|
||||
body: { ...reqBody, model: combo.name },
|
||||
combo,
|
||||
allCombos: [combo],
|
||||
isModelAvailable: async (modelStr: string) => {
|
||||
checkCount++;
|
||||
if (checkCount === 1) {
|
||||
throw new Error("DB connection failed");
|
||||
}
|
||||
return true;
|
||||
},
|
||||
relayOptions: undefined,
|
||||
signal: undefined,
|
||||
settings: {},
|
||||
log: makeLog(),
|
||||
handleSingleModel: async (_body: unknown, modelStr: string) => {
|
||||
calls.push(modelStr);
|
||||
return okResponse(modelStr);
|
||||
},
|
||||
});
|
||||
|
||||
assert.equal(response.status, 200);
|
||||
// m1 pre-screen failed (treated as available), m2 pre-screen succeeded
|
||||
// Both should be tried, m1 first since it's available
|
||||
assert.equal(calls.length, 1);
|
||||
assert.equal(calls[0], "p1/m1");
|
||||
});
|
||||
|
||||
test("pre-screen: only runs for priority strategy", async () => {
|
||||
const checkOrder: string[] = [];
|
||||
|
||||
const combo = await combosDb.createCombo({
|
||||
name: "prescreen-strategy",
|
||||
strategy: "round-robin",
|
||||
models: ["p1/m1", "p2/m2"],
|
||||
});
|
||||
|
||||
const response = await handleComboChat({
|
||||
body: { ...reqBody, model: combo.name },
|
||||
combo,
|
||||
allCombos: [combo],
|
||||
isModelAvailable: async (modelStr: string) => {
|
||||
checkOrder.push(modelStr);
|
||||
return true;
|
||||
},
|
||||
relayOptions: undefined,
|
||||
signal: undefined,
|
||||
settings: {},
|
||||
log: makeLog(),
|
||||
handleSingleModel: async (_body: unknown, modelStr: string) => {
|
||||
return okResponse(modelStr);
|
||||
},
|
||||
});
|
||||
|
||||
assert.equal(response.status, 200);
|
||||
assert.equal(checkOrder.length >= 1, true);
|
||||
});
|
||||
|
||||
test("pre-screen: backward compatible with all targets available", async () => {
|
||||
const calls: string[] = [];
|
||||
|
||||
const combo = await combosDb.createCombo({
|
||||
name: "prescreen-compat",
|
||||
strategy: "priority",
|
||||
models: ["p1/m1", "p2/m2", "p3/m3"],
|
||||
});
|
||||
|
||||
const response = await handleComboChat({
|
||||
body: { ...reqBody, model: combo.name },
|
||||
combo,
|
||||
allCombos: [combo],
|
||||
isModelAvailable: async () => true,
|
||||
relayOptions: undefined,
|
||||
signal: undefined,
|
||||
settings: {},
|
||||
log: makeLog(),
|
||||
handleSingleModel: async (_body: unknown, modelStr: string) => {
|
||||
calls.push(modelStr);
|
||||
return okResponse(modelStr);
|
||||
},
|
||||
});
|
||||
|
||||
assert.equal(response.status, 200);
|
||||
assert.equal(calls.length, 1);
|
||||
assert.equal(calls[0], "p1/m1");
|
||||
});
|
||||
@@ -14,6 +14,12 @@ const {
|
||||
seedConnection,
|
||||
settingsDb,
|
||||
} = harness;
|
||||
const { preScreenTargets } = await import(
|
||||
"../../open-sse/services/combo.ts"
|
||||
);
|
||||
const { getCircuitBreaker } = await import(
|
||||
"../../src/shared/utils/circuitBreaker.ts"
|
||||
);
|
||||
|
||||
test.beforeEach(async () => {
|
||||
await resetStorage();
|
||||
@@ -105,3 +111,49 @@ test("combo failover skips the cooled provider target on the next request", asyn
|
||||
assert.equal(openaiCalls, 1);
|
||||
assert.equal(claudeCalls, 2);
|
||||
});
|
||||
|
||||
test("pre-screen marks target unavailable when circuit breaker is OPEN", async () => {
|
||||
const breaker = getCircuitBreaker("openai", { failureThreshold: 1, resetTimeout: 60_000 });
|
||||
try {
|
||||
await breaker.execute(async () => {
|
||||
throw new Error("simulated failure");
|
||||
});
|
||||
} catch {
|
||||
// expected
|
||||
}
|
||||
|
||||
const targets = [
|
||||
{
|
||||
kind: "model" as const,
|
||||
stepId: "step-1",
|
||||
executionKey: "openai/gpt-4o",
|
||||
modelStr: "openai/gpt-4o",
|
||||
provider: "openai",
|
||||
providerId: "conn-1",
|
||||
connectionId: "conn-1",
|
||||
weight: 1,
|
||||
label: null,
|
||||
},
|
||||
{
|
||||
kind: "model" as const,
|
||||
stepId: "step-2",
|
||||
executionKey: "claude/claude-3-5-sonnet-20241022",
|
||||
modelStr: "claude/claude-3-5-sonnet-20241022",
|
||||
provider: "claude",
|
||||
providerId: "conn-2",
|
||||
connectionId: "conn-2",
|
||||
weight: 1,
|
||||
label: null,
|
||||
},
|
||||
];
|
||||
|
||||
const results = await preScreenTargets(targets as any);
|
||||
|
||||
const openaiResult = results.get("openai/gpt-4o");
|
||||
assert.ok(openaiResult, "openai target should have a pre-screen result");
|
||||
assert.equal(openaiResult.available, false, "open-circuit-breaker target should be unavailable");
|
||||
|
||||
const claudeResult = results.get("claude/claude-3-5-sonnet-20241022");
|
||||
assert.ok(claudeResult, "claude target should have a pre-screen result");
|
||||
assert.equal(claudeResult.available, true, "closed-circuit-breaker target should be available");
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user