BCAI/easyai-ai-gateway
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: add runner failover rules and cache affinity

Browse Source
This commit is contained in:
wangbo
2026-06-28 20:50:23 +08:00
parent 41bb3a6525
commit 0e675e259c
31 changed files with 2939 additions and 470 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/internal/runner/service.go
+25 -8
View File
@@ -145,7 +145,15 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
return Result{}, err
}
}
candidates, err := s.store.ListModelCandidates(ctx, task.Model, modelType, user)
runnerPolicy, err := s.store.GetActiveRunnerPolicy(ctx)
if err != nil {
return Result{}, err
}
cacheAffinityKey := buildCacheAffinityKey(task.Kind, modelType, body)
candidates, err := s.store.ListModelCandidates(ctx, task.Model, modelType, user, store.ListModelCandidatesOptions{
CacheAffinityKey: cacheAffinityKey,
CacheAffinityPolicy: runnerPolicy.CacheAffinityPolicy,
})
if err != nil {
s.recordFailedAttempt(ctx, failedAttemptRecord{
Task: task,
@@ -278,10 +286,6 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
return Result{}, err
}
runnerPolicy, err := s.store.GetActiveRunnerPolicy(ctx)
if err != nil {
return Result{}, err
}
maxPlatforms := maxPlatformsForCandidates(candidates, runnerPolicy)
maxFailoverDuration := maxFailoverDurationForCandidates(candidates, runnerPolicy)
singleSourceProtected := singleSourceProtectionActive(candidates, maxPlatforms, runnerPolicy)
@@ -320,7 +324,7 @@ candidatesLoop:
break candidatesLoop
}
candidateBody := preprocessing.Body
response, err := s.runCandidate(ctx, task, user, candidateBody, preprocessing.Log, candidate, nextAttemptNo, onDelta, singleSourceProtected)
response, err := s.runCandidate(ctx, task, user, candidateBody, preprocessing.Log, candidate, nextAttemptNo, onDelta, singleSourceProtected, runnerPolicy.CacheAffinityPolicy)
if err == nil {
attemptNo = nextAttemptNo
billings := s.billings(ctx, user, task.Kind, candidateBody, candidate, response, isSimulation(task, candidate))
@@ -490,7 +494,7 @@ candidatesLoop:
if !decision.Retry {
break
}
s.applyFailoverAction(ctx, task.ID, candidate, decision, isSimulation(task, candidate), singleSourceProtected)
s.applyFailoverAction(ctx, task.ID, candidate, task.Model, 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,
@@ -531,7 +535,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, singleSourceProtected bool) (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, cacheAffinityPolicy map[string]any) (clients.Response, error) {
simulated := isSimulation(task, candidate)
baseAttemptMetrics := mergeMetrics(attemptMetrics(candidate, attemptNo, simulated), parameterPreprocessingMetrics(preprocessing))
reservations := s.rateLimitReservations(ctx, user, candidate, body)
@@ -763,6 +767,19 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
}); err != nil {
return clients.Response{}, fmt.Errorf("finish task attempt: %w", err)
}
if err := s.store.RecordCacheAffinityObservation(context.WithoutCancel(ctx), store.CacheAffinityObservationInput{
CacheAffinityKey: candidate.CacheAffinity.Key,
CacheAffinityPolicy: cacheAffinityPolicy,
PlatformID: candidate.PlatformID,
PlatformModelID: candidate.PlatformModelID,
ClientID: candidate.ClientID,
ModelType: candidate.ModelType,
InputTokens: response.Usage.InputTokens,
CachedInputTokens: response.Usage.CachedInputTokens,
CachedInputTokensKnown: response.Usage.CachedInputTokensKnown,
}); err != nil {
s.logger.Warn("record cache affinity observation failed", "error", err, "clientId", candidate.ClientID)
}
return response, nil
}