diff --git a/packages/console/app/src/routes/zen/util/handler.ts b/packages/console/app/src/routes/zen/util/handler.ts index 44dc5f0cae..320fccfe1c 100644 --- a/packages/console/app/src/routes/zen/util/handler.ts +++ b/packages/console/app/src/routes/zen/util/handler.ts @@ -284,8 +284,7 @@ export async function handler( const costInfo = calculateCost(modelInfo, usageInfo) await trialLimiter?.track(usageInfo) await modelTpmLimiter?.track(providerInfo.id, providerInfo.model, usageInfo) - if (providerInfo.budgetPriority !== undefined) - await providerBudgetTracker?.track(providerInfo.id, providerInfo.budgetPriority, costInfo.totalCostInCent) + await providerBudgetTracker?.track(providerInfo.id, providerInfo.budgetPriority, costInfo.totalCostInCent) await trackUsage(sessionId, billingSource, authInfo, modelInfo, providerInfo, usageInfo, costInfo) await reload(billingSource, authInfo, costInfo) json.cost = calculateOccurredCost(billingSource, costInfo) @@ -346,12 +345,7 @@ export async function handler( timestampLastByte, usageInfo, ) - if (providerInfo.budgetPriority !== undefined) - await providerBudgetTracker?.track( - providerInfo.id, - providerInfo.budgetPriority, - costInfo.totalCostInCent, - ) + await providerBudgetTracker?.track(providerInfo.id, providerInfo.budgetPriority, costInfo.totalCostInCent) await trackUsage(sessionId, billingSource, authInfo, modelInfo, providerInfo, usageInfo, costInfo) await reload(billingSource, authInfo, costInfo) const cost = calculateOccurredCost(billingSource, costInfo) @@ -518,7 +512,12 @@ export async function handler( stickyProviderId: string | undefined, modelTpmLimits: Record | undefined, modelTpsLimits: Record | undefined, - providerBudget: { qualify: (providerId: string, priority: number) => boolean } | undefined, + providerBudget: + | { + qualify: (providerId: string, priority: number) => boolean + prefer: (providerId: string, priority: number) => boolean + } + | undefined, ) { const modelProvider = (() => { // Byok is top priority b/c if user set their own API key, we should use it @@ -581,15 +580,19 @@ export async function handler( const stickProvider = allProviders.find((provider) => provider.id === stickyProviderId) if (!stickProvider) return provider - // stick provider exists + selected provider is API type => use sticky provider - if (!provider.tpsGoal) return stickProvider + const preferBudgetProvider = + provider.budgetPriority !== undefined && providerBudget?.prefer(provider.id, provider.budgetPriority) - // stick provier exists + selected provider is GPU type + GPU not idle => use selected provider - const tps = modelTpsLimits?.[`${provider.id}/${provider.model}/${provider.tpsGoal}`] ?? { - qualify: 0, - unqualify: 0, - } - if (tps.qualify <= tps.unqualify * 3) return stickProvider + const preferTpsProvider = (() => { + if (!provider.tpsGoal) return false + const tps = modelTpsLimits?.[`${provider.id}/${provider.model}/${provider.tpsGoal}`] ?? { + qualify: 0, + unqualify: 0, + } + return tps.qualify > tps.unqualify * 3 + })() + + if (!preferBudgetProvider && !preferTpsProvider) return stickProvider return provider } diff --git a/packages/console/app/src/routes/zen/util/providerBudgetTracker.ts b/packages/console/app/src/routes/zen/util/providerBudgetTracker.ts index 14dc7f31f2..a21f70b538 100644 --- a/packages/console/app/src/routes/zen/util/providerBudgetTracker.ts +++ b/packages/console/app/src/routes/zen/util/providerBudgetTracker.ts @@ -55,6 +55,7 @@ export function createProviderBudgetTracker( let effectiveBudget: Record> = {} // Cumulative current-minute spend through each priority, per provider. let spentThroughPriority: Record> = {} + let previousSpentThroughPriority: Record> = {} return { // Returns whether a provider at a given priority still has budget headroom. @@ -87,6 +88,7 @@ export function createProviderBudgetTracker( effectiveBudget = {} spentThroughPriority = {} + previousSpentThroughPriority = {} Object.entries(maxPriorityByProvider).forEach(([providerId, maxPriority]) => { const providerBudget = budgetByProvider[providerId] if (providerBudget === undefined) return @@ -96,11 +98,13 @@ export function createProviderBudgetTracker( let previousRunning = 0 effectiveBudget[providerId] = {} spentThroughPriority[providerId] = {} + previousSpentThroughPriority[providerId] = {} Array.from({ length: maxPriority }, (_, index) => index + 1).forEach((priority) => { currentRunning += current[providerId]?.[priority] ?? 0 effectiveBudget[providerId][priority] = Math.max(0, budget - previousRunning) previousRunning += previous[providerId]?.[priority] ?? 0 spentThroughPriority[providerId][priority] = currentRunning + previousSpentThroughPriority[providerId][priority] = previousRunning }) }) @@ -114,9 +118,18 @@ export function createProviderBudgetTracker( const spentThroughCurrentPriority = spentThroughPriority[providerId]?.[priority] ?? 0 return spentThroughCurrentPriority < budget }, + prefer: (providerId: string, priority: number) => { + const providerBudget = budgetByProvider[providerId] + if (providerBudget === undefined) return false + const budget = centsToMicroCents(providerBudget * 100) + const previousUsage = previousSpentThroughPriority[providerId]?.[priority] + if (previousUsage === undefined) return false + return previousUsage < budget * 0.8 + }, } }, - track: async (provider: string, priority: number, costInCent: number) => { + track: async (provider: string, priority: number | undefined, costInCent: number) => { + if (priority === undefined) return const config = tracked.find((item) => item.id === provider && item.budgetPriority === priority) if (!config) return if (config.budgetContribution === undefined) return