BCAI/easyai-ai-gateway
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: add runner failover rules and cache affinity

Browse Source
This commit is contained in:
wangbo
2026-06-28 20:50:23 +08:00
parent 41bb3a6525
commit 0e675e259c
31 changed files with 2939 additions and 470 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/internal/store/candidates.go
+253 -20
View File
@@ -3,6 +3,7 @@ package store
import (
"context"
"fmt"
"math"
"sort"
"strings"
"unicode"
@@ -10,11 +11,17 @@ import (
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
)
func (s *Store) ListModelCandidates(ctx context.Context, model string, modelType string, user *auth.User) ([]RuntimeModelCandidate, error) {
type ListModelCandidatesOptions struct {
CacheAffinityKey string
CacheAffinityPolicy map[string]any
}
func (s *Store) ListModelCandidates(ctx context.Context, model string, modelType string, user *auth.User, options ...ListModelCandidatesOptions) ([]RuntimeModelCandidate, error) {
exactModel := strings.TrimSpace(model)
modelMatchKey := normalizeModelMatchKey(exactModel)
listOptions := normalizeListModelCandidatesOptions(modelType, options...)
rows, err := s.pool.Query(ctx, `
SELECT p.id::text, p.platform_key, p.name, p.provider,
SELECT p.id::text, p.platform_key, p.name, p.provider,
COALESCE(NULLIF(p.config->>'specType', ''), NULLIF(cp.provider_type, ''), NULLIF(p.config->>'sourceSpecType', ''), p.provider) AS spec_type,
COALESCE(p.base_url, ''),
p.auth_type, p.credentials, p.config, p.default_pricing_mode,
@@ -35,17 +42,27 @@ SELECT p.id::text, p.platform_key, p.name, p.provider,
COALESCE(queued.waiting, 0)::float8,
COALESCE(rpm.used_value, 0)::float8, COALESCE(rpm.reserved_value, 0)::float8,
COALESCE(tpm.used_value, 0)::float8, COALESCE(tpm.reserved_value, 0)::float8,
COALESCE(s.running_count, 0)::float8,
COALESCE(s.waiting_count, 0)::float8,
COALESCE(s.limiter_ratio, 0)::float8,
COALESCE(EXTRACT(EPOCH FROM s.last_assigned_at), 0)::float8
FROM platform_models m
JOIN integration_platforms p ON p.id = m.platform_id
LEFT JOIN model_catalog_providers cp ON cp.provider_key = p.provider OR cp.provider_code = p.provider
LEFT JOIN base_model_catalog b ON b.id = m.base_model_id
LEFT JOIN model_runtime_policy_sets rp ON rp.id = COALESCE(m.runtime_policy_set_id, b.runtime_policy_set_id)
LEFT JOIN runtime_client_states s
ON s.client_id = p.platform_key || ':' || $2::text || ':' || COALESCE(NULLIF(m.provider_model_name, ''), m.model_name)
COALESCE(s.running_count, 0)::float8,
COALESCE(s.waiting_count, 0)::float8,
COALESCE(s.limiter_ratio, 0)::float8,
COALESCE(EXTRACT(EPOCH FROM s.last_assigned_at), 0)::float8,
COALESCE(ca.request_count, 0)::float8,
COALESCE(ca.input_tokens, 0)::float8,
COALESCE(ca.cached_input_tokens, 0)::float8,
COALESCE(ca.ema_hit_ratio, 0)::float8,
COALESCE(ca.last_hit_ratio, 0)::float8,
COALESCE(EXTRACT(EPOCH FROM ca.last_observed_at), 0)::float8
FROM platform_models m
JOIN integration_platforms p ON p.id = m.platform_id
LEFT JOIN model_catalog_providers cp ON cp.provider_key = p.provider OR cp.provider_code = p.provider
LEFT JOIN base_model_catalog b ON b.id = m.base_model_id
LEFT JOIN model_runtime_policy_sets rp ON rp.id = COALESCE(m.runtime_policy_set_id, b.runtime_policy_set_id)
LEFT JOIN runtime_client_states s
ON s.client_id = p.platform_key || ':' || $2::text || ':' || COALESCE(NULLIF(m.provider_model_name, ''), m.model_name)
LEFT JOIN gateway_cache_affinity_stats ca
ON ca.client_id = p.platform_key || ':' || $2::text || ':' || COALESCE(NULLIF(m.provider_model_name, ''), m.model_name)
AND ca.cache_affinity_key = NULLIF($4::text, '')
AND ($5::int <= 0 OR ca.last_observed_at >= now() - ($5::int * interval '1 second'))
LEFT JOIN (
SELECT scope_key, SUM(lease_value) AS active
FROM gateway_concurrency_leases
@@ -154,11 +171,11 @@ WHERE p.status = 'enabled'
)
)
)
ORDER BY effective_priority ASC,
COALESCE(s.running_count, 0) ASC,
COALESCE(s.waiting_count, 0) ASC,
COALESCE(s.last_assigned_at, to_timestamp(0)) ASC,
m.created_at ASC`, exactModel, modelType, modelMatchKey)
ORDER BY effective_priority ASC,
COALESCE(s.running_count, 0) ASC,
COALESCE(s.waiting_count, 0) ASC,
COALESCE(s.last_assigned_at, to_timestamp(0)) ASC,
m.created_at ASC`, exactModel, modelType, modelMatchKey, listOptions.CacheAffinityKey, cacheAffinityStaleAfterSeconds(listOptions.CacheAffinityPolicy))
if err != nil {
return nil, err
}
@@ -194,6 +211,12 @@ ORDER BY effective_priority ASC,
var stateWaitingCount float64
var stateLimiterRatio float64
var lastAssignedUnix float64
var cacheRequestCount float64
var cacheInputTokens float64
var cacheCachedInputTokens float64
var cacheEMAHitRatio float64
var cacheLastHitRatio float64
var cacheLastObservedUnix float64
if err := rows.Scan(
&item.PlatformID,
&item.PlatformKey,
@@ -246,6 +269,12 @@ ORDER BY effective_priority ASC,
&stateWaitingCount,
&stateLimiterRatio,
&lastAssignedUnix,
&cacheRequestCount,
&cacheInputTokens,
&cacheCachedInputTokens,
&cacheEMAHitRatio,
&cacheLastHitRatio,
&cacheLastObservedUnix,
); err != nil {
return nil, err
}
@@ -272,6 +301,14 @@ ORDER BY effective_priority ASC,
item.RunningCount = stateRunningCount
item.WaitingCount = maxFloat(queuedWaiting, stateWaitingCount)
item.LastAssignedUnix = lastAssignedUnix
applyRuntimeCandidateCacheAffinity(&item, listOptions, runtimeCandidateCacheAffinityInput{
RequestCount: int(cacheRequestCount),
InputTokens: int(cacheInputTokens),
CachedInputTokens: int(cacheCachedInputTokens),
EMAHitRatio: cacheEMAHitRatio,
LastHitRatio: cacheLastHitRatio,
LastObservedUnix: cacheLastObservedUnix,
})
applyRuntimeCandidateLoad(&item, runtimeCandidateLoadInput{
Policy: effectiveModelRateLimitPolicy(item.PlatformRateLimitPolicy, item.RuntimeRateLimitPolicy, item.RuntimePolicySetID, item.RuntimePolicyOverride, item.ModelRateLimitPolicy),
ConcurrentActive: concurrentActive,
@@ -383,6 +420,182 @@ type runtimeCandidateLoadInput struct {
StateLimiterRatio float64
}
type runtimeCandidateCacheAffinityInput struct {
RequestCount int
InputTokens int
CachedInputTokens int
EMAHitRatio float64
LastHitRatio float64
LastObservedUnix float64
}
func normalizeListModelCandidatesOptions(modelType string, options ...ListModelCandidatesOptions) ListModelCandidatesOptions {
if len(options) == 0 {
return ListModelCandidatesOptions{}
}
out := options[0]
out.CacheAffinityKey = strings.TrimSpace(out.CacheAffinityKey)
if out.CacheAffinityKey == "" || !cacheAffinityPolicyEnabled(out.CacheAffinityPolicy, modelType) {
out.CacheAffinityKey = ""
}
return out
}
func applyRuntimeCandidateCacheAffinity(candidate *RuntimeModelCandidate, options ListModelCandidatesOptions, input runtimeCandidateCacheAffinityInput) {
affinity := RuntimeCandidateCacheAffinity{
Key: options.CacheAffinityKey,
RequestCount: input.RequestCount,
InputTokens: input.InputTokens,
CachedInputTokens: input.CachedInputTokens,
EMAHitRatio: boundedRatio(input.EMAHitRatio),
LastHitRatio: boundedRatio(input.LastHitRatio),
LastObservedUnix: input.LastObservedUnix,
AdjustedPriority: float64(candidate.PlatformPriority),
}
minSamples := cacheAffinityMinSamples(options.CacheAffinityPolicy)
hasObservedCachedHit := affinity.CachedInputTokens > 0 || affinity.LastHitRatio > 0
hasSampledAffinity := input.RequestCount >= minSamples && affinity.EMAHitRatio > 0
if options.CacheAffinityKey == "" || input.RequestCount <= 0 || (!hasObservedCachedHit && !hasSampledAffinity) {
candidate.CacheAffinity = affinity
return
}
affinity.Confidence = float64(input.RequestCount) / float64(minSamples)
if affinity.Confidence > 1 {
affinity.Confidence = 1
}
if hasObservedCachedHit && affinity.Confidence < 1 {
affinity.Confidence = 1
}
affinity.Score = runtimeCandidateCacheAffinityScore(affinity)
maxBoost := cacheAffinityMaxPriorityBoost(options.CacheAffinityPolicy)
affinity.Boost = affinity.Score * maxBoost
if affinity.Boost > maxBoost {
affinity.Boost = maxBoost
}
if hasObservedCachedHit || affinity.Score > 0 {
affinity.Applied = true
}
if affinity.Boost > 0 {
affinity.AdjustedPriority = float64(candidate.PlatformPriority) - affinity.Boost
}
candidate.CacheAffinity = affinity
}
func cacheAffinityPolicyEnabled(policy map[string]any, modelType string) bool {
if enabled, ok := policy["enabled"].(bool); ok && !enabled {
return false
}
modelTypes := stringListFromAny(policy["modelTypes"])
if len(modelTypes) == 0 {
modelTypes = []string{"chat", "text_generate", "responses"}
}
modelType = strings.TrimSpace(modelType)
for _, candidate := range modelTypes {
if strings.EqualFold(strings.TrimSpace(candidate), modelType) {
return true
}
}
return false
}
func cacheAffinityMinSamples(policy map[string]any) int {
value := intFromStorePolicy(policy, "minSamples", 3)
if value <= 0 {
return 1
}
return value
}
func cacheAffinityMinInputTokens(policy map[string]any) int {
value := intFromStorePolicy(policy, "minInputTokens", 512)
if value < 0 {
return 0
}
return value
}
func cacheAffinityStaleAfterSeconds(policy map[string]any) int {
return intFromStorePolicy(policy, "staleAfterSeconds", 86400)
}
func cacheAffinityMaxPriorityBoost(policy map[string]any) float64 {
value := floatFromStorePolicy(policy, "maxPriorityBoost", 20)
if value < 0 {
return 0
}
return value
}
func cacheAffinityEMAAlpha(policy map[string]any) float64 {
value := floatFromStorePolicy(policy, "emaAlpha", 0.35)
if value <= 0 || value > 1 {
return 0.35
}
return value
}
func runtimeCandidateCacheAffinityScore(affinity RuntimeCandidateCacheAffinity) float64 {
hitRatio := affinity.EMAHitRatio
if affinity.LastHitRatio > hitRatio {
hitRatio = affinity.LastHitRatio
}
hitRatio = boundedRatio(hitRatio)
confidence := boundedRatio(affinity.Confidence)
if hitRatio <= 0 || confidence <= 0 {
return 0
}
tokenFactor := 0.0
if affinity.CachedInputTokens > 0 {
tokenFactor = float64(affinity.CachedInputTokens) / 8192
if tokenFactor > 1 {
tokenFactor = 1
}
}
return math.Round(hitRatio*confidence*(1+tokenFactor)*1_000_000) / 1_000_000
}
func intFromStorePolicy(policy map[string]any, key string, fallback int) int {
if policy == nil {
return fallback
}
switch value := policy[key].(type) {
case int:
return value
case int64:
return int(value)
case float64:
return int(value)
default:
return fallback
}
}
func floatFromStorePolicy(policy map[string]any, key string, fallback float64) float64 {
if policy == nil {
return fallback
}
switch value := policy[key].(type) {
case int:
return float64(value)
case int64:
return float64(value)
case float64:
return value
default:
return fallback
}
}
func boundedRatio(value float64) float64 {
if value < 0 {
return 0
}
if value > 1 {
return 1
}
return value
}
func applyRuntimeCandidateLoad(candidate *RuntimeModelCandidate, input runtimeCandidateLoadInput) {
rpmLimit := rateLimitForMetric(input.Policy, "rpm")
tpmLimitValue := tpmLimit(input.Policy)
@@ -422,6 +635,9 @@ func sortRuntimeModelCandidates(items []RuntimeModelCandidate) {
hasNonFull := false
for index := range items {
items[index].LoadAvoided = false
if !items[index].CacheAffinity.Applied && items[index].CacheAffinity.AdjustedPriority == 0 {
items[index].CacheAffinity.AdjustedPriority = float64(items[index].PlatformPriority)
}
if runtimeCandidateFull(items[index]) {
hasFull = true
} else {
@@ -439,8 +655,25 @@ func sortRuntimeModelCandidates(items []RuntimeModelCandidate) {
if aFull != bFull {
return !aFull
}
if items[i].PlatformPriority != items[j].PlatformPriority {
return items[i].PlatformPriority < items[j].PlatformPriority
if items[i].CacheAffinity.Applied != items[j].CacheAffinity.Applied {
return items[i].CacheAffinity.Applied
}
if items[i].CacheAffinity.Applied && items[j].CacheAffinity.Applied {
if items[i].CacheAffinity.Score != items[j].CacheAffinity.Score {
return items[i].CacheAffinity.Score > items[j].CacheAffinity.Score
}
if items[i].CacheAffinity.EMAHitRatio != items[j].CacheAffinity.EMAHitRatio {
return items[i].CacheAffinity.EMAHitRatio > items[j].CacheAffinity.EMAHitRatio
}
if items[i].CacheAffinity.LastHitRatio != items[j].CacheAffinity.LastHitRatio {
return items[i].CacheAffinity.LastHitRatio > items[j].CacheAffinity.LastHitRatio
}
if items[i].CacheAffinity.CachedInputTokens != items[j].CacheAffinity.CachedInputTokens {
return items[i].CacheAffinity.CachedInputTokens > items[j].CacheAffinity.CachedInputTokens
}
}
if items[i].CacheAffinity.AdjustedPriority != items[j].CacheAffinity.AdjustedPriority {
return items[i].CacheAffinity.AdjustedPriority < items[j].CacheAffinity.AdjustedPriority
}
if items[i].LoadRatio != items[j].LoadRatio {
return items[i].LoadRatio < items[j].LoadRatio