zen: budget

This commit is contained in:
Frank 2026-06-29 05:46:53 -04:00
parent 078385d386
commit f90d154465
2 changed files with 34 additions and 18 deletions

View file

@ -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<string, number> | undefined,
modelTpsLimits: Record<string, { qualify: number; unqualify: number }> | 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
}

View file

@ -55,6 +55,7 @@ export function createProviderBudgetTracker(
let effectiveBudget: Record<string, Record<number, number>> = {}
// Cumulative current-minute spend through each priority, per provider.
let spentThroughPriority: Record<string, Record<number, number>> = {}
let previousSpentThroughPriority: Record<string, Record<number, number>> = {}
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