feat: collect model performance metrics (#4635)

This commit is contained in:
Calcium-Ion
2026-05-06 13:55:23 +08:00
committed by GitHub
parent 8b2b03d276
commit 9acf5fecae
26 changed files with 1078 additions and 120 deletions
+94
View File
@@ -0,0 +1,94 @@
package perfmetrics
import (
"fmt"
"strconv"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/setting/perf_metrics_setting"
)
func flushLoop() {
for {
interval := perf_metrics_setting.GetFlushIntervalMinutes()
time.Sleep(time.Duration(interval) * time.Minute)
setting := perf_metrics_setting.GetSetting()
if !setting.Enabled {
continue
}
flushCompletedBuckets()
cleanupExpiredMetrics(setting.RetentionDays)
}
}
func flushCompletedBuckets() {
currentBucket := bucketStart(time.Now().Unix())
hotBuckets.Range(func(key, value any) bool {
k := key.(bucketKey)
if k.bucketTs >= currentBucket {
return true
}
bucket := value.(*atomicBucket)
drained := bucket.drain()
if drained.requestCount == 0 {
deleteOldEmptyBucket(k, key)
return true
}
err := model.UpsertPerfMetric(&model.PerfMetric{
ModelName: k.model,
Group: k.group,
BucketTs: k.bucketTs,
RequestCount: drained.requestCount,
SuccessCount: drained.successCount,
TotalLatencyMs: drained.totalLatencyMs,
TtftSumMs: drained.ttftSumMs,
TtftCount: drained.ttftCount,
})
if err != nil {
bucket.addCounters(drained)
common.SysError(fmt.Sprintf("failed to flush perf metric bucket model=%s group=%s bucket=%d: %s", k.model, k.group, k.bucketTs, err.Error()))
return true
}
deleteOldEmptyBucket(k, key)
return true
})
}
func deleteOldEmptyBucket(k bucketKey, rawKey any) {
if k.bucketTs < bucketStart(time.Now().Add(-24*time.Hour).Unix()) {
hotBuckets.Delete(rawKey)
}
}
func cleanupExpiredMetrics(retentionDays int) {
if retentionDays <= 0 {
return
}
cutoff := time.Now().Add(-time.Duration(retentionDays) * 24 * time.Hour).Unix()
if err := model.DeletePerfMetricsBefore(cutoff); err != nil {
common.SysError("failed to cleanup expired perf metrics: " + err.Error())
}
}
func redisCounters(values map[string]string) counters {
return counters{
requestCount: parseRedisInt(values["req"]),
successCount: parseRedisInt(values["ok"]),
totalLatencyMs: parseRedisInt(values["lat"]),
ttftSumMs: parseRedisInt(values["ttft"]),
ttftCount: parseRedisInt(values["ttft_n"]),
}
}
func parseRedisInt(value string) int64 {
if value == "" {
return 0
}
parsed, _ := strconv.ParseInt(value, 10, 64)
return parsed
}
+261
View File
@@ -0,0 +1,261 @@
package perfmetrics
import (
"context"
"fmt"
"sort"
"sync"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/setting/perf_metrics_setting"
)
var hotBuckets sync.Map
const seriesSchema = "dbcd0a3c01b55203"
func Init() {
go flushLoop()
}
func RecordRelaySample(info *relaycommon.RelayInfo, success bool) {
if info == nil {
return
}
now := time.Now()
hasTtft := info.IsStream && info.HasSendResponse()
ttftMs := int64(0)
if hasTtft {
ttftMs = info.FirstResponseTime.Sub(info.StartTime).Milliseconds()
}
Record(Sample{
Model: info.OriginModelName,
Group: info.UsingGroup,
LatencyMs: now.Sub(info.StartTime).Milliseconds(),
TtftMs: ttftMs,
HasTtft: hasTtft,
Success: success,
})
}
func Record(sample Sample) {
setting := perf_metrics_setting.GetSetting()
if !setting.Enabled || sample.Model == "" {
return
}
if sample.Group == "" {
sample.Group = "default"
}
if sample.LatencyMs < 0 {
sample.LatencyMs = 0
}
key := bucketKey{
model: sample.Model,
group: sample.Group,
bucketTs: bucketStart(time.Now().Unix()),
}
actual, _ := hotBuckets.LoadOrStore(key, &atomicBucket{})
actual.(*atomicBucket).add(sample)
recordRedis(key, sample)
}
func Query(params QueryParams) (QueryResult, error) {
if params.Hours <= 0 {
params.Hours = 24
}
if params.Hours > 24*30 {
params.Hours = 24 * 30
}
endTs := time.Now().Unix()
startTs := endTs - int64(params.Hours)*3600
merged := map[bucketKey]counters{}
rows, err := model.GetPerfMetrics(params.Model, params.Group, startTs, endTs)
if err != nil {
return QueryResult{}, err
}
for _, row := range rows {
mergeCounters(merged, bucketKey{
model: row.ModelName,
group: row.Group,
bucketTs: row.BucketTs,
}, counters{
requestCount: row.RequestCount,
successCount: row.SuccessCount,
totalLatencyMs: row.TotalLatencyMs,
ttftSumMs: row.TtftSumMs,
ttftCount: row.TtftCount,
})
}
hotBuckets.Range(func(key, value any) bool {
k := key.(bucketKey)
if k.model != params.Model || k.bucketTs < startTs || k.bucketTs > endTs {
return true
}
if params.Group != "" && k.group != params.Group {
return true
}
mergeCounters(merged, k, value.(*atomicBucket).snapshot())
return true
})
return buildQueryResult(params.Model, merged), nil
}
func bucketStart(ts int64) int64 {
bucketSeconds := perf_metrics_setting.GetBucketSeconds()
if bucketSeconds <= 0 {
bucketSeconds = 3600
}
return ts - (ts % bucketSeconds)
}
func mergeCounters(merged map[bucketKey]counters, key bucketKey, value counters) {
if value.requestCount == 0 {
return
}
current := merged[key]
current.requestCount += value.requestCount
current.successCount += value.successCount
current.totalLatencyMs += value.totalLatencyMs
current.ttftSumMs += value.ttftSumMs
current.ttftCount += value.ttftCount
merged[key] = current
}
func buildQueryResult(modelName string, merged map[bucketKey]counters) QueryResult {
groupBuckets := map[string]map[int64]counters{}
for key, value := range merged {
if value.requestCount == 0 {
continue
}
if _, ok := groupBuckets[key.group]; !ok {
groupBuckets[key.group] = map[int64]counters{}
}
groupBuckets[key.group][key.bucketTs] = value
}
groups := make([]string, 0, len(groupBuckets))
for group := range groupBuckets {
groups = append(groups, group)
}
sort.Strings(groups)
results := make([]GroupResult, 0, len(groups))
for _, group := range groups {
buckets := groupBuckets[group]
timestamps := make([]int64, 0, len(buckets))
for ts := range buckets {
timestamps = append(timestamps, ts)
}
sort.Slice(timestamps, func(i, j int) bool {
return timestamps[i] < timestamps[j]
})
total := counters{}
series := make([]BucketPoint, 0, len(timestamps))
for _, ts := range timestamps {
value := buckets[ts]
total.requestCount += value.requestCount
total.successCount += value.successCount
total.totalLatencyMs += value.totalLatencyMs
total.ttftSumMs += value.ttftSumMs
total.ttftCount += value.ttftCount
series = append(series, bucketPoint(ts, value))
}
results = append(results, GroupResult{
Group: group,
AvgTtftMs: avg(total.ttftSumMs, total.ttftCount),
AvgLatencyMs: avg(total.totalLatencyMs, total.requestCount),
SuccessRate: successRate(total),
RequestCount: total.requestCount,
SuccessCount: total.successCount,
TtftCount: total.ttftCount,
Series: series,
})
}
return QueryResult{
ModelName: modelName,
SeriesSchema: seriesSchema,
Groups: results,
}
}
func bucketPoint(ts int64, value counters) BucketPoint {
return BucketPoint{
Ts: ts,
AvgTtftMs: avg(value.ttftSumMs, value.ttftCount),
AvgLatencyMs: avg(value.totalLatencyMs, value.requestCount),
SuccessRate: successRate(value),
Count: value.requestCount,
SuccessCount: value.successCount,
TtftCount: value.ttftCount,
}
}
func avg(sum int64, count int64) int64 {
if count <= 0 {
return 0
}
return sum / count
}
func successRate(value counters) float64 {
if value.requestCount <= 0 {
return 0
}
return float64(value.successCount) / float64(value.requestCount) * 100
}
func recordRedis(key bucketKey, sample Sample) {
if !common.RedisEnabled || common.RDB == nil {
return
}
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
redisKey := redisBucketKey(key)
pipe := common.RDB.TxPipeline()
pipe.HIncrBy(ctx, redisKey, "req", 1)
if sample.Success {
pipe.HIncrBy(ctx, redisKey, "ok", 1)
}
if sample.LatencyMs > 0 {
pipe.HIncrBy(ctx, redisKey, "lat", sample.LatencyMs)
}
if sample.HasTtft && sample.TtftMs >= 0 {
pipe.HIncrBy(ctx, redisKey, "ttft", sample.TtftMs)
pipe.HIncrBy(ctx, redisKey, "ttft_n", 1)
}
pipe.Expire(ctx, redisKey, time.Hour)
_, _ = pipe.Exec(ctx)
}
func mergeRedisActiveBuckets(merged map[bucketKey]counters, params QueryParams, startTs int64, endTs int64) {
if !common.RedisEnabled || common.RDB == nil || params.Model == "" || params.Group == "" {
return
}
active := bucketStart(time.Now().Unix())
if active < startTs || active > endTs {
return
}
key := bucketKey{model: params.Model, group: params.Group, bucketTs: active}
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
values, err := common.RDB.HGetAll(ctx, redisBucketKey(key)).Result()
if err != nil || len(values) == 0 {
return
}
mergeCounters(merged, key, redisCounters(values))
}
func redisBucketKey(key bucketKey) string {
return fmt.Sprintf("perf:%s:%s:%d", key.model, key.group, key.bucketTs)
}
+124
View File
@@ -0,0 +1,124 @@
package perfmetrics
import "sync/atomic"
type Store interface {
Record(sample Sample)
Query(params QueryParams) (QueryResult, error)
}
type Sample struct {
Model string
Group string
LatencyMs int64
TtftMs int64
HasTtft bool
Success bool
}
type QueryParams struct {
Model string
Group string
Hours int
}
type BucketPoint struct {
Ts int64 `json:"ts"`
AvgTtftMs int64 `json:"avg_ttft_ms"`
AvgLatencyMs int64 `json:"avg_latency_ms"`
SuccessRate float64 `json:"success_rate"`
Count int64 `json:"count"`
SuccessCount int64 `json:"success_count"`
TtftCount int64 `json:"ttft_count"`
}
type GroupResult struct {
Group string `json:"group"`
AvgTtftMs int64 `json:"avg_ttft_ms"`
AvgLatencyMs int64 `json:"avg_latency_ms"`
SuccessRate float64 `json:"success_rate"`
RequestCount int64 `json:"request_count"`
SuccessCount int64 `json:"success_count"`
TtftCount int64 `json:"ttft_count"`
Series []BucketPoint `json:"series"`
}
type QueryResult struct {
ModelName string `json:"model_name"`
SeriesSchema string `json:"series_schema"`
Groups []GroupResult `json:"groups"`
}
type bucketKey struct {
model string
group string
bucketTs int64
}
type counters struct {
requestCount int64
successCount int64
totalLatencyMs int64
ttftSumMs int64
ttftCount int64
}
type atomicBucket struct {
requestCount atomic.Int64
successCount atomic.Int64
totalLatencyMs atomic.Int64
ttftSumMs atomic.Int64
ttftCount atomic.Int64
}
func (b *atomicBucket) add(sample Sample) {
b.requestCount.Add(1)
if sample.Success {
b.successCount.Add(1)
}
if sample.LatencyMs > 0 {
b.totalLatencyMs.Add(sample.LatencyMs)
}
if sample.HasTtft && sample.TtftMs >= 0 {
b.ttftSumMs.Add(sample.TtftMs)
b.ttftCount.Add(1)
}
}
func (b *atomicBucket) snapshot() counters {
return counters{
requestCount: b.requestCount.Load(),
successCount: b.successCount.Load(),
totalLatencyMs: b.totalLatencyMs.Load(),
ttftSumMs: b.ttftSumMs.Load(),
ttftCount: b.ttftCount.Load(),
}
}
func (b *atomicBucket) drain() counters {
return counters{
requestCount: b.requestCount.Swap(0),
successCount: b.successCount.Swap(0),
totalLatencyMs: b.totalLatencyMs.Swap(0),
ttftSumMs: b.ttftSumMs.Swap(0),
ttftCount: b.ttftCount.Swap(0),
}
}
func (b *atomicBucket) addCounters(c counters) {
if c.requestCount != 0 {
b.requestCount.Add(c.requestCount)
}
if c.successCount != 0 {
b.successCount.Add(c.successCount)
}
if c.totalLatencyMs != 0 {
b.totalLatencyMs.Add(c.totalLatencyMs)
}
if c.ttftSumMs != 0 {
b.ttftSumMs.Add(c.ttftSumMs)
}
if c.ttftCount != 0 {
b.ttftCount.Add(c.ttftCount)
}
}