feat: publish sanitized custom.97 source
This commit is contained in:
1022 files changed
+69903
-19207
No files matched your search
@@ -6,13 +6,16 @@ import (
|
||||
"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"
|
||||
@@ -80,8 +83,8 @@ func TestTokenRoutingSingleChannelRetryCooldownAndRecovery(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestAutomaticGroupPreferenceRanking(t *testing.T) {
|
||||
original := []automaticCandidate{{group: "cheap", cost: .1, ttft: 900, success: 90, ratio: .2}, {group: "fast", cost: .2, ttft: 100, success: 95, ratio: .3}, {group: "stable", cost: .3, ttft: 400, 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", "stable": "stable"} {
|
||||
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)
|
||||
@@ -93,6 +96,45 @@ func TestAutomaticGroupPreferenceRanking(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
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 {
|
||||
@@ -132,13 +174,13 @@ func TestAutomaticGroupUsesAllAuthorizedModelGroups(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestPoolFailoverRespectsRetrySafetyBoundaries(t *testing.T) {
|
||||
for _, code := range []int{400, 401, 429, 500, 502, 503, 504, 524} {
|
||||
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.NewErrorWithStatusCode(errors.New("test upstream"), apiTypes.ErrorCodeBadResponseStatusCode, code)
|
||||
require.Equal(t, code == 500 || code == 502 || code == 503, IsPoolFailoverError(c, err))
|
||||
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)
|
||||
@@ -211,3 +253,129 @@ func TestTokenRoutingPoolPreferenceExclusionAndAffinityIsolation(t *testing.T) {
|
||||
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)
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user