增加单一源保护全局策略

This commit is contained in:
2026-06-15 00:14:40 +08:00
parent bffd4ecb98
commit b860ef37e8
12 changed files with 210 additions and 14 deletions
@@ -37,6 +37,40 @@ func TestFailoverBudgetDefaults(t *testing.T) {
}
}
func TestSingleSourceProtectionDefaultsOn(t *testing.T) {
runnerPolicy := store.RunnerPolicy{Status: "active"}
if !singleSourceProtectionActive([]store.RuntimeModelCandidate{{PlatformID: "only-source"}}, 99, runnerPolicy) {
t.Fatal("single source protection should default to enabled")
}
}
func TestSingleSourceProtectionCanBeDisabled(t *testing.T) {
runnerPolicy := store.RunnerPolicy{
Status: "active",
SingleSourcePolicy: map[string]any{"enabled": false},
}
if singleSourceProtectionActive([]store.RuntimeModelCandidate{{PlatformID: "only-source"}}, 99, runnerPolicy) {
t.Fatal("single source protection should respect the global toggle")
}
}
func TestSingleSourceProtectionUsesEffectivePlatformBudget(t *testing.T) {
candidates := []store.RuntimeModelCandidate{
{PlatformID: "first-source"},
{PlatformID: "second-source"},
}
runnerPolicy := store.RunnerPolicy{Status: "active"}
if !singleSourceProtectionActive(candidates, 1, runnerPolicy) {
t.Fatal("maxPlatforms=1 should make this execution a single-source run")
}
if singleSourceProtectionActive(candidates, 2, runnerPolicy) {
t.Fatal("multiple effective sources should allow disable and cooldown actions")
}
}
func TestFailoverTimeBudgetExceeded(t *testing.T) {
if !failoverTimeBudgetExceeded(time.Now().Add(-601*time.Second), 10*time.Minute) {
t.Fatal("failover time budget should stop retries after the configured duration")
+25 -4
View File
@@ -9,7 +9,7 @@ import (
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func (s *Service) applyCandidateFailurePolicies(ctx context.Context, taskID string, candidate store.RuntimeModelCandidate, cause error, simulated bool) {
func (s *Service) applyCandidateFailurePolicies(ctx context.Context, taskID string, candidate store.RuntimeModelCandidate, cause error, simulated bool, singleSourceProtected bool) {
code := clients.ErrorCode(cause)
message := ""
if cause != nil {
@@ -18,7 +18,9 @@ func (s *Service) applyCandidateFailurePolicies(ctx context.Context, taskID stri
autoDisablePolicy := effectiveRuntimePolicy(candidate.AutoDisablePolicy, candidate.RuntimePolicyOverride, "autoDisablePolicy")
if failurePolicyMatches(autoDisablePolicy, code, message) && intFromPolicy(autoDisablePolicy, "threshold") <= 1 {
if err := s.store.DisableCandidatePlatform(ctx, candidate.PlatformID); err == nil {
if singleSourceProtected {
s.emitSingleSourceProtected(ctx, taskID, candidate, "auto_disable", code, simulated)
} else if err := s.store.DisableCandidatePlatform(ctx, candidate.PlatformID); err == nil {
_ = s.emit(ctx, taskID, "task.policy.auto_disabled", "running", "auto_disable", 0.48, "candidate platform disabled by failure policy", map[string]any{
"platformId": candidate.PlatformID,
"platformModelId": candidate.PlatformModelID,
@@ -30,7 +32,9 @@ func (s *Service) applyCandidateFailurePolicies(ctx context.Context, taskID stri
degradePolicy := effectiveRuntimePolicy(candidate.DegradePolicy, candidate.RuntimePolicyOverride, "degradePolicy")
if failurePolicyMatches(degradePolicy, code, message) {
cooldownSeconds := intFromPolicy(degradePolicy, "cooldownSeconds")
if err := s.store.CooldownCandidatePlatformModel(ctx, candidate.PlatformModelID, cooldownSeconds); err == nil {
if singleSourceProtected {
s.emitSingleSourceProtected(ctx, taskID, candidate, "degrade_cooldown", code, simulated)
} else if err := s.store.CooldownCandidatePlatformModel(ctx, candidate.PlatformModelID, cooldownSeconds); err == nil {
_ = s.emit(ctx, taskID, "task.policy.degraded", "running", "degrade", 0.5, "candidate model cooled down by failure policy", map[string]any{
"platformId": candidate.PlatformID,
"platformModelId": candidate.PlatformModelID,
@@ -41,9 +45,13 @@ func (s *Service) applyCandidateFailurePolicies(ctx context.Context, taskID stri
}
}
func (s *Service) applyFailoverAction(ctx context.Context, taskID string, candidate store.RuntimeModelCandidate, decision failoverDecision, simulated bool) {
func (s *Service) applyFailoverAction(ctx context.Context, taskID string, candidate store.RuntimeModelCandidate, decision failoverDecision, simulated bool, singleSourceProtected bool) {
switch decision.Action {
case "disable_and_next":
if singleSourceProtected {
s.emitSingleSourceProtected(ctx, taskID, candidate, decision.Action, decision.Info.Code, simulated)
return
}
if err := s.store.DisableCandidatePlatform(ctx, candidate.PlatformID); err == nil {
_ = s.emit(ctx, taskID, "task.policy.failover_disabled", "running", "failover", 0.51, "candidate platform disabled by runner failover policy", addPolicyTracePayload(map[string]any{
"platformId": candidate.PlatformID,
@@ -53,6 +61,10 @@ func (s *Service) applyFailoverAction(ctx context.Context, taskID string, candid
}, decision.Match, decision.Info), simulated)
}
case "cooldown_and_next":
if singleSourceProtected {
s.emitSingleSourceProtected(ctx, taskID, candidate, decision.Action, decision.Info.Code, simulated)
return
}
if err := s.store.CooldownCandidatePlatformModel(ctx, candidate.PlatformModelID, decision.CooldownSeconds); err == nil {
_ = s.emit(ctx, taskID, "task.policy.failover_cooled_down", "running", "failover", 0.51, "candidate model cooled down by runner failover policy", addPolicyTracePayload(map[string]any{
"platformId": candidate.PlatformID,
@@ -65,6 +77,15 @@ func (s *Service) applyFailoverAction(ctx context.Context, taskID string, candid
}
}
func (s *Service) emitSingleSourceProtected(ctx context.Context, taskID string, candidate store.RuntimeModelCandidate, action string, code string, simulated bool) {
_ = s.emit(ctx, taskID, "task.policy.single_source_protected", "running", "policy_guard", 0.5, "single source protected from disable or cooldown", map[string]any{
"platformId": candidate.PlatformID,
"platformModelId": candidate.PlatformModelID,
"action": action,
"code": code,
}, simulated)
}
func (s *Service) applyPriorityDemotePolicy(ctx context.Context, taskID string, attemptNo int, runnerPolicy store.RunnerPolicy, candidate store.RuntimeModelCandidate, requestedModel string, cause error, simulated bool) {
if errors.Is(cause, store.ErrRateLimited) {
return
+41 -4
View File
@@ -240,6 +240,7 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
}
maxPlatforms := maxPlatformsForCandidates(candidates, runnerPolicy)
maxFailoverDuration := maxFailoverDurationForCandidates(candidates, runnerPolicy)
singleSourceProtected := singleSourceProtectionActive(candidates, maxPlatforms, runnerPolicy)
var lastErr error
var lastCandidate store.RuntimeModelCandidate
var lastPreprocessing *parameterPreprocessingLog
@@ -275,7 +276,7 @@ candidatesLoop:
break candidatesLoop
}
candidateBody := preprocessing.Body
response, err := s.runCandidate(ctx, task, user, candidateBody, preprocessing.Log, candidate, nextAttemptNo, onDelta)
response, err := s.runCandidate(ctx, task, user, candidateBody, preprocessing.Log, candidate, nextAttemptNo, onDelta, singleSourceProtected)
if err == nil {
attemptNo = nextAttemptNo
billings := s.billings(ctx, user, task.Kind, candidateBody, candidate, response, isSimulation(task, candidate))
@@ -445,7 +446,7 @@ candidatesLoop:
if !decision.Retry {
break
}
s.applyFailoverAction(ctx, task.ID, candidate, decision, isSimulation(task, candidate))
s.applyFailoverAction(ctx, task.ID, candidate, decision, isSimulation(task, candidate), singleSourceProtected)
if err := s.emit(ctx, task.ID, "task.retrying", "running", "retry", 0.55, "retrying next client", addPolicyTracePayload(map[string]any{
"attempt": attemptNo,
"action": decision.Action,
@@ -486,7 +487,7 @@ candidatesLoop:
return Result{Task: failed, Output: failed.Result}, lastErr
}
func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user *auth.User, body map[string]any, preprocessing parameterPreprocessingLog, candidate store.RuntimeModelCandidate, attemptNo int, onDelta clients.StreamDelta) (clients.Response, error) {
func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user *auth.User, body map[string]any, preprocessing parameterPreprocessingLog, candidate store.RuntimeModelCandidate, attemptNo int, onDelta clients.StreamDelta, singleSourceProtected bool) (clients.Response, error) {
simulated := isSimulation(task, candidate)
baseAttemptMetrics := mergeMetrics(attemptMetrics(candidate, attemptNo, simulated), parameterPreprocessingMetrics(preprocessing))
reservations := s.rateLimitReservations(ctx, user, candidate, body)
@@ -638,7 +639,7 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
ErrorMessage: err.Error(),
})
_ = s.emit(ctx, task.ID, "task.attempt.failed", "running", "attempt_failed", 0.45, err.Error(), map[string]any{"attempt": attemptNo, "retryable": retryable, "requestId": requestID, "statusCode": clients.ErrorResponseMetadata(err).StatusCode, "metrics": metrics}, simulated)
s.applyCandidateFailurePolicies(ctx, task.ID, candidate, err, simulated)
s.applyCandidateFailurePolicies(ctx, task.ID, candidate, err, simulated, singleSourceProtected)
return clients.Response{}, err
}
uploadedResult, err := s.uploadGeneratedAssets(ctx, task.ID, task.Kind, response.Result)
@@ -1099,6 +1100,42 @@ func maxPlatformsForCandidates(candidates []store.RuntimeModelCandidate, runnerP
return maxPlatforms
}
func singleSourceProtectionActive(candidates []store.RuntimeModelCandidate, maxPlatforms int, runnerPolicy store.RunnerPolicy) bool {
if !runnerSingleSourceProtectionEnabled(runnerPolicy) {
return false
}
return runtimeCandidateSourceCount(candidates, maxPlatforms) <= 1
}
func runnerSingleSourceProtectionEnabled(runnerPolicy store.RunnerPolicy) bool {
if enabled, ok := runnerPolicy.SingleSourcePolicy["enabled"].(bool); ok {
return enabled
}
return true
}
func runtimeCandidateSourceCount(candidates []store.RuntimeModelCandidate, maxPlatforms int) int {
limit := len(candidates)
if maxPlatforms > 0 && maxPlatforms < limit {
limit = maxPlatforms
}
sources := make(map[string]bool, limit)
for index := 0; index < limit; index++ {
key := strings.TrimSpace(candidates[index].PlatformID)
if key == "" {
key = strings.TrimSpace(candidates[index].PlatformModelID)
}
if key == "" {
key = strings.TrimSpace(candidates[index].ClientID)
}
if key == "" {
key = strconv.Itoa(index)
}
sources[key] = true
}
return len(sources)
}
func maxFailoverDurationForCandidates(candidates []store.RuntimeModelCandidate, runnerPolicy store.RunnerPolicy) time.Duration {
seconds := intFromPolicy(runnerPolicy.FailoverPolicy, "maxDurationSeconds")
if seconds <= 0 {