Files
new-api-custom/service/token_routing_test.go
T

382 lines
18 KiB
Go

package service
import (
"context"
"errors"
"fmt"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/model"
perfmetrics "github.com/QuantumNous/new-api/pkg/perf_metrics"
apiTypes "github.com/QuantumNous/new-api/relaykit/types"
"github.com/QuantumNous/new-api/setting/operation_setting"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"math"
"net/http"
"net/http/httptest"
"slices"
"testing"
"time"
)
func TestTokenRoutingSingleChannelRetryCooldownAndRecovery(t *testing.T) {
db := setupChannelSelectAutoGroupsTest(t)
createChannelSelectAutoGroupsChannel(t, db, 6911, "default", "retry-model")
createChannelSelectAutoGroupsChannel(t, db, 6912, "vip", "retry-model")
model.InitChannelCache()
ResetGroupAvailability()
t.Cleanup(func() { ResetGroupAvailability(); _ = poolCooldowns().Purge() })
makeCtx := func(token int) *gin.Context {
c := setupGroupPoolContext(t)
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
c.Set("id", 99)
c.Set("token_id", token)
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []model.TokenGroupEntry{{Group: "default", Priority: 0}, {Group: "vip", Priority: 1}}})
return c
}
selectID := func(c *gin.Context) int {
t.Helper()
ch, _, err := selectTokenRoutingPool(&RetryParam{Ctx: c, ModelName: "retry-model", Retry: common.GetPointer(0)})
require.NoError(t, err)
require.NotNil(t, ch)
AppendUsedChannel(c, ch.Id)
return ch.Id
}
c := makeCtx(991)
require.Equal(t, 6911, selectID(c))
RecordPoolRetryableFailure(c, "retry-model", "default")
require.Equal(t, 6911, selectID(c), "a one-channel group must actually get its second attempt")
RecordPoolRetryableFailure(c, "retry-model", "default")
require.Equal(t, 6912, selectID(c))
require.Equal(t, 6912, selectID(makeCtx(991)))
require.Equal(t, 6911, selectID(makeCtx(992)))
key := poolCooldownKey(c, "retry-model", "default")
state, found, err := poolCooldowns().Get(key)
require.NoError(t, err)
require.True(t, found)
require.InDelta(t, int64((10*time.Minute)/time.Millisecond), state.Until-time.Now().UnixMilli(), 2000)
require.NoError(t, poolCooldowns().SetWithTTL(key, poolCooldown{Until: time.Now().Add(-time.Second).UnixMilli()}, 0))
trial := makeCtx(991)
require.Equal(t, 6911, selectID(trial))
require.Equal(t, 6912, selectID(makeCtx(991)), "concurrent request waits for recovery probe")
RecordPoolRetryableFailure(trial, "retry-model", "default")
require.Equal(t, 6912, selectID(trial), "failed probe switches without another primary attempt")
require.NoError(t, poolCooldowns().SetWithTTL(key, poolCooldown{Until: time.Now().Add(-time.Second).UnixMilli()}, 0))
recovery := makeCtx(991)
require.Equal(t, 6911, selectID(recovery))
RecordPoolSuccess(recovery, "retry-model", "default")
_, found, err = poolCooldowns().Get(key)
require.NoError(t, err)
require.False(t, found)
next := makeCtx(991)
require.Equal(t, 6911, selectID(next))
RecordPoolRetryableFailure(next, "retry-model", "default")
require.Equal(t, 6911, selectID(next), "recovered group returns to two-attempt behavior")
mixed := makeCtx(993)
require.Equal(t, 6911, selectID(mixed))
RecordPoolRetryableFailure(mixed, "retry-model", "default")
require.Equal(t, 6911, selectID(mixed))
require.Equal(t, 6912, selectID(mixed), "the one same-group retry must not repeat when a different error triggers another selection")
}
func TestAutomaticGroupPreferenceRanking(t *testing.T) {
original := []automaticCandidate{{group: "cheap", cost: .1, ttft: 900, throughput: 20, success: 90, ratio: .2}, {group: "fast", cost: .2, ttft: 100, throughput: 30, success: 95, ratio: .3}, {group: "stable", cost: .3, ttft: 400, throughput: 80, success: 99, ratio: .4}, {group: "unknown", cost: math.Inf(1), ttft: math.Inf(1), success: -1, ratio: .01}}
for preference, want := range map[string]string{"cost": "unknown", "speed": "fast", "throughput": "stable", "stable": "stable"} {
t.Run(preference, func(t *testing.T) {
items := slices.Clone(original)
sortAutomaticCandidates(items, preference)
require.Equal(t, want, items[0].group)
if preference != "cost" {
require.Equal(t, "unknown", items[len(items)-1].group)
}
})
}
}
func TestAutomaticThroughputUsesRequestedModelCallsAndTestFallback(t *testing.T) {
db := setupChannelSelectAutoGroupsTest(t)
require.NoError(t, db.AutoMigrate(&model.Log{}))
previousLogs := model.LOG_DB
model.LOG_DB = db
t.Cleanup(func() { model.LOG_DB = previousLogs; groupModelPerformanceCache.Purge(); automaticMetricCache.Purge() })
for _, group := range []string{"default", "vip"} {
id := 6951
if group == "vip" {
id = 6952
}
createChannelSelectAutoGroupsChannel(t, db, id, group, "throughput-model")
}
model.InitChannelCache()
logs := []model.Log{
{Type: model.LogTypeConsume, TokenId: 1, Group: "default", ModelName: "throughput-model", IsStream: true, UseTime: 2, CompletionTokens: 20, Other: `{"frt":100}`},
{Type: model.LogTypeConsume, ChannelId: 6952, ModelName: "throughput-model", Content: "Model test", IsStream: true, UseTime: 2, CompletionTokens: 100, Other: `{"frt":500}`},
{Type: model.LogTypeConsume, TokenId: 1, Group: "default", ModelName: "another-model", IsStream: true, UseTime: 1, CompletionTokens: 999, Other: `{"frt":1}`},
}
require.NoError(t, db.Create(&logs).Error)
automaticMetricCache.Set("throughput-model", perfmetrics.QueryResult{})
makeCtx := func() *gin.Context {
c := setupGroupPoolContext(t)
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "smart", AutoPreference: "throughput"})
return c
}
c := makeCtx()
assert.Equal(t, []model.TokenGroupEntry{{Group: "vip", Priority: 0}, {Group: "default", Priority: 1}}, ResolveRequestPool(c, "throughput-model"))
// New traffic is reflected after snapshot refresh, but never reorders retries
// of an already-ranked request.
require.NoError(t, db.Create(&model.Log{Type: model.LogTypeConsume, TokenId: 1, Group: "vip", ModelName: "throughput-model", IsStream: true, UseTime: 10, CompletionTokens: 5, Other: `{"frt":100}`}).Error)
groupModelPerformanceCache.Purge()
assert.Equal(t, "vip", ResolveRequestPool(c, "throughput-model")[0].Group)
assert.Equal(t, "default", ResolveRequestPool(makeCtx(), "throughput-model")[0].Group)
policy := &model.TokenRoutingPolicy{Version: 1, Mode: "smart", AutoPreference: "throughput"}
require.NoError(t, policy.Validate(5, func(string) bool { return true }))
}
func TestAutomaticGroupUsesAllAuthorizedModelGroups(t *testing.T) {
db := setupChannelSelectAutoGroupsTest(t)
for _, fixture := range []struct {
id int
group, name string
}{{6921, "default", "smart-model"}, {6922, "vip", "smart-model"}, {6923, "secret", "smart-model"}, {6924, "default", "another-model"}} {
createChannelSelectAutoGroupsChannel(t, db, fixture.id, fixture.group, fixture.name)
}
model.InitChannelCache()
groupModelCostSnapshot.Lock()
oldCosts, oldTime := groupModelCostSnapshot.values, groupModelCostSnapshot.updated
cheap, expensive := 0.1, 0.3
groupModelCostSnapshot.values = map[string]map[string]model.GroupModelCost{"default": {"smart-model": {USDPerMillion: &expensive}}, "vip": {"smart-model": {USDPerMillion: &cheap}}}
groupModelCostSnapshot.updated = time.Now()
groupModelCostSnapshot.Unlock()
automaticMetricCache.Set("smart-model", perfmetrics.QueryResult{})
t.Cleanup(func() {
groupModelCostSnapshot.Lock()
groupModelCostSnapshot.values, groupModelCostSnapshot.updated = oldCosts, oldTime
groupModelCostSnapshot.Unlock()
automaticMetricCache.Purge()
})
makeCtx := func() *gin.Context {
c := setupGroupPoolContext(t)
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "smart", AutoPreference: "cost"})
return c
}
result := ResolveRequestPool(makeCtx(), "smart-model")
require.Equal(t, []model.TokenGroupEntry{{Group: "default", Priority: 0}, {Group: "vip", Priority: 1}}, result)
list := ResolveRequestPool(makeCtx(), "")
require.Len(t, list, 2)
for _, entry := range list {
require.NotEqual(t, "secret", entry.Group)
}
require.Empty(t, ResolveRequestPool(makeCtx(), "unknown-model"))
}
func TestPoolFailoverRespectsRetrySafetyBoundaries(t *testing.T) {
for _, code := range []int{400, 401, 403, 429, 500, 502, 503, 504, 524} {
t.Run(fmt.Sprint(code), func(t *testing.T) {
c := setupGroupPoolContext(t)
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "smart"})
err := apiTypes.NewOpenAIError(errors.New("test upstream"), apiTypes.ErrorCodeBadResponseStatusCode, code)
require.Equal(t, code == 403 || code == 500 || code == 502 || code == 503, IsPoolFailoverError(c, err))
ctx, cancel := context.WithCancel(c.Request.Context())
cancel()
c.Request = c.Request.WithContext(ctx)
require.False(t, IsPoolFailoverError(c, err))
})
}
c := setupGroupPoolContext(t)
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "smart"})
_, err := c.Writer.WriteString("already streaming")
require.NoError(t, err)
failure := apiTypes.NewErrorWithStatusCode(errors.New("test upstream"), apiTypes.ErrorCodeBadResponseStatusCode, 500)
require.False(t, IsPoolFailoverError(c, failure))
require.Equal(t, "stop", DecideRelayRetry(c, failure, 5).Action)
}
func TestTokenRoutingPoolPreferenceExclusionAndAffinityIsolation(t *testing.T) {
db := setupChannelSelectAutoGroupsTest(t)
const modelName = "pool-model"
createChannelSelectAutoGroupsChannel(t, db, 6101, "default", modelName)
createChannelSelectAutoGroupsChannel(t, db, 6102, "vip", modelName)
model.InitChannelCache()
ResetGroupAvailability()
t.Cleanup(ResetGroupAvailability)
config := operation_setting.GetChannelAffinitySetting()
previous := *config
*config = operation_setting.ChannelAffinitySetting{Enabled: true, DefaultTTLSeconds: 3600, SessionMode: "prefer", Rules: []operation_setting.ChannelAffinityRule{{Name: "pool-test", ModelRegex: []string{".*"}, KeySources: []operation_setting.ChannelAffinityKeySource{{Type: "request_header", Key: "Session-Id"}}, SessionMode: "prefer"}}}
t.Cleanup(func() { *config = previous; _ = poolBindings().Purge() })
ctx := func(token int, session string) *gin.Context {
c := setupGroupPoolContext(t)
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
c.Request.Header.Set("Session-Id", session)
c.Set("id", 100)
c.Set("token_id", token)
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []model.TokenGroupEntry{{Group: "default", Priority: 0}, {Group: "vip", Priority: 1}}, ModelRules: []model.ModelRoutingRule{{Model: modelName, Groups: []model.TokenGroupEntry{{Group: "vip", Priority: 0}}}}})
return c
}
c := ctx(100, "session-a")
param := &RetryParam{Ctx: c, TokenGroup: "default", ModelName: modelName, Retry: common.GetPointer(0)}
first, group, err := SelectChannelForRequest(c, modelName, param)
require.Nil(t, err)
require.NotNil(t, first)
require.Equal(t, 6102, first.Id)
require.Equal(t, "vip", group)
AppendUsedChannel(c, first.Id)
second, group, err2 := CacheGetRandomSatisfiedChannel(param)
require.NoError(t, err2)
require.NotNil(t, second)
require.Equal(t, 6101, second.Id)
c.Set("channel_id", second.Id)
common.SetContextKey(c, constant.ContextKeyAutoGroup, group)
RequestPolicy(c).Successful = true
RecordPoolAffinity(c)
next := ctx(100, "session-a")
bound, boundGroup := PreferredPoolChannel(next, modelName)
require.NotNil(t, bound)
require.Equal(t, 6101, bound.Id)
require.Equal(t, "default", boundGroup)
isolated, _ := PreferredPoolChannel(ctx(101, "session-a"), modelName)
require.Nil(t, isolated)
isolated, _ = PreferredPoolChannel(ctx(100, "session-b"), modelName)
require.Nil(t, isolated)
isolated, _ = PreferredPoolChannel(ctx(100, "session-a"), "another-model")
require.Nil(t, isolated)
ClearPoolAffinity(next)
bound, _ = PreferredPoolChannel(ctx(100, "session-a"), modelName)
require.Nil(t, bound)
AppendUsedChannel(c, second.Id)
last, _, err2 := CacheGetRandomSatisfiedChannel(param)
require.NoError(t, err2)
require.Nil(t, last)
}
func TestPoolUpstreamForbiddenFailover(t *testing.T) {
for _, mode := range []string{"smart", "custom"} {
t.Run(mode, func(t *testing.T) {
db := setupChannelSelectAutoGroupsTest(t)
const modelName = "forbidden-pool-model"
createChannelSelectAutoGroupsChannel(t, db, 6931, "default", modelName)
createChannelSelectAutoGroupsChannel(t, db, 6932, "vip", modelName)
model.InitChannelCache()
config := operation_setting.GetChannelAffinitySetting()
oldConfig := *config
oldRedis, oldLogDB := common.RedisEnabled, model.LOG_DB
common.RedisEnabled, model.LOG_DB = false, nil
*config = operation_setting.ChannelAffinitySetting{Enabled: true, DefaultTTLSeconds: 3600, SessionMode: "strict", Rules: []operation_setting.ChannelAffinityRule{{Name: "forbidden-pool", ModelRegex: []string{".*"}, KeySources: []operation_setting.ChannelAffinityKeySource{{Type: "request_header", Key: "Session-Id"}}, SessionMode: "strict"}}}
automaticMetricCache.Set(modelName, perfmetrics.QueryResult{})
require.NoError(t, poolBindings().Purge())
require.NoError(t, poolCooldowns().Purge())
t.Cleanup(func() {
_ = poolBindings().Purge()
_ = poolCooldowns().Purge()
automaticMetricCache.Purge()
*config = oldConfig
common.RedisEnabled, model.LOG_DB = oldRedis, oldLogDB
})
makeCtx := func() *gin.Context {
c := setupGroupPoolContext(t)
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
c.Request.Header.Set("Session-Id", "forbidden-session")
common.SetContextKey(c, constant.ContextKeyTokenGroup, "auto")
c.Set("id", 103)
c.Set("token_id", 104)
c.Set("original_model", modelName)
policy := &model.TokenRoutingPolicy{Version: 1, Mode: mode, AutoPreference: "cost", AllowDegraded: true}
if mode == "custom" {
policy.Groups = []model.TokenGroupEntry{{Group: "default", Priority: 0}, {Group: "vip", Priority: 1}}
}
c.Set(model.TokenRoutingContextKey, policy)
return c
}
// Establish the same successful session binding used by real requests.
first := makeCtx()
channel, group, selectionErr := SelectChannelForRequest(first, modelName, &RetryParam{Ctx: first, ModelName: modelName, Retry: common.GetPointer(0)})
require.Nil(t, selectionErr)
require.NotNil(t, channel)
require.Equal(t, "default", group)
first.Set("channel_id", channel.Id)
RequestPolicy(first).Successful = true
RecordPoolAffinity(first)
current := makeCtx()
param := &RetryParam{Ctx: current, ModelName: modelName, Retry: common.GetPointer(0)}
channel, group, selectionErr = SelectChannelForRequest(current, modelName, param)
require.Nil(t, selectionErr)
require.Equal(t, "default", group)
require.True(t, ShouldSkipRetryAfterChannelAffinityFailure(current))
budget := PoolRetryBudget(current, modelName)
for attempt := range 2 {
require.Equal(t, 6931, channel.Id)
AppendUsedChannel(current, channel.Id)
response := httptest.NewRecorder()
response.WriteHeader(http.StatusForbidden)
_, err := response.WriteString(`{"error":{"message":"upstream access denied","type":"permission_error","code":"forbidden"}}`)
require.NoError(t, err)
upstreamErr := RelayErrorHandler(current, response.Result(), false)
require.True(t, IsPoolFailoverError(current, upstreamErr))
RecordPoolRetryableFailure(current, modelName, group)
require.Equal(t, "retry", DecideRelayRetry(current, upstreamErr, budget-attempt).Action)
param.IncreaseRetry()
channel, group, selectionErr = SelectChannelForRequest(current, modelName, param)
require.Nil(t, selectionErr)
require.NotNil(t, channel)
}
assert.Equal(t, "vip", group)
assert.Equal(t, 6932, channel.Id)
current.Set("channel_id", channel.Id)
RequestPolicy(current).Successful = true
RecordPoolAffinity(current)
next := makeCtx()
bound, boundGroup := PreferredPoolChannel(next, modelName)
require.NotNil(t, bound)
assert.Equal(t, "vip", boundGroup)
assert.Equal(t, 6932, bound.Id)
assert.True(t, poolGroupSkipped(next, modelName, "default"))
})
}
}
func TestPoolForbiddenRetryBoundaries(t *testing.T) {
for _, scenario := range []string{"local", "skip_retry", "disabled_status", "pinned", "response_started", "cancelled", "no_policy", "budget_exhausted"} {
t.Run(scenario, func(t *testing.T) {
c := setupGroupPoolContext(t)
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
common.SetContextKey(c, constant.ContextKeyTokenGroup, "auto")
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "smart"})
c.Set(ginKeyChannelAffinitySkipRetry, true)
upstreamErr := apiTypes.NewOpenAIError(errors.New("upstream forbidden"), apiTypes.ErrorCodeBadResponseStatusCode, http.StatusForbidden)
switch scenario {
case "local":
upstreamErr = apiTypes.NewErrorWithStatusCode(errors.New("local access denied"), apiTypes.ErrorCodeAccessDenied, http.StatusForbidden)
case "skip_retry":
upstreamErr = apiTypes.NewOpenAIError(errors.New("upstream forbidden"), apiTypes.ErrorCodeBadResponseStatusCode, http.StatusForbidden, apiTypes.ErrOptionWithSkipRetry())
case "disabled_status":
previous := operation_setting.AutomaticRetryStatusCodesToString()
require.NoError(t, operation_setting.AutomaticRetryStatusCodesFromString("500-503"))
t.Cleanup(func() { require.NoError(t, operation_setting.AutomaticRetryStatusCodesFromString(previous)) })
case "pinned":
GetChannelConstraints(c).AddPin(dto.ChannelPin{ChannelId: 6931, Source: dto.PinSourceToken, RetryMode: dto.PinRetrySingleAttempt})
case "response_started":
_, err := c.Writer.WriteString("already streaming")
require.NoError(t, err)
case "cancelled":
ctx, cancel := context.WithCancel(c.Request.Context())
cancel()
c.Request = c.Request.WithContext(ctx)
case "no_policy":
c.Set(model.TokenRoutingContextKey, (*model.TokenRoutingPolicy)(nil))
}
assert.Equal(t, scenario == "budget_exhausted", IsPoolFailoverError(c, upstreamErr))
remaining := 3
if scenario == "budget_exhausted" {
remaining = 0
}
assert.Equal(t, "stop", DecideRelayRetry(c, upstreamErr, remaining).Action)
})
}
}