fix(runner): 保留同平台重试的并发租约
同平台 retry_same 之前会在首个 attempt 失败时释放任务级 admission 租约,下一次 attempt 继续续租旧租约并进入 upstream_submission_unknown。 现在明确区分 attempt 自有租约与任务级 admission 租约,只在 attempt 结束时释放前者;补充回归测试覆盖同平台重试继续续租 admission 租约。 验证:env -u AI_GATEWAY_TEST_DATABASE_URL go test ./... -count=1;go vet ./...;gofmt;git diff --check。
This commit is contained in:
@@ -31,3 +31,29 @@ func TestLeaseRenewalFatalityDistinguishesTransientDatabaseErrors(t *testing.T)
|
|||||||
t.Fatal("concurrency lease renewal error at expiry must be fatal")
|
t.Fatal("concurrency lease renewal error at expiry must be fatal")
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestBindAttemptConcurrencyLeasesPreservesAdmissionLeaseAcrossRetry(t *testing.T) {
|
||||||
|
attemptLease := store.ConcurrencyLease{ID: "attempt-lease", TTL: time.Minute}
|
||||||
|
admissionLeases := []store.ConcurrencyLease{
|
||||||
|
{ID: "admission-user-group", TTL: 2 * time.Minute},
|
||||||
|
{ID: "admission-worker-capacity", TTL: 2 * time.Minute},
|
||||||
|
}
|
||||||
|
|
||||||
|
active, release := bindAttemptConcurrencyLeases([]store.ConcurrencyLease{attemptLease}, admissionLeases)
|
||||||
|
if len(active) != 3 || active[0].ID != attemptLease.ID || active[1].ID != admissionLeases[0].ID || active[2].ID != admissionLeases[1].ID {
|
||||||
|
t.Fatalf("active leases must renew attempt and admission ownership together: %+v", active)
|
||||||
|
}
|
||||||
|
if len(release) != 1 || release[0].ID != attemptLease.ID {
|
||||||
|
t.Fatalf("attempt cleanup must not release task-scoped admission leases: %+v", release)
|
||||||
|
}
|
||||||
|
|
||||||
|
// A same-platform retry receives the same admission lease IDs. The first
|
||||||
|
// attempt cleanup must leave them renewable by the next attempt.
|
||||||
|
retryActive, retryRelease := bindAttemptConcurrencyLeases(nil, admissionLeases)
|
||||||
|
if len(retryActive) != 2 || retryActive[0].ID != admissionLeases[0].ID || retryActive[1].ID != admissionLeases[1].ID {
|
||||||
|
t.Fatalf("retry must keep using the live admission leases: %+v", retryActive)
|
||||||
|
}
|
||||||
|
if len(retryRelease) != 0 {
|
||||||
|
t.Fatalf("retry without attempt-owned leases must not release admission leases: %+v", retryRelease)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -1357,20 +1357,18 @@ func (s *Service) runCandidate(
|
|||||||
clientErr := &clients.ClientError{Code: "rate_limit", Message: err.Error(), Retryable: retryable}
|
clientErr := &clients.ClientError{Code: "rate_limit", Message: err.Error(), Retryable: retryable}
|
||||||
return clients.Response{}, &localRateLimitError{clientErr: clientErr, cause: err, retryAfter: localRateLimitRetryAfter(err)}
|
return clients.Response{}, &localRateLimitError{clientErr: clientErr, cause: err, retryAfter: localRateLimitRetryAfter(err)}
|
||||||
}
|
}
|
||||||
|
attemptOwnedLeases := append([]store.ConcurrencyLease(nil), limitResult.Leases...)
|
||||||
if admittedPlatformModelID == candidate.PlatformModelID {
|
if admittedPlatformModelID == candidate.PlatformModelID {
|
||||||
limitResult.Leases = append(limitResult.Leases, admittedLeases...)
|
limitResult.Leases, attemptOwnedLeases = bindAttemptConcurrencyLeases(limitResult.Leases, admittedLeases)
|
||||||
}
|
}
|
||||||
rateReservationsFinalized := false
|
rateReservationsFinalized := false
|
||||||
retainAdmittedLeases := false
|
|
||||||
defer func() {
|
defer func() {
|
||||||
if !rateReservationsFinalized {
|
if !rateReservationsFinalized {
|
||||||
_ = s.store.ReleaseRateLimitReservations(context.WithoutCancel(ctx), limitResult.Reservations, "attempt_failed")
|
_ = s.store.ReleaseRateLimitReservations(context.WithoutCancel(ctx), limitResult.Reservations, "attempt_failed")
|
||||||
}
|
}
|
||||||
}()
|
}()
|
||||||
defer func() {
|
defer func() {
|
||||||
if !retainAdmittedLeases {
|
_ = s.store.ReleaseConcurrencyLeases(context.WithoutCancel(ctx), attemptOwnedLeases)
|
||||||
_ = s.store.ReleaseConcurrencyLeases(context.WithoutCancel(ctx), limitResult.Leases)
|
|
||||||
}
|
|
||||||
}()
|
}()
|
||||||
|
|
||||||
attemptID, err := s.store.CreateTaskAttempt(ctx, store.CreateTaskAttemptInput{
|
attemptID, err := s.store.CreateTaskAttempt(ctx, store.CreateTaskAttemptInput{
|
||||||
@@ -1798,12 +1796,20 @@ func (s *Service) runCandidate(
|
|||||||
}); err != nil {
|
}); err != nil {
|
||||||
s.logger.Warn("record cache affinity observation failed", "error", err, "clientId", candidate.ClientID)
|
s.logger.Warn("record cache affinity observation failed", "error", err, "clientId", candidate.ClientID)
|
||||||
}
|
}
|
||||||
if admittedPlatformModelID == candidate.PlatformModelID {
|
|
||||||
retainAdmittedLeases = true
|
|
||||||
}
|
|
||||||
return response, nil
|
return response, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// bindAttemptConcurrencyLeases combines task-scoped admission leases with
|
||||||
|
// leases acquired by the current attempt. Only the latter are returned for
|
||||||
|
// attempt cleanup: admission owns its leases across retry_same and releases
|
||||||
|
// them when the task changes binding or reaches a terminal state.
|
||||||
|
func bindAttemptConcurrencyLeases(attemptLeases []store.ConcurrencyLease, admissionLeases []store.ConcurrencyLease) ([]store.ConcurrencyLease, []store.ConcurrencyLease) {
|
||||||
|
attemptOwned := append([]store.ConcurrencyLease(nil), attemptLeases...)
|
||||||
|
active := append([]store.ConcurrencyLease(nil), attemptLeases...)
|
||||||
|
active = append(active, admissionLeases...)
|
||||||
|
return active, attemptOwned
|
||||||
|
}
|
||||||
|
|
||||||
func (s *Service) observeProviderQuotaWait(metric string) {
|
func (s *Service) observeProviderQuotaWait(metric string) {
|
||||||
observer, ok := s.billingMetrics.(interface {
|
observer, ok := s.billingMetrics.(interface {
|
||||||
ObserveProviderQuotaWait(string)
|
ObserveProviderQuotaWait(string)
|
||||||
|
|||||||
Reference in New Issue
Block a user