Skip to content

Commit 678c3ae

Browse files
author
QTom
committed
feat: integrate RPM scheduling checks into account selection flow
1 parent c1c31ed commit 678c3ae

1 file changed

Lines changed: 101 additions & 6 deletions

File tree

‎backend/internal/service/gateway_service.go‎

Lines changed: 101 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1157,6 +1157,7 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
11571157
return nil, errors.New("no available accounts")
11581158
}
11591159
ctx = s.withWindowCostPrefetch(ctx, accounts)
1160+
ctx = s.withRPMPrefetch(ctx, accounts)
11601161

11611162
isExcluded := func(accountID int64) bool {
11621163
if excludedIDs == nil {
@@ -1232,6 +1233,10 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
12321233
filteredWindowCost++
12331234
continue
12341235
}
1236+
// RPM 检查（非粘性会话路径）
1237+
if !s.isAccountSchedulableForRPM(ctx, account, false) {
1238+
continue
1239+
}
12351240
routingCandidates = append(routingCandidates, account)
12361241
}
12371242

@@ -1255,7 +1260,9 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
12551260
s.isAccountAllowedForPlatform(stickyAccount, platform, useMixed) &&
12561261
(requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, stickyAccount, requestedModel)) &&
12571262
s.isAccountSchedulableForModelSelection(ctx, stickyAccount, requestedModel) &&
1258-
s.isAccountSchedulableForWindowCost(ctx, stickyAccount, true) { // 粘性会话窗口费用检查
1263+
s.isAccountSchedulableForWindowCost(ctx, stickyAccount, true) &&
1264+
1265+
s.isAccountSchedulableForRPM(ctx, stickyAccount, true) { // 粘性会话窗口费用+RPM 检查
12591266
result, err := s.tryAcquireAccountSlot(ctx, stickyAccountID, stickyAccount.Concurrency)
12601267
if err == nil && result.Acquired {
12611268
// 会话数量限制检查
@@ -1409,7 +1416,9 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
14091416
s.isAccountAllowedForPlatform(account, platform, useMixed) &&
14101417
(requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) &&
14111418
s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) &&
1412-
s.isAccountSchedulableForWindowCost(ctx, account, true) { // 粘性会话窗口费用检查
1419+
s.isAccountSchedulableForWindowCost(ctx, account, true) &&
1420+
1421+
s.isAccountSchedulableForRPM(ctx, account, true) { // 粘性会话窗口费用+RPM 检查
14131422
result, err := s.tryAcquireAccountSlot(ctx, accountID, account.Concurrency)
14141423
if err == nil && result.Acquired {
14151424
// 会话数量限制检查
@@ -1475,6 +1484,10 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro
14751484
if !s.isAccountSchedulableForWindowCost(ctx, acc, false) {
14761485
continue
14771486
}
1487+
// RPM 检查（非粘性会话路径）
1488+
if !s.isAccountSchedulableForRPM(ctx, acc, false) {
1489+
continue
1490+
}
14781491
candidates = append(candidates, acc)
14791492
}
14801493

@@ -2158,6 +2171,76 @@ checkSchedulability:
21582171
return true
21592172
}
21602173

2174+
// rpmPrefetchContextKey is the context key for prefetched RPM counts.
2175+
type rpmPrefetchContextKeyType struct{}
2176+
2177+
var rpmPrefetchContextKey = rpmPrefetchContextKeyType{}
2178+
2179+
func rpmFromPrefetchContext(ctx context.Context, accountID int64) (int, bool) {
2180+
if v, ok := ctx.Value(rpmPrefetchContextKey).(map[int64]int); ok {
2181+
count, found := v[accountID]
2182+
return count, found
2183+
}
2184+
return 0, false
2185+
}
2186+
2187+
// withRPMPrefetch 批量预取所有候选账号的 RPM 计数
2188+
func (s *GatewayService) withRPMPrefetch(ctx context.Context, accounts []Account) context.Context {
2189+
if s.rpmCache == nil {
2190+
return ctx
2191+
}
2192+
2193+
var ids []int64
2194+
for i := range accounts {
2195+
if accounts[i].IsAnthropicOAuthOrSetupToken() && accounts[i].GetBaseRPM() > 0 {
2196+
ids = append(ids, accounts[i].ID)
2197+
}
2198+
}
2199+
if len(ids) == 0 {
2200+
return ctx
2201+
}
2202+
2203+
counts, err := s.rpmCache.GetRPMBatch(ctx, ids)
2204+
if err != nil {
2205+
return ctx // 失败开放
2206+
}
2207+
return context.WithValue(ctx, rpmPrefetchContextKey, counts)
2208+
}
2209+
2210+
// isAccountSchedulableForRPM 检查账号是否可根据 RPM 进行调度
2211+
// 仅适用于 Anthropic OAuth/SetupToken 账号
2212+
func (s *GatewayService) isAccountSchedulableForRPM(ctx context.Context, account *Account, isSticky bool) bool {
2213+
if !account.IsAnthropicOAuthOrSetupToken() {
2214+
return true
2215+
}
2216+
baseRPM := account.GetBaseRPM()
2217+
if baseRPM <= 0 {
2218+
return true
2219+
}
2220+
2221+
// 尝试从预取缓存获取
2222+
var currentRPM int
2223+
if count, ok := rpmFromPrefetchContext(ctx, account.ID); ok {
2224+
currentRPM = count
2225+
} else if s.rpmCache != nil {
2226+
if count, err := s.rpmCache.GetRPM(ctx, account.ID); err == nil {
2227+
currentRPM = count
2228+
}
2229+
// 失败开放：GetRPM 错误时允许调度
2230+
}
2231+
2232+
schedulability := account.CheckRPMSchedulability(currentRPM)
2233+
switch schedulability {
2234+
case WindowCostSchedulable:
2235+
return true
2236+
case WindowCostStickyOnly:
2237+
return isSticky
2238+
case WindowCostNotSchedulable:
2239+
return false
2240+
}
2241+
return true
2242+
}
2243+
21612244
// checkAndRegisterSession 检查并注册会话，用于会话数量限制
21622245
// 仅适用于 Anthropic OAuth/SetupToken 账号
21632246
// sessionID: 会话标识符（使用粘性会话的 hash）
@@ -2492,7 +2575,7 @@ func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context,
24922575
if clearSticky {
24932576
_ = s.cache.DeleteSessionAccountID(ctx, derefGroupID(groupID), sessionHash)
24942577
}
2495-
if !clearSticky && s.isAccountInGroup(account, groupID) && account.Platform == platform && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) {
2578+
if !clearSticky && s.isAccountInGroup(account, groupID) && account.Platform == platform && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) && s.isAccountSchedulableForRPM(ctx, account, true) {
24962579
if s.debugModelRoutingEnabled() {
24972580
logger.LegacyPrintf("service.gateway", "[ModelRoutingDebug] legacy routed sticky hit: group_id=%v model=%s session=%s account=%d", derefGroupID(groupID), requestedModel, shortSessionHash(sessionHash), accountID)
24982581
}
@@ -2542,6 +2625,9 @@ func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context,
25422625
if !s.isAccountSchedulableForModelSelection(ctx, acc, requestedModel) {
25432626
continue
25442627
}
2628+
if !s.isAccountSchedulableForRPM(ctx, acc, false) {
2629+
continue
2630+
}
25452631
if selected == nil {
25462632
selected = acc
25472633
continue
@@ -2592,7 +2678,7 @@ func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context,
25922678
if clearSticky {
25932679
_ = s.cache.DeleteSessionAccountID(ctx, derefGroupID(groupID), sessionHash)
25942680
}
2595-
if !clearSticky && s.isAccountInGroup(account, groupID) && account.Platform == platform && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) {
2681+
if !clearSticky && s.isAccountInGroup(account, groupID) && account.Platform == platform && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) && s.isAccountSchedulableForRPM(ctx, account, true) {
25962682
return account, nil
25972683
}
25982684
}
@@ -2631,6 +2717,9 @@ func (s *GatewayService) selectAccountForModelWithPlatform(ctx context.Context,
26312717
if !s.isAccountSchedulableForModelSelection(ctx, acc, requestedModel) {
26322718
continue
26332719
}
2720+
if !s.isAccountSchedulableForRPM(ctx, acc, false) {
2721+
continue
2722+
}
26342723
if selected == nil {
26352724
selected = acc
26362725
continue
@@ -2700,7 +2789,7 @@ func (s *GatewayService) selectAccountWithMixedScheduling(ctx context.Context, g
27002789
if clearSticky {
27012790
_ = s.cache.DeleteSessionAccountID(ctx, derefGroupID(groupID), sessionHash)
27022791
}
2703-
if !clearSticky && s.isAccountInGroup(account, groupID) && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) {
2792+
if !clearSticky && s.isAccountInGroup(account, groupID) && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) && s.isAccountSchedulableForRPM(ctx, account, true) {
27042793
if account.Platform == nativePlatform || (account.Platform == PlatformAntigravity && account.IsMixedSchedulingEnabled()) {
27052794
if s.debugModelRoutingEnabled() {
27062795
logger.LegacyPrintf("service.gateway", "[ModelRoutingDebug] legacy mixed routed sticky hit: group_id=%v model=%s session=%s account=%d", derefGroupID(groupID), requestedModel, shortSessionHash(sessionHash), accountID)
@@ -2752,6 +2841,9 @@ func (s *GatewayService) selectAccountWithMixedScheduling(ctx context.Context, g
27522841
if !s.isAccountSchedulableForModelSelection(ctx, acc, requestedModel) {
27532842
continue
27542843
}
2844+
if !s.isAccountSchedulableForRPM(ctx, acc, false) {
2845+
continue
2846+
}
27552847
if selected == nil {
27562848
selected = acc
27572849
continue
@@ -2802,7 +2894,7 @@ func (s *GatewayService) selectAccountWithMixedScheduling(ctx context.Context, g
28022894
if clearSticky {
28032895
_ = s.cache.DeleteSessionAccountID(ctx, derefGroupID(groupID), sessionHash)
28042896
}
2805-
if !clearSticky && s.isAccountInGroup(account, groupID) && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) {
2897+
if !clearSticky && s.isAccountInGroup(account, groupID) && (requestedModel == "" || s.isModelSupportedByAccountWithContext(ctx, account, requestedModel)) && s.isAccountSchedulableForModelSelection(ctx, account, requestedModel) && s.isAccountSchedulableForRPM(ctx, account, true) {
28062898
if account.Platform == nativePlatform || (account.Platform == PlatformAntigravity && account.IsMixedSchedulingEnabled()) {
28072899
return account, nil
28082900
}
@@ -2843,6 +2935,9 @@ func (s *GatewayService) selectAccountWithMixedScheduling(ctx context.Context, g
28432935
if !s.isAccountSchedulableForModelSelection(ctx, acc, requestedModel) {
28442936
continue
28452937
}
2938+
if !s.isAccountSchedulableForRPM(ctx, acc, false) {
2939+
continue
2940+
}
28462941
if selected == nil {
28472942
selected = acc
28482943
continue

0 commit comments

Comments
 (0)