231 lines
11 KiB
Go
231 lines
11 KiB
Go
package service
|
|
|
|
import (
|
|
"net/http/httptest"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/constant"
|
|
"github.com/QuantumNous/new-api/model"
|
|
perfmetrics "github.com/QuantumNous/new-api/pkg/perf_metrics"
|
|
"github.com/QuantumNous/new-api/setting"
|
|
"github.com/QuantumNous/new-api/setting/operation_setting"
|
|
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestPoolAffinityManualRoutingSwitch(t *testing.T) {
|
|
for _, change := range []string{"priority", "model_rule", "blacklist"} {
|
|
t.Run(change, func(t *testing.T) {
|
|
db := setupChannelSelectAutoGroupsTest(t)
|
|
const modelName = "manual-switch-model"
|
|
createChannelSelectAutoGroupsChannel(t, db, 6201, "default", modelName)
|
|
createChannelSelectAutoGroupsChannel(t, db, 6202, "vip", modelName)
|
|
model.InitChannelCache()
|
|
ResetGroupAvailability()
|
|
t.Cleanup(ResetGroupAvailability)
|
|
previousRedis := common.RedisEnabled
|
|
common.RedisEnabled = false
|
|
t.Cleanup(func() { common.RedisEnabled = previousRedis })
|
|
require.NoError(t, poolBindings().Purge())
|
|
t.Cleanup(func() { _ = poolBindings().Purge() })
|
|
config := operation_setting.GetChannelAffinitySetting()
|
|
previousConfig := *config
|
|
*config = operation_setting.ChannelAffinitySetting{Enabled: true, DefaultTTLSeconds: 3600, SessionMode: "strict", Rules: []operation_setting.ChannelAffinityRule{{Name: "switch-test", ModelRegex: []string{".*"}, KeySources: []operation_setting.ChannelAffinityKeySource{{Type: "request_header", Key: "Session-Id"}}, SessionMode: "strict"}}}
|
|
t.Cleanup(func() { *config = previousConfig })
|
|
previousLogDB := model.LOG_DB
|
|
model.LOG_DB = db
|
|
t.Cleanup(func() { model.LOG_DB = previousLogDB })
|
|
require.NoError(t, db.AutoMigrate(&model.Log{}))
|
|
ctx := func(switched bool) *gin.Context {
|
|
c := setupGroupPoolContext(t)
|
|
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
|
|
c.Request.Header.Set("Session-Id", "same-session")
|
|
c.Set("id", 100)
|
|
c.Set("token_id", 100)
|
|
c.Set("original_model", modelName)
|
|
policy := &model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []model.TokenGroupEntry{{Group: "default", Priority: 0}, {Group: "vip", Priority: 1}}}
|
|
if change == "blacklist" {
|
|
policy.Mode = "smart"
|
|
policy.Groups = nil
|
|
automaticMetricCache.Set(modelName, perfmetrics.QueryResult{})
|
|
}
|
|
if switched {
|
|
if change == "priority" {
|
|
policy.Groups[0].Priority = 1
|
|
policy.Groups[1].Priority = 0
|
|
} else if change == "blacklist" {
|
|
policy.ExcludedGroups = []string{"default"}
|
|
} else {
|
|
policy.ModelRules = []model.ModelRoutingRule{{Model: modelName, Groups: []model.TokenGroupEntry{{Group: "vip", Priority: 0}}}}
|
|
}
|
|
}
|
|
c.Set(model.TokenRoutingContextKey, policy)
|
|
return c
|
|
}
|
|
old := ctx(false)
|
|
ch, _, selectErr := SelectChannelForRequest(old, modelName, &RetryParam{Ctx: old, TokenGroup: "default", ModelName: modelName, Retry: common.GetPointer(0)})
|
|
require.Nil(t, selectErr)
|
|
require.NotNil(t, ch)
|
|
assert.Equal(t, 6201, ch.Id)
|
|
old.Set("channel_id", ch.Id)
|
|
common.SetContextKey(old, constant.ContextKeyAutoGroup, "default")
|
|
RequestPolicy(old).Successful = true
|
|
RecordPoolAffinity(old)
|
|
unchanged, _ := PreferredPoolChannel(ctx(false), modelName)
|
|
require.NotNil(t, unchanged)
|
|
assert.Equal(t, 6201, unchanged.Id)
|
|
|
|
details := model.NewLogOther()
|
|
appendRequestDetails(old, nil, details)
|
|
body, err := common.Marshal(details.Snapshot())
|
|
require.NoError(t, err)
|
|
log := model.Log{UserId: 100, TokenId: 100, ModelName: modelName, CorrelationKey: old.GetString("log_correlation_key"), Type: model.LogTypeConsume, CreatedAt: time.Now().Unix(), ChannelId: 6201, Group: "default", Other: string(body)}
|
|
require.NoError(t, db.Create(&log).Error)
|
|
|
|
// The edited policy must take precedence over the existing binding.
|
|
switched := ctx(true)
|
|
selected, group, selectErr := SelectChannelForRequest(switched, modelName, &RetryParam{Ctx: switched, TokenGroup: "default", ModelName: modelName, Retry: common.GetPointer(0)})
|
|
require.Nil(t, selectErr)
|
|
require.NotNil(t, selected)
|
|
assert.Equal(t, 6202, selected.Id)
|
|
assert.Equal(t, "vip", group)
|
|
switched.Set("channel_id", selected.Id)
|
|
common.SetContextKey(switched, constant.ContextKeyAutoGroup, group)
|
|
RequestPolicy(switched).Successful = true
|
|
RecordPoolAffinity(switched)
|
|
unchanged, _ = PreferredPoolChannel(ctx(true), modelName)
|
|
require.NotNil(t, unchanged)
|
|
assert.Equal(t, 6202, unchanged.Id)
|
|
|
|
// A request started before the edit may finish after the new one.
|
|
RecordPoolAffinity(old)
|
|
stale, _ := PreferredPoolChannel(ctx(true), modelName)
|
|
assert.Nil(t, stale)
|
|
|
|
// Cache expiry/restart must not resurrect the old successful log.
|
|
require.NoError(t, poolBindings().Purge())
|
|
stale, _ = PreferredPoolChannel(ctx(true), modelName)
|
|
assert.Nil(t, stale)
|
|
|
|
// A matching historical policy remains recoverable after cache expiry.
|
|
require.NoError(t, poolBindings().Purge())
|
|
recovered, _ := PreferredPoolChannel(ctx(false), modelName)
|
|
require.NotNil(t, recovered)
|
|
assert.Equal(t, 6201, recovered.Id)
|
|
|
|
// Legacy bindings carry no policy and must be selected afresh once.
|
|
require.NoError(t, poolBindings().SetWithTTL(old.GetString("pool_affinity_key"), poolBinding{Channel: 6201, Group: "default"}, time.Hour))
|
|
legacy, _ := PreferredPoolChannel(ctx(false), modelName)
|
|
assert.Nil(t, legacy)
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestAutomaticPoolKeepsHealthySession(t *testing.T) {
|
|
for _, scenario := range []string{"cheaper_recovered", "ranking_changed", "candidate_added", "degraded", "cooldown", "permission_removed", "channel_disabled", "policy_changed", "new_session", "session_off", "custom_recovery"} {
|
|
t.Run(scenario, func(t *testing.T) {
|
|
db := setupChannelSelectAutoGroupsTest(t)
|
|
const modelName = "smart-sticky-model"
|
|
createChannelSelectAutoGroupsChannel(t, db, 6301, "default", modelName)
|
|
createChannelSelectAutoGroupsChannel(t, db, 6302, "vip", modelName)
|
|
model.InitChannelCache()
|
|
oldRedis, oldLogDB := common.RedisEnabled, model.LOG_DB
|
|
common.RedisEnabled, model.LOG_DB = false, nil
|
|
config := operation_setting.GetChannelAffinitySetting()
|
|
oldConfig := *config
|
|
*config = operation_setting.ChannelAffinitySetting{Enabled: true, DefaultTTLSeconds: 3600, SessionMode: "prefer", Rules: []operation_setting.ChannelAffinityRule{{Name: "smart-sticky", ModelRegex: []string{".*"}, KeySources: []operation_setting.ChannelAffinityKeySource{{Type: "request_header", Key: "Session-Id"}}, SessionMode: "prefer"}}}
|
|
require.NoError(t, poolBindings().Purge())
|
|
require.NoError(t, poolCooldowns().Purge())
|
|
automaticMetricCache.Set(modelName, perfmetrics.QueryResult{})
|
|
t.Cleanup(func() {
|
|
*config = oldConfig
|
|
common.RedisEnabled, model.LOG_DB = oldRedis, oldLogDB
|
|
_ = poolBindings().Purge()
|
|
_ = poolCooldowns().Purge()
|
|
automaticMetricCache.Purge()
|
|
})
|
|
ctx := func() *gin.Context {
|
|
c := setupGroupPoolContext(t)
|
|
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
|
|
c.Request.Header.Set("Session-Id", "same-session")
|
|
c.Set("id", 101)
|
|
c.Set("token_id", 102)
|
|
c.Set("original_model", modelName)
|
|
policy := &model.TokenRoutingPolicy{Version: 1, Mode: "smart", AutoPreference: "cost", SkipDegradedNonAstra: true}
|
|
if scenario == "custom_recovery" {
|
|
policy.Mode = "custom"
|
|
policy.Groups = []model.TokenGroupEntry{{Group: "default", Priority: 0}, {Group: "vip", Priority: 1}}
|
|
}
|
|
c.Set(model.TokenRoutingContextKey, policy)
|
|
c.Set("pool_group_fingerprints", map[string]AstraSamplingHistory{})
|
|
c.Set("pool_group_visions", map[string]DetectionHistory{})
|
|
return c
|
|
}
|
|
// A temporary failure sends the session to the more expensive group.
|
|
first := ctx()
|
|
cheapKey := poolCooldownKey(first, modelName, "default")
|
|
require.NoError(t, poolCooldowns().SetWithTTL(cheapKey, poolCooldown{Until: time.Now().Add(time.Minute).UnixMilli()}, 0))
|
|
if scenario == "candidate_added" {
|
|
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"vip":"VIP"}`))
|
|
}
|
|
ch, group, selectErr := SelectChannelForRequest(first, modelName, &RetryParam{Ctx: first, ModelName: modelName, Retry: common.GetPointer(0)})
|
|
require.Nil(t, selectErr)
|
|
require.NotNil(t, ch)
|
|
require.Equal(t, "vip", group)
|
|
first.Set("channel_id", ch.Id)
|
|
RequestPolicy(first).Successful = true
|
|
RecordPoolAffinity(first)
|
|
require.NoError(t, poolCooldowns().SetWithTTL(cheapKey, poolCooldown{Until: time.Now().Add(-time.Second).UnixMilli()}, 0))
|
|
next := ctx()
|
|
want := "vip"
|
|
switch scenario {
|
|
case "ranking_changed":
|
|
// The initial pool already placed default first. Reverse it to prove ranking changes are irrelevant.
|
|
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":4,"vip":3}`))
|
|
case "candidate_added":
|
|
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP"}`))
|
|
case "degraded":
|
|
next.Set("pool_group_visions", map[string]DetectionHistory{"vip": {Quality: "degraded"}})
|
|
want = "default"
|
|
case "cooldown":
|
|
require.NoError(t, poolCooldowns().SetWithTTL(poolCooldownKey(next, modelName, "vip"), poolCooldown{Until: time.Now().Add(time.Minute).UnixMilli()}, 0))
|
|
want = "default"
|
|
case "permission_removed":
|
|
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default"}`))
|
|
want = "default"
|
|
case "channel_disabled":
|
|
require.NoError(t, db.Model(&model.Channel{}).Where("id = ?", 6302).Update("status", common.ChannelStatusAutoDisabled).Error)
|
|
model.InitChannelCache()
|
|
want = "default"
|
|
case "policy_changed":
|
|
RequestRoutingPolicy(next).AutoPreference = "stable"
|
|
want = "default"
|
|
case "new_session":
|
|
next.Request.Header.Set("Session-Id", "new-session")
|
|
want = "default"
|
|
case "session_off":
|
|
config.Rules[0].SessionMode = "off"
|
|
config.SessionMode = "off"
|
|
want = "default"
|
|
case "custom_recovery":
|
|
want = "default"
|
|
}
|
|
selected, selectedGroup, selectErr := SelectChannelForRequest(next, modelName, &RetryParam{Ctx: next, ModelName: modelName, Retry: common.GetPointer(0)})
|
|
require.Nil(t, selectErr)
|
|
require.NotNil(t, selected)
|
|
assert.Equal(t, want, selectedGroup)
|
|
if want == "vip" {
|
|
assert.Equal(t, 6302, selected.Id)
|
|
state, found, err := poolCooldowns().Get(cheapKey)
|
|
require.NoError(t, err)
|
|
require.True(t, found)
|
|
assert.Empty(t, state.ProbeID, "keeping a healthy route must not claim a cheaper group's recovery probe")
|
|
}
|
|
})
|
|
}
|
|
}
|