117 lines
6.3 KiB
Go
117 lines
6.3 KiB
Go
package model
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"os"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/QuantumNous/new-api/common"
|
|
"github.com/glebarez/sqlite"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
"gorm.io/driver/mysql"
|
|
"gorm.io/driver/postgres"
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/schema"
|
|
)
|
|
|
|
func TestGroupModelPerformanceDatabaseMatrix(t *testing.T) {
|
|
for _, dialect := range []string{"sqlite", "mysql", "postgres"} {
|
|
t.Run(dialect, func(t *testing.T) {
|
|
var driver gorm.Dialector
|
|
switch dialect {
|
|
case "sqlite":
|
|
driver = sqlite.Open(":memory:")
|
|
case "mysql":
|
|
if os.Getenv("TEST_MYSQL_DSN") == "" {
|
|
t.Skip("TEST_MYSQL_DSN unset")
|
|
}
|
|
driver = mysql.Open(os.Getenv("TEST_MYSQL_DSN"))
|
|
case "postgres":
|
|
if os.Getenv("TEST_POSTGRES_DSN") == "" {
|
|
t.Skip("TEST_POSTGRES_DSN unset")
|
|
}
|
|
driver = postgres.Open(os.Getenv("TEST_POSTGRES_DSN"))
|
|
}
|
|
db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("performance_%d_", time.Now().UnixNano())}})
|
|
require.NoError(t, err)
|
|
conn, err := db.DB()
|
|
require.NoError(t, err)
|
|
conn.SetMaxOpenConns(1)
|
|
previousDB, previousLogs := DB, LOG_DB
|
|
DB, LOG_DB = db, db
|
|
t.Cleanup(func() {
|
|
for _, table := range []any{&Log{}, &Ability{}, &Channel{}} {
|
|
require.NoError(t, db.Migrator().DropTable(table))
|
|
}
|
|
DB, LOG_DB = previousDB, previousLogs
|
|
require.NoError(t, conn.Close())
|
|
})
|
|
query := "select version()"
|
|
if dialect == "sqlite" {
|
|
query = "select sqlite_version()"
|
|
}
|
|
var version string
|
|
require.NoError(t, db.Raw(query).Scan(&version).Error)
|
|
t.Logf("database version: %s", version)
|
|
require.NoError(t, db.AutoMigrate(&Log{}, &Ability{}, &Channel{}))
|
|
require.NoError(t, db.Create(&Channel{Id: 1, Status: common.ChannelStatusEnabled, Group: "primary,backup"}).Error)
|
|
require.NoError(t, db.Create(&Channel{Id: 2, Status: common.ChannelStatusEnabled, Group: "primary"}).Error)
|
|
require.NoError(t, db.Create(&Channel{Id: 3, Status: common.ChannelStatusAutoDisabled, Group: "primary"}).Error)
|
|
for _, channel := range []int{1, 2, 3} {
|
|
require.NoError(t, db.Create(&Ability{ChannelId: channel, Group: "primary", Model: "m", Enabled: true}).Error)
|
|
}
|
|
ctx := context.Background()
|
|
empty, err := GetGroupModelPerformance(ctx, "primary", "m")
|
|
require.NoError(t, err)
|
|
assert.Nil(t, empty.Throughput)
|
|
assert.Nil(t, empty.FirstToken)
|
|
logs := []Log{
|
|
{ChannelId: 1, Type: LogTypeConsume, Content: "模型测试", ModelName: "m", IsStream: true, CreatedAt: 1, UseTime: 10, CompletionTokens: 100, Other: `{"frt":1000}`},
|
|
{ChannelId: 1, Type: LogTypeConsume, Content: "Model test", ModelName: "m", IsStream: true, CreatedAt: 2, UseTime: 1, CompletionTokens: 20, Other: `{"frt":500,"latency_ms":500}`},
|
|
{ChannelId: 2, Type: LogTypeConsume, Content: "Model test", ModelName: "m", IsStream: true, CreatedAt: 2, UseTime: 10, CompletionTokens: 200, Other: `{"frt":1500}`},
|
|
{ChannelId: 3, Type: LogTypeConsume, Content: "Model test", ModelName: "m", IsStream: true, CreatedAt: 3, UseTime: 1, CompletionTokens: 999, Other: `{"frt":1}`},
|
|
{ChannelId: 1, Type: LogTypeConsume, Content: "Model test", ModelName: "other", IsStream: true, CreatedAt: 3, UseTime: 1, CompletionTokens: 999, Other: `{"frt":1}`},
|
|
{ChannelId: 1, Type: LogTypeConsume, Content: "Model test", ModelName: "m", IsStream: true, CreatedAt: 4, UseTime: 1, CompletionTokens: 999, Other: `{"frt":1,"stream_status":{"status":"error"}}`},
|
|
}
|
|
require.NoError(t, db.Create(&logs).Error)
|
|
fallback, err := GetGroupModelPerformance(ctx, "primary", "m")
|
|
require.NoError(t, err)
|
|
require.NotNil(t, fallback.Throughput)
|
|
assert.Equal(t, ModelPerformanceMetric{Value: 30, Samples: 2, Source: "test", UpdatedAt: 2}, *fallback.Throughput)
|
|
assert.Equal(t, 1000.0, fallback.FirstToken.Value)
|
|
calls := make([]Log, 0, 31)
|
|
for i := range 31 {
|
|
calls = append(calls, Log{ChannelId: 1, Type: LogTypeConsume, TokenId: 8, Group: "primary", ModelName: "m", IsStream: true, CreatedAt: int64(100 + i), UseTime: i + 1, CompletionTokens: (i + 1) * (i + 1), Other: fmt.Sprintf(`{"frt":%d}`, (i+1)*100)})
|
|
}
|
|
require.NoError(t, db.Create(&calls).Error)
|
|
// Invalid, another group's, and synthetic requests must not displace real samples.
|
|
noise := []Log{
|
|
{Type: LogTypeConsume, TokenId: 8, Group: "backup", ModelName: "m", IsStream: true, CreatedAt: 999, UseTime: 1, CompletionTokens: 999, Other: `{"frt":1}`},
|
|
{Type: LogTypeConsume, TokenId: 8, Group: "primary", ModelName: "other", IsStream: true, CreatedAt: 999, UseTime: 1, CompletionTokens: 999, Other: `{"frt":1}`},
|
|
{Type: LogTypeConsume, TokenId: 8, Group: "primary", ModelName: "m", IsStream: true, CreatedAt: 999, UseTime: 1, CompletionTokens: 999, Other: `{"frt":1,"stream_status":{"status":"error"}}`},
|
|
{Type: LogTypeConsume, TokenId: 8, Group: "primary", ModelName: "m", IsStream: true, CreatedAt: 999, UseTime: 1, CompletionTokens: 999, Other: `{"frt":1,"stream_status":{"status":"ok","end_reason":"client_gone"}}`},
|
|
}
|
|
require.NoError(t, db.Create(&noise).Error)
|
|
actual, err := GetGroupModelPerformance(ctx, "primary", "m")
|
|
require.NoError(t, err)
|
|
assert.Equal(t, ModelPerformanceMetric{Value: 16.5, Samples: 30, Source: "calls", UpdatedAt: 130}, *actual.Throughput)
|
|
assert.Equal(t, ModelPerformanceMetric{Value: 1650, Samples: 30, Source: "calls", UpdatedAt: 130}, *actual.FirstToken)
|
|
// Synchronous requests contribute throughput, not fabricated first-token timing.
|
|
require.NoError(t, db.Create(&Log{Type: LogTypeConsume, TokenId: 8, Group: "primary", ModelName: "m", CreatedAt: 1000, CompletionTokens: 10, Other: `{"latency_ms":500,"frt":500}`}).Error)
|
|
actual, err = GetGroupModelPerformance(ctx, "primary", "m")
|
|
require.NoError(t, err)
|
|
assert.InDelta(t, (16.5*30-2+20)/30, actual.Throughput.Value, 0.0001)
|
|
assert.Equal(t, 1650.0, actual.FirstToken.Value)
|
|
require.NoError(t, db.Model(&Log{}).Where(map[string]any{"group": "primary", "model_name": "m"}).Where("token_id > 0").Update("is_stream", false).Error)
|
|
partial, err := GetGroupModelPerformance(ctx, "primary", "m")
|
|
require.NoError(t, err)
|
|
assert.Equal(t, "calls", partial.Throughput.Source)
|
|
assert.Equal(t, "test", partial.FirstToken.Source)
|
|
assert.Equal(t, 1000.0, partial.FirstToken.Value)
|
|
})
|
|
}
|
|
}
|