feat: sync deployed v1.0.0-rc.38-custom.67
This commit is contained in:
cneicy committed
1 parent
d09216bbde
commit
8e6ed0317b
154 files changed
+39540
-20509
No files matched your search
@@ -1,16 +1,160 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"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"
|
||||
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/require"
|
||||
"math"
|
||||
"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, 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"} {
|
||||
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 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, 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))
|
||||
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"
|
||||
|
||||
Reference in new issue
Block a user