wangbo
9e4fc7362d
使用 PostgreSQL 统一同步与异步非文本任务的并发准入、持久化等待和 Worker 容量分配，并将生产 API 与独立 Worker 角色拆分。 补充策略管理、共享契约、OpenAPI、Kubernetes 双节点 Worker 清单及跨节点验收脚本；未默认启用任何生产 queue_size 策略。 已在原基线完成 Go、前端、迁移、Shell、Kustomize 与长任务容量验收；合入最新主干后将重新执行发布门禁。
287 lines
8.5 KiB
Go
287 lines
8.5 KiB
Go
package runner
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
|
)
|
|
|
|
type localRateLimitError struct {
|
|
clientErr *clients.ClientError
|
|
cause error
|
|
retryAfter time.Duration
|
|
}
|
|
|
|
func (e *localRateLimitError) Error() string {
|
|
if e == nil || e.clientErr == nil {
|
|
return store.ErrRateLimited.Error()
|
|
}
|
|
return e.clientErr.Error()
|
|
}
|
|
|
|
func (e *localRateLimitError) Unwrap() []error {
|
|
if e == nil || e.clientErr == nil {
|
|
if e != nil && e.cause != nil {
|
|
return []error{e.cause}
|
|
}
|
|
return []error{store.ErrRateLimited}
|
|
}
|
|
if e.cause != nil {
|
|
return []error{e.clientErr, e.cause}
|
|
}
|
|
return []error{e.clientErr, store.ErrRateLimited}
|
|
}
|
|
|
|
func localRateLimitRetryAfter(err error) time.Duration {
|
|
var limitErr *localRateLimitError
|
|
if errors.As(err, &limitErr) && limitErr.retryAfter > 0 {
|
|
return limitErr.retryAfter
|
|
}
|
|
return store.RateLimitRetryAfter(err)
|
|
}
|
|
|
|
func isLocalRateLimitError(err error) bool {
|
|
var limitErr *localRateLimitError
|
|
return errors.As(err, &limitErr)
|
|
}
|
|
|
|
func (s *Service) rateLimitReservations(ctx context.Context, user *auth.User, candidate store.RuntimeModelCandidate, body map[string]any) []store.RateLimitReservation {
|
|
out := make([]store.RateLimitReservation, 0)
|
|
out = append(out, reservationsFromPolicy(
|
|
"platform_model",
|
|
candidate.PlatformModelID,
|
|
firstNonEmptyString(candidate.DisplayName, candidate.ModelAlias, candidate.ModelName),
|
|
map[string]any{
|
|
"platformId": candidate.PlatformID,
|
|
"platformName": candidate.PlatformName,
|
|
"modelAlias": candidate.ModelAlias,
|
|
"modelName": candidate.ModelName,
|
|
},
|
|
effectiveRateLimitPolicy(candidate),
|
|
body,
|
|
)...)
|
|
if group, err := s.store.ResolveUserGroupPolicy(ctx, user); err == nil && group.ID != "" {
|
|
out = append(out, reservationsFromPolicy(
|
|
"user_group",
|
|
group.ID,
|
|
firstNonEmptyString(group.Name, group.GroupKey),
|
|
map[string]any{
|
|
"groupKey": group.GroupKey,
|
|
"name": group.Name,
|
|
},
|
|
store.NormalizeRateLimitPolicy(group.RateLimitPolicy),
|
|
body,
|
|
)...)
|
|
}
|
|
return out
|
|
}
|
|
|
|
func (s *Service) admissionScopes(ctx context.Context, user *auth.User, candidate store.RuntimeModelCandidate) ([]store.AdmissionScope, string) {
|
|
scopes := make([]store.AdmissionScope, 0, 2)
|
|
platformPolicy := effectiveRateLimitPolicy(candidate)
|
|
if scope, ok := admissionScopeFromPolicy(
|
|
"platform_model",
|
|
candidate.PlatformModelID,
|
|
firstNonEmptyString(candidate.DisplayName, candidate.ModelAlias, candidate.ModelName),
|
|
map[string]any{
|
|
"platformId": candidate.PlatformID,
|
|
"platformName": candidate.PlatformName,
|
|
"modelAlias": candidate.ModelAlias,
|
|
"modelName": candidate.ModelName,
|
|
},
|
|
platformPolicy,
|
|
); ok {
|
|
scopes = append(scopes, scope)
|
|
}
|
|
groupID := ""
|
|
if group, err := s.store.ResolveUserGroupPolicy(ctx, user); err == nil && group.ID != "" {
|
|
groupPolicy := store.NormalizeRateLimitPolicy(group.RateLimitPolicy)
|
|
if scope, ok := admissionScopeFromPolicy(
|
|
"user_group",
|
|
group.ID,
|
|
firstNonEmptyString(group.Name, group.GroupKey),
|
|
map[string]any{"groupKey": group.GroupKey, "name": group.Name},
|
|
groupPolicy,
|
|
); ok {
|
|
groupID = group.ID
|
|
scopes = append(scopes, scope)
|
|
}
|
|
}
|
|
return scopes, groupID
|
|
}
|
|
|
|
func admissionScopeFromPolicy(scopeType string, scopeKey string, scopeName string, metadata map[string]any, policy map[string]any) (store.AdmissionScope, bool) {
|
|
concurrentLimit, hasConcurrent := store.RateLimitPolicyMetric(policy, "concurrent")
|
|
queueRule, hasQueue := store.QueueRuleFromPolicy(policy)
|
|
if !hasConcurrent && !hasQueue {
|
|
return store.AdmissionScope{}, false
|
|
}
|
|
scope := store.AdmissionScope{
|
|
ScopeType: scopeType,
|
|
ScopeKey: scopeKey,
|
|
ScopeName: scopeName,
|
|
ScopeMetadata: metadata,
|
|
ConcurrentLimit: concurrentLimit,
|
|
Amount: 1,
|
|
LeaseTTLSeconds: 120,
|
|
Policy: policy,
|
|
}
|
|
if hasQueue {
|
|
scope.QueueLimit = queueRule.Limit
|
|
scope.MaxWaitSeconds = queueRule.MaxWaitSeconds
|
|
}
|
|
rules, _ := store.NormalizeRateLimitPolicy(policy)["rules"].([]any)
|
|
for _, rawRule := range rules {
|
|
rule, _ := rawRule.(map[string]any)
|
|
if strings.TrimSpace(stringFromMap(rule, "metric")) != "concurrent" {
|
|
continue
|
|
}
|
|
if ttl := int(floatFromAny(rule["leaseTtlSeconds"])); ttl > 0 {
|
|
scope.LeaseTTLSeconds = ttl
|
|
}
|
|
break
|
|
}
|
|
return scope, true
|
|
}
|
|
|
|
func effectiveRateLimitPolicy(candidate store.RuntimeModelCandidate) map[string]any {
|
|
return store.EffectiveRateLimitPolicy(store.EffectiveRateLimitPolicyInput{
|
|
BasePolicy: candidate.BaseRateLimitPolicy,
|
|
PlatformPolicy: candidate.PlatformRateLimitPolicy,
|
|
RuntimePolicy: candidate.RuntimeRateLimitPolicy,
|
|
RuntimePolicyExplicit: candidate.RuntimePolicyExplicit,
|
|
RuntimePolicyOverride: candidate.RateLimitRuntimeOverride,
|
|
ModelPolicy: candidate.ModelRateLimitPolicy,
|
|
ModelPolicyMode: candidate.ModelRateLimitPolicyMode,
|
|
})
|
|
}
|
|
|
|
func effectiveRetryPolicy(candidate store.RuntimeModelCandidate) map[string]any {
|
|
policy := candidate.PlatformRetryPolicy
|
|
if len(candidate.RuntimeRetryPolicy) > 0 {
|
|
policy = mergeMap(policy, candidate.RuntimeRetryPolicy)
|
|
}
|
|
if nested, ok := candidate.RuntimePolicyOverride["retryPolicy"].(map[string]any); ok && len(nested) > 0 {
|
|
policy = mergeMap(policy, nested)
|
|
}
|
|
if len(candidate.ModelRetryPolicy) > 0 {
|
|
policy = mergeMap(policy, candidate.ModelRetryPolicy)
|
|
}
|
|
return policy
|
|
}
|
|
|
|
func reservationsFromPolicy(scopeType string, scopeKey string, scopeName string, scopeMetadata map[string]any, policy map[string]any, body map[string]any) []store.RateLimitReservation {
|
|
if scopeKey == "" || !hasRules(policy) {
|
|
return nil
|
|
}
|
|
rules, _ := policy["rules"].([]any)
|
|
out := make([]store.RateLimitReservation, 0, len(rules))
|
|
estimatedTokens := estimateRequestTokens(body)
|
|
for _, rawRule := range rules {
|
|
rule, _ := rawRule.(map[string]any)
|
|
metric := strings.TrimSpace(stringFromMap(rule, "metric"))
|
|
limit := floatFromAny(rule["limit"])
|
|
if metric == "" || metric == "queue_size" || limit <= 0 {
|
|
continue
|
|
}
|
|
amount := 1.0
|
|
if strings.HasPrefix(metric, "tpm") {
|
|
amount = float64(estimatedTokens)
|
|
}
|
|
out = append(out, store.RateLimitReservation{
|
|
ScopeType: scopeType,
|
|
ScopeKey: scopeKey,
|
|
ScopeName: scopeName,
|
|
ScopeMetadata: scopeMetadata,
|
|
Metric: metric,
|
|
Limit: limit,
|
|
Amount: amount,
|
|
WindowSeconds: int(floatFromAny(rule["windowSeconds"])),
|
|
LeaseTTLSeconds: int(floatFromAny(rule["leaseTtlSeconds"])),
|
|
Policy: policy,
|
|
})
|
|
}
|
|
return out
|
|
}
|
|
|
|
func hasRules(policy map[string]any) bool {
|
|
rules, _ := policy["rules"].([]any)
|
|
return len(rules) > 0
|
|
}
|
|
|
|
func estimateRequestTokens(body map[string]any) int {
|
|
var text strings.Builder
|
|
appendTokenEstimateText(&text, body["prompt"])
|
|
appendTokenEstimateText(&text, body["input"])
|
|
appendTokenEstimateText(&text, body["query"])
|
|
appendTokenEstimateText(&text, body["documents"])
|
|
for _, item := range contentItems(body["content"]) {
|
|
if stringFromAny(item["type"]) == "text" {
|
|
appendTokenEstimateText(&text, item["text"])
|
|
}
|
|
}
|
|
if messages, ok := body["messages"].([]any); ok {
|
|
for _, raw := range messages {
|
|
message, _ := raw.(map[string]any)
|
|
switch content := message["content"].(type) {
|
|
case string:
|
|
appendTokenEstimateText(&text, content)
|
|
case []any:
|
|
for _, rawPart := range content {
|
|
part, _ := rawPart.(map[string]any)
|
|
appendTokenEstimateText(&text, part["text"])
|
|
}
|
|
}
|
|
}
|
|
}
|
|
estimatedText := text.String()
|
|
if estimatedText == "" {
|
|
return 1
|
|
}
|
|
return len([]rune(estimatedText))/4 + 1
|
|
}
|
|
|
|
func appendTokenEstimateText(out *strings.Builder, value any) {
|
|
switch typed := value.(type) {
|
|
case string:
|
|
out.WriteString(typed)
|
|
case []string:
|
|
for _, item := range typed {
|
|
out.WriteString(item)
|
|
}
|
|
case []any:
|
|
for _, item := range typed {
|
|
appendTokenEstimateText(out, item)
|
|
}
|
|
case map[string]any:
|
|
for _, key := range []string{"text", "content", "query", "document"} {
|
|
if text := stringFromAny(typed[key]); text != "" {
|
|
out.WriteString(text)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func tokenUsageAmounts(usage clients.Usage) map[string]float64 {
|
|
out := map[string]float64{}
|
|
if usage.InputTokens > 0 {
|
|
out["tpm_input"] = float64(usage.InputTokens)
|
|
}
|
|
if usage.OutputTokens > 0 {
|
|
out["tpm_output"] = float64(usage.OutputTokens)
|
|
}
|
|
total := usage.TotalTokens
|
|
if total <= 0 {
|
|
total = usage.InputTokens + usage.OutputTokens
|
|
}
|
|
if total > 0 {
|
|
out["tpm_total"] = float64(total)
|
|
}
|
|
return out
|
|
}
|