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

fix: treat local rate limit queueing separately

Browse Source
This commit is contained in:
wangbo
2026-05-12 10:34:06 +08:00
parent ddfd4f9035
commit 2a91b31d12
7 changed files with 107 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/internal/runner/service.go
+40 -20
View File
@@ -121,6 +121,7 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
maxFailoverDuration := maxFailoverDurationForCandidates(candidates, runnerPolicy)
attemptNo := task.AttemptCount
var lastErr error
candidatesLoop:
for index, candidate := range candidates {
if index >= maxPlatforms {
break
@@ -128,9 +129,10 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
clientAttempts := clientAttemptsForCandidate(candidate)
var candidateErr error
for clientAttempt := 1; clientAttempt <= clientAttempts; clientAttempt++ {
attemptNo++
response, err := s.runCandidate(ctx, task, user, body, candidate, attemptNo, onDelta)
nextAttemptNo := attemptNo + 1
response, err := s.runCandidate(ctx, task, user, body, candidate, nextAttemptNo, onDelta)
if err == nil {
attemptNo = nextAttemptNo
billings := s.billings(ctx, user, task.Kind, body, candidate, response, isSimulation(task, candidate))
record := buildSuccessRecord(task, user, body, candidate, response, billings, isSimulation(task, candidate))
record.Metrics = s.withAttemptHistory(ctx, task.ID, record.Metrics)
@@ -174,6 +176,19 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
}
return Result{Task: finished, Output: response.Result}, nil
}
if isLocalRateLimitError(err) {
lastErr = err
candidateErr = err
if task.AsyncMode && store.RateLimitRetryable(err) {
queued, delay, queueErr := s.requeueRateLimitedTask(ctx, task, err)
if queueErr != nil {
return Result{}, queueErr
}
return Result{Task: queued, Output: queued.Result}, &TaskQueuedError{Delay: delay}
}
break candidatesLoop
}
attemptNo = nextAttemptNo
lastErr = err
candidateErr = err
retryDecision := retryDecisionForCandidate(candidate, err)
@@ -297,9 +312,21 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user *auth.User, body map[string]any, candidate store.RuntimeModelCandidate, attemptNo int, onDelta clients.StreamDelta) (clients.Response, error) {
simulated := isSimulation(task, candidate)
if err := s.emit(ctx, task.ID, "task.attempt.started", "running", "submitting", 0.25, "client attempt started", map[string]any{"attempt": attemptNo, "clientId": candidate.ClientID}, simulated); err != nil {
return clients.Response{}, fmt.Errorf("emit attempt started: %w", err)
reservations := s.rateLimitReservations(ctx, user, candidate, body)
limitResult, err := s.store.ReserveRateLimits(ctx, task.ID, "", reservations)
if err != nil {
retryable := store.RateLimitRetryable(err)
clientErr := &clients.ClientError{Code: "rate_limit", Message: err.Error(), Retryable: retryable}
return clients.Response{}, &localRateLimitError{clientErr: clientErr, cause: err, retryAfter: localRateLimitRetryAfter(err)}
}
rateReservationsFinalized := false
defer func() {
if !rateReservationsFinalized {
_ = s.store.ReleaseRateLimitReservations(context.WithoutCancel(ctx), limitResult.Reservations, "attempt_failed")
}
}()
defer s.store.ReleaseConcurrencyLeases(context.WithoutCancel(ctx), limitResult.LeaseIDs)
attemptID, err := s.store.CreateTaskAttempt(ctx, store.CreateTaskAttemptInput{
TaskID: task.ID,
AttemptNo: attemptNo,
@@ -315,28 +342,21 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
if err != nil {
return clients.Response{}, fmt.Errorf("create task attempt: %w", err)
}
reservations := s.rateLimitReservations(ctx, user, candidate, body)
limitResult, err := s.store.ReserveRateLimits(ctx, task.ID, attemptID, reservations)
if err != nil {
retryable := store.RateLimitRetryable(err)
clientErr := &clients.ClientError{Code: "rate_limit", Message: err.Error(), Retryable: retryable}
if err := s.store.AttachRateLimitResultToAttempt(ctx, attemptID, limitResult); err != nil {
clientErr := &clients.ClientError{Code: "runtime_error", Message: err.Error(), Retryable: false}
_ = s.store.FinishTaskAttempt(ctx, store.FinishTaskAttemptInput{
AttemptID: attemptID,
Status: "failed",
Retryable: retryable,
Metrics: mergeMetrics(attemptMetrics(candidate, attemptNo, simulated), map[string]any{"error": err.Error(), "retryable": retryable, "retryAfterMs": localRateLimitRetryAfter(err).Milliseconds(), "trace": []any{failureTraceEntry(clientErr, retryable)}}),
ErrorCode: "rate_limit",
Retryable: false,
Metrics: mergeMetrics(attemptMetrics(candidate, attemptNo, simulated), map[string]any{"error": err.Error(), "retryable": false, "trace": []any{failureTraceEntry(clientErr, false)}}),
ErrorCode: clients.ErrorCode(clientErr),
ErrorMessage: err.Error(),
})
return clients.Response{}, &localRateLimitError{clientErr: clientErr, cause: err, retryAfter: localRateLimitRetryAfter(err)}
return clients.Response{}, fmt.Errorf("attach rate limit reservations to attempt: %w", err)
}
if err := s.emit(ctx, task.ID, "task.attempt.started", "running", "submitting", 0.25, "client attempt started", map[string]any{"attempt": attemptNo, "clientId": candidate.ClientID}, simulated); err != nil {
return clients.Response{}, fmt.Errorf("emit attempt started: %w", err)
}
rateReservationsFinalized := false
defer func() {
if !rateReservationsFinalized {
_ = s.store.ReleaseRateLimitReservations(context.WithoutCancel(ctx), limitResult.Reservations, "attempt_failed")
}
}()
defer s.store.ReleaseConcurrencyLeases(context.WithoutCancel(ctx), limitResult.LeaseIDs)
if err := s.store.RecordClientAssignment(ctx, candidate); err != nil {
return clients.Response{}, fmt.Errorf("record client assignment: %w", err)