feat: 算力引擎 Go 后端核心（new-api）
- 统一 OpenAI 兼容 /v1 中继 + 渠道/额度/令牌/流水 - 各领域包：relay 模型网关、model 数据层、controller 管理 API
This commit is contained in:
@@ -0,0 +1,98 @@
|
||||
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,
|
||||
OutputTokens: drained.outputTokens,
|
||||
GenerationMs: drained.generationMs,
|
||||
})
|
||||
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"]),
|
||||
outputTokens: parseRedisInt(values["out"]),
|
||||
generationMs: parseRedisInt(values["gen_ms"]),
|
||||
}
|
||||
}
|
||||
|
||||
func parseRedisInt(value string) int64 {
|
||||
if value == "" {
|
||||
return 0
|
||||
}
|
||||
parsed, _ := strconv.ParseInt(value, 10, 64)
|
||||
return parsed
|
||||
}
|
||||
@@ -0,0 +1,428 @@
|
||||
package perfmetrics
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"math"
|
||||
"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
|
||||
|
||||
// seriesSchema is a stable client cache/schema marker. Do not change it when
|
||||
// hiding fields or making response-only privacy hardening changes.
|
||||
const seriesSchema = "dbcd0a3c01b55203"
|
||||
|
||||
func Init() {
|
||||
go flushLoop()
|
||||
}
|
||||
|
||||
func RecordRelaySample(info *relaycommon.RelayInfo, success bool, outputTokens int64) {
|
||||
if info == nil {
|
||||
return
|
||||
}
|
||||
now := time.Now()
|
||||
hasTtft := info.IsStream && info.HasSendResponse()
|
||||
ttftMs := int64(0)
|
||||
if hasTtft {
|
||||
ttftMs = info.FirstResponseTime.Sub(info.StartTime).Milliseconds()
|
||||
}
|
||||
latencyMs := now.Sub(info.StartTime).Milliseconds()
|
||||
generationMs := latencyMs
|
||||
if hasTtft {
|
||||
generationMs = now.Sub(info.FirstResponseTime).Milliseconds()
|
||||
}
|
||||
if generationMs <= 0 {
|
||||
generationMs = latencyMs
|
||||
}
|
||||
Record(Sample{
|
||||
Model: info.OriginModelName,
|
||||
Group: info.UsingGroup,
|
||||
LatencyMs: latencyMs,
|
||||
TtftMs: ttftMs,
|
||||
HasTtft: hasTtft,
|
||||
Success: success,
|
||||
OutputTokens: outputTokens,
|
||||
GenerationMs: generationMs,
|
||||
})
|
||||
}
|
||||
|
||||
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,
|
||||
outputTokens: row.OutputTokens,
|
||||
generationMs: row.GenerationMs,
|
||||
})
|
||||
}
|
||||
|
||||
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 QuerySummaryAll(hours int, groups []string) (SummaryAllResult, error) {
|
||||
if hours <= 0 {
|
||||
hours = 24
|
||||
}
|
||||
if hours > 24*30 {
|
||||
hours = 24 * 30
|
||||
}
|
||||
endTs := time.Now().Unix()
|
||||
startTs := endTs - int64(hours)*3600
|
||||
allowedGroups := allowedGroupSet(groups)
|
||||
|
||||
rows, err := model.GetPerfMetricsSummaryBucketsAll(startTs, endTs, groups)
|
||||
if err != nil {
|
||||
return SummaryAllResult{}, err
|
||||
}
|
||||
|
||||
totals := map[string]counters{}
|
||||
modelBuckets := map[string]map[int64]counters{}
|
||||
for _, row := range rows {
|
||||
value := counters{
|
||||
requestCount: row.RequestCount,
|
||||
successCount: row.SuccessCount,
|
||||
totalLatencyMs: row.TotalLatencyMs,
|
||||
outputTokens: row.OutputTokens,
|
||||
generationMs: row.GenerationMs,
|
||||
}
|
||||
mergeModelTotals(totals, row.ModelName, value)
|
||||
mergeModelBucket(modelBuckets, row.ModelName, row.BucketTs, value)
|
||||
}
|
||||
|
||||
hotBuckets.Range(func(key, value any) bool {
|
||||
k := key.(bucketKey)
|
||||
if k.bucketTs < startTs || k.bucketTs > endTs {
|
||||
return true
|
||||
}
|
||||
if allowedGroups != nil {
|
||||
if _, ok := allowedGroups[k.group]; !ok {
|
||||
return true
|
||||
}
|
||||
}
|
||||
snap := value.(*atomicBucket).snapshot()
|
||||
if snap.requestCount == 0 {
|
||||
return true
|
||||
}
|
||||
mergeModelTotals(totals, k.model, snap)
|
||||
mergeModelBucket(modelBuckets, k.model, k.bucketTs, snap)
|
||||
return true
|
||||
})
|
||||
|
||||
models := make([]ModelSummary, 0, len(totals))
|
||||
for name, total := range totals {
|
||||
if total.requestCount == 0 {
|
||||
continue
|
||||
}
|
||||
avgLatency := total.totalLatencyMs / total.requestCount
|
||||
successRate := float64(total.successCount) / float64(total.requestCount) * 100
|
||||
avgTps := 0.0
|
||||
if total.generationMs > 0 {
|
||||
avgTps = float64(total.outputTokens) / (float64(total.generationMs) / 1000.0)
|
||||
}
|
||||
models = append(models, ModelSummary{
|
||||
ModelName: name,
|
||||
AvgLatencyMs: avgLatency,
|
||||
SuccessRate: math.Round(successRate*100) / 100,
|
||||
AvgTps: math.Round(avgTps*100) / 100,
|
||||
RecentSuccessRates: recentSuccessRates(modelBuckets[name], 3),
|
||||
RequestCount: total.requestCount,
|
||||
})
|
||||
}
|
||||
sort.Slice(models, func(i, j int) bool {
|
||||
return models[i].RequestCount > models[j].RequestCount
|
||||
})
|
||||
|
||||
return SummaryAllResult{Models: models}, nil
|
||||
}
|
||||
|
||||
func mergeModelTotals(totals map[string]counters, modelName string, value counters) {
|
||||
if value.requestCount == 0 {
|
||||
return
|
||||
}
|
||||
current := totals[modelName]
|
||||
current.requestCount += value.requestCount
|
||||
current.successCount += value.successCount
|
||||
current.totalLatencyMs += value.totalLatencyMs
|
||||
current.ttftSumMs += value.ttftSumMs
|
||||
current.ttftCount += value.ttftCount
|
||||
current.outputTokens += value.outputTokens
|
||||
current.generationMs += value.generationMs
|
||||
totals[modelName] = current
|
||||
}
|
||||
|
||||
func mergeModelBucket(modelBuckets map[string]map[int64]counters, modelName string, bucketTs int64, value counters) {
|
||||
if value.requestCount == 0 {
|
||||
return
|
||||
}
|
||||
if _, ok := modelBuckets[modelName]; !ok {
|
||||
modelBuckets[modelName] = map[int64]counters{}
|
||||
}
|
||||
current := modelBuckets[modelName][bucketTs]
|
||||
current.requestCount += value.requestCount
|
||||
current.successCount += value.successCount
|
||||
current.totalLatencyMs += value.totalLatencyMs
|
||||
current.ttftSumMs += value.ttftSumMs
|
||||
current.ttftCount += value.ttftCount
|
||||
current.outputTokens += value.outputTokens
|
||||
current.generationMs += value.generationMs
|
||||
modelBuckets[modelName][bucketTs] = current
|
||||
}
|
||||
|
||||
func recentSuccessRates(buckets map[int64]counters, limit int) []float64 {
|
||||
if len(buckets) == 0 || limit <= 0 {
|
||||
return nil
|
||||
}
|
||||
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]
|
||||
})
|
||||
if len(timestamps) > limit {
|
||||
timestamps = timestamps[len(timestamps)-limit:]
|
||||
}
|
||||
rates := make([]float64, 0, len(timestamps))
|
||||
for _, ts := range timestamps {
|
||||
rates = append(rates, math.Round(successRate(buckets[ts])*100)/100)
|
||||
}
|
||||
return rates
|
||||
}
|
||||
|
||||
func allowedGroupSet(groups []string) map[string]struct{} {
|
||||
if groups == nil {
|
||||
return nil
|
||||
}
|
||||
allowed := make(map[string]struct{}, len(groups))
|
||||
for _, group := range groups {
|
||||
allowed[group] = struct{}{}
|
||||
}
|
||||
return allowed
|
||||
}
|
||||
|
||||
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
|
||||
current.outputTokens += value.outputTokens
|
||||
current.generationMs += value.generationMs
|
||||
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
|
||||
total.outputTokens += value.outputTokens
|
||||
total.generationMs += value.generationMs
|
||||
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),
|
||||
AvgTps: avgTps(total),
|
||||
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),
|
||||
AvgTps: avgTps(value),
|
||||
}
|
||||
}
|
||||
|
||||
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 avgTps(value counters) float64 {
|
||||
if value.outputTokens <= 0 || value.generationMs <= 0 {
|
||||
return 0
|
||||
}
|
||||
return float64(value.outputTokens) / (float64(value.generationMs) / 1000)
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
if sample.OutputTokens > 0 && sample.GenerationMs > 0 {
|
||||
pipe.HIncrBy(ctx, redisKey, "out", sample.OutputTokens)
|
||||
pipe.HIncrBy(ctx, redisKey, "gen_ms", sample.GenerationMs)
|
||||
}
|
||||
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)
|
||||
}
|
||||
@@ -0,0 +1,153 @@
|
||||
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
|
||||
OutputTokens int64
|
||||
GenerationMs int64
|
||||
}
|
||||
|
||||
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"`
|
||||
AvgTps float64 `json:"avg_tps"`
|
||||
}
|
||||
|
||||
type GroupResult struct {
|
||||
Group string `json:"group"`
|
||||
AvgTtftMs int64 `json:"avg_ttft_ms"`
|
||||
AvgLatencyMs int64 `json:"avg_latency_ms"`
|
||||
SuccessRate float64 `json:"success_rate"`
|
||||
AvgTps float64 `json:"avg_tps"`
|
||||
Series []BucketPoint `json:"series"`
|
||||
}
|
||||
|
||||
type QueryResult struct {
|
||||
ModelName string `json:"model_name"`
|
||||
SeriesSchema string `json:"series_schema"`
|
||||
Groups []GroupResult `json:"groups"`
|
||||
}
|
||||
|
||||
type ModelSummary struct {
|
||||
ModelName string `json:"model_name"`
|
||||
AvgLatencyMs int64 `json:"avg_latency_ms"`
|
||||
SuccessRate float64 `json:"success_rate"`
|
||||
AvgTps float64 `json:"avg_tps"`
|
||||
RecentSuccessRates []float64 `json:"recent_success_rates,omitempty"`
|
||||
RequestCount int64 `json:"-"`
|
||||
}
|
||||
|
||||
type SummaryAllResult struct {
|
||||
Models []ModelSummary `json:"models"`
|
||||
}
|
||||
|
||||
type bucketKey struct {
|
||||
model string
|
||||
group string
|
||||
bucketTs int64
|
||||
}
|
||||
|
||||
type counters struct {
|
||||
requestCount int64
|
||||
successCount int64
|
||||
totalLatencyMs int64
|
||||
ttftSumMs int64
|
||||
ttftCount int64
|
||||
outputTokens int64
|
||||
generationMs int64
|
||||
}
|
||||
|
||||
type atomicBucket struct {
|
||||
requestCount atomic.Int64
|
||||
successCount atomic.Int64
|
||||
totalLatencyMs atomic.Int64
|
||||
ttftSumMs atomic.Int64
|
||||
ttftCount atomic.Int64
|
||||
outputTokens atomic.Int64
|
||||
generationMs 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)
|
||||
}
|
||||
if sample.OutputTokens > 0 && sample.GenerationMs > 0 {
|
||||
b.outputTokens.Add(sample.OutputTokens)
|
||||
b.generationMs.Add(sample.GenerationMs)
|
||||
}
|
||||
}
|
||||
|
||||
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(),
|
||||
outputTokens: b.outputTokens.Load(),
|
||||
generationMs: b.generationMs.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),
|
||||
outputTokens: b.outputTokens.Swap(0),
|
||||
generationMs: b.generationMs.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)
|
||||
}
|
||||
if c.outputTokens != 0 {
|
||||
b.outputTokens.Add(c.outputTokens)
|
||||
}
|
||||
if c.generationMs != 0 {
|
||||
b.generationMs.Add(c.generationMs)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user