96 lines
4.5 KiB
Go
96 lines
4.5 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/QuantumNous/new-api/model"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/stretchr/testify/require"
|
|
"net/http/httptest"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestPoolWithoutAvailableBackupDoesNotCoolDown(t *testing.T) {
|
|
db := setupChannelSelectAutoGroupsTest(t)
|
|
createChannelSelectAutoGroupsChannel(t, db, 7801, "default", "m")
|
|
createChannelSelectAutoGroupsChannel(t, db, 7802, "vip", "other")
|
|
model.InitChannelCache()
|
|
ResetGroupAvailability()
|
|
t.Cleanup(func() { ResetGroupAvailability(); _ = poolCooldowns().Purge() })
|
|
makeCtx := func() *gin.Context {
|
|
c := setupGroupPoolContext(t)
|
|
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
|
|
c.Set("token_id", 7801)
|
|
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []model.TokenGroupEntry{{Group: "default"}, {Group: "vip", Priority: 1}}})
|
|
return c
|
|
}
|
|
c := makeCtx()
|
|
RecordPoolRetryableFailure(c, "m", "default")
|
|
RecordPoolRetryableFailure(c, "m", "default")
|
|
_, found, err := poolCooldowns().Get(poolCooldownKey(c, "m", "default"))
|
|
require.NoError(t, err)
|
|
require.False(t, found)
|
|
// An old cooldown and global health state must not strand a single eligible group.
|
|
require.NoError(t, poolCooldowns().SetWithTTL(poolCooldownKey(c, "m", "default"), poolCooldown{Until: time.Now().Add(10 * time.Minute).UnixMilli()}, 0))
|
|
for i := 0; i < 4; i++ {
|
|
RecordGroupOutcome("default", false)
|
|
}
|
|
next := makeCtx()
|
|
ch, group, err := selectTokenRoutingPool(&RetryParam{Ctx: next, ModelName: "m", Retry: common.GetPointer(0)})
|
|
require.NoError(t, err)
|
|
require.NotNil(t, ch)
|
|
require.Equal(t, "default", group)
|
|
AppendUsedChannel(next, ch.Id)
|
|
RecordPoolRetryableFailure(next, "m", "default")
|
|
ch, _, err = selectTokenRoutingPool(&RetryParam{Ctx: next, ModelName: "m", Retry: common.GetPointer(1)})
|
|
require.NoError(t, err)
|
|
require.NotNil(t, ch)
|
|
AppendUsedChannel(next, ch.Id)
|
|
RecordPoolRetryableFailure(next, "m", "default")
|
|
ch, _, err = selectTokenRoutingPool(&RetryParam{Ctx: next, ModelName: "m", Retry: common.GetPointer(2)})
|
|
require.NoError(t, err)
|
|
require.Nil(t, ch, "do not loop indefinitely inside one request")
|
|
}
|
|
func TestAutomaticSpeedUsesLatestExactModelStreamingTest(t *testing.T) {
|
|
db := setupChannelSelectAutoGroupsTest(t)
|
|
require.NoError(t, db.AutoMigrate(&model.Log{}))
|
|
previous := model.LOG_DB
|
|
model.LOG_DB = db
|
|
t.Cleanup(func() { model.LOG_DB = previous })
|
|
now := time.Now().Unix()
|
|
samples := []model.Log{
|
|
{ChannelId: 1, Type: 2, TokenId: 0, Content: "模型测试", ModelName: "m", IsStream: true, CreatedAt: now - 20, Other: `{"frt":500}`},
|
|
{ChannelId: 1, Type: 2, TokenId: 0, Content: "模型测试", ModelName: "m", IsStream: true, CreatedAt: now - 10, Other: `{"frt":60000}`},
|
|
{ChannelId: 1, Type: 2, TokenId: 9, Content: "模型测试", ModelName: "m", IsStream: true, CreatedAt: now, Other: `{"frt":1}`},
|
|
{ChannelId: 2, Type: 2, TokenId: 0, Content: "模型测试", ModelName: "m", IsStream: true, CreatedAt: now - 5, Other: `{"frt":1200}`},
|
|
{ChannelId: 2, Type: 2, TokenId: 0, Content: "模型测试", ModelName: "other", IsStream: true, CreatedAt: now, Other: `{"frt":10}`},
|
|
{ChannelId: 3, Type: 2, TokenId: 0, Content: "模型测试", ModelName: "m", IsStream: false, CreatedAt: now, Other: `{"frt":10}`},
|
|
{ChannelId: 4, Type: 2, TokenId: 0, Content: "模型测试", ModelName: "m", IsStream: true, CreatedAt: now - 2*3600, Other: `{"frt":10}`},
|
|
}
|
|
require.NoError(t, db.Create(&samples).Error)
|
|
speeds := latestModelTestSpeeds(context.Background(), "m")
|
|
require.Equal(t, map[int]float64{1: 60000, 2: 1200}, speeds)
|
|
}
|
|
|
|
func TestProbeAndStreamingCooldownsAreIsolated(t *testing.T) {
|
|
db := setupChannelSelectAutoGroupsTest(t)
|
|
createChannelSelectAutoGroupsChannel(t, db, 7901, "default", "m")
|
|
createChannelSelectAutoGroupsChannel(t, db, 7902, "vip", "m")
|
|
model.InitChannelCache()
|
|
t.Cleanup(func() { _ = poolCooldowns().Purge() })
|
|
makeCtx := func(stream bool) *gin.Context {
|
|
c := setupGroupPoolContext(t)
|
|
c.Request = httptest.NewRequest("POST", "/v1/responses", nil)
|
|
c.Set("token_id", 7901)
|
|
c.Set("pool_request_stream", stream)
|
|
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []model.TokenGroupEntry{{Group: "default"}, {Group: "vip", Priority: 1}}})
|
|
return c
|
|
}
|
|
probe := makeCtx(false)
|
|
RecordPoolRetryableFailure(probe, "m", "default")
|
|
RecordPoolRetryableFailure(probe, "m", "default")
|
|
require.True(t, poolGroupSkipped(makeCtx(false), "m", "default"))
|
|
require.False(t, poolGroupSkipped(makeCtx(true), "m", "default"))
|
|
}
|