diff --git a/apps/api/internal/runner/admission.go b/apps/api/internal/runner/admission.go index 2291b0c..1e7bfae 100644 --- a/apps/api/internal/runner/admission.go +++ b/apps/api/internal/runner/admission.go @@ -103,6 +103,7 @@ func (s *Service) buildTaskAdmissionPlanForCurrentBinding( CacheAffinityKey: cacheAffinityKeys.Primary, CacheAffinityKeys: cacheAffinityKeys.Lookup, CacheAffinityPolicy: runnerPolicy.CacheAffinityPolicy, + QuotaScopePrefix: acceptanceQuotaScopePrefix(task), }) if err == nil { candidates, err = filterCandidatesByRequestedPlatform(candidates, body) @@ -120,6 +121,7 @@ func (s *Service) buildTaskAdmissionPlanForCurrentBinding( return taskAdmissionPlan{}, err } candidates, _ = pinCandidatesToTaskAdmission(candidates, admission) + var candidateRateLimitErr error for _, candidate := range candidates { available, availabilityErr := s.store.RuntimeCandidateAvailable(ctx, candidate.PlatformID, candidate.PlatformModelID) if availabilityErr != nil { @@ -151,6 +153,13 @@ func (s *Service) buildTaskAdmissionPlanForCurrentBinding( s.rateLimitReservations(ctx, user, candidate, body), ) if err := s.store.CheckRateLimits(ctx, reservations); err != nil { + var limitErr *store.RateLimitExceededError + if errors.As(err, &limitErr) && limitErr.ScopeType == "platform_model" { + if candidateRateLimitErr == nil { + candidateRateLimitErr = err + } + continue + } return taskAdmissionPlan{}, err } return taskAdmissionPlan{ @@ -162,6 +171,9 @@ func (s *Service) buildTaskAdmissionPlanForCurrentBinding( Eligible: true, }, nil } + if candidateRateLimitErr != nil { + return taskAdmissionPlan{}, candidateRateLimitErr + } return taskAdmissionPlan{}, store.ErrNoModelCandidate } @@ -239,6 +251,13 @@ func acceptanceScopeKey(task store.GatewayTask, scopeKey string) string { return "acceptance:" + runID + ":" + scopeKey } +func acceptanceQuotaScopePrefix(task store.GatewayTask) string { + if task.RunMode != "acceptance" { + return "" + } + return acceptanceScopeKey(task, "") +} + func (s *Service) loadAsyncTaskAdmission(ctx context.Context, task store.GatewayTask) (*store.TaskAdmission, error) { if !task.AsyncMode { return nil, nil diff --git a/apps/api/internal/runner/admission_test.go b/apps/api/internal/runner/admission_test.go index fbdecf5..683bf89 100644 --- a/apps/api/internal/runner/admission_test.go +++ b/apps/api/internal/runner/admission_test.go @@ -117,6 +117,19 @@ func TestAcceptanceInfrastructureReservationsIsolateEveryMetric(t *testing.T) { } } +func TestAcceptanceQuotaScopePrefixMatchesIsolatedCounters(t *testing.T) { + acceptanceTask := store.GatewayTask{RunMode: "acceptance", AcceptanceRunID: "run-1"} + if got := acceptanceQuotaScopePrefix(acceptanceTask); got != "acceptance:run-1:" { + t.Fatalf("acceptance quota scope prefix=%q", got) + } + if got := acceptanceQuotaScopePrefix(store.GatewayTask{RunMode: "acceptance_canary", AcceptanceRunID: "run-1"}); got != "" { + t.Fatalf("canary must use production quota scope, got %q", got) + } + if got := acceptanceQuotaScopePrefix(store.GatewayTask{RunMode: "production"}); got != "" { + t.Fatalf("production quota scope changed, got %q", got) + } +} + func TestAcceptanceAdmissionScopesLeaveProductionPolicyUnchanged(t *testing.T) { input := []store.AdmissionScope{{ ScopeType: "platform_model", diff --git a/apps/api/internal/runner/service.go b/apps/api/internal/runner/service.go index 1b3af7d..fb25b3a 100644 --- a/apps/api/internal/runner/service.go +++ b/apps/api/internal/runner/service.go @@ -372,6 +372,7 @@ func (s *Service) executeWithToken(ctx context.Context, task store.GatewayTask, CacheAffinityKey: cacheAffinityKeys.Primary, CacheAffinityKeys: cacheAffinityKeys.Lookup, CacheAffinityPolicy: runnerPolicy.CacheAffinityPolicy, + QuotaScopePrefix: acceptanceQuotaScopePrefix(task), }) if err == nil { candidates, err = filterCandidatesByRequestedPlatform(candidates, body) diff --git a/apps/api/internal/store/candidates.go b/apps/api/internal/store/candidates.go index 9ad6200..a14ce34 100644 --- a/apps/api/internal/store/candidates.go +++ b/apps/api/internal/store/candidates.go @@ -17,6 +17,7 @@ type ListModelCandidatesOptions struct { CacheAffinityKey string CacheAffinityKeys []string CacheAffinityPolicy map[string]any + QuotaScopePrefix string } func (s *Store) ListModelCandidates(ctx context.Context, model string, modelType string, user *auth.User, options ...ListModelCandidatesOptions) ([]RuntimeModelCandidate, error) { @@ -102,7 +103,7 @@ LEFT JOIN ( AND released_at IS NULL AND expires_at > statement_timestamp() GROUP BY scope_key - ) con ON con.scope_key = m.id::text + ) con ON con.scope_key = $5::text || m.id::text LEFT JOIN ( SELECT queued_sources.platform_model_id, COUNT(DISTINCT queued_sources.task_id) AS waiting FROM ( @@ -137,7 +138,7 @@ LEFT JOIN ( AND metric = 'rpm' AND reset_at > now() ORDER BY scope_key, window_start DESC - ) rpm ON rpm.scope_key = m.id::text + ) rpm ON rpm.scope_key = $5::text || m.id::text LEFT JOIN ( SELECT scope_key, SUM(used_value) AS used_value, SUM(reserved_value) AS reserved_value FROM gateway_rate_limit_counters @@ -145,7 +146,7 @@ LEFT JOIN ( AND metric LIKE 'tpm%' AND reset_at > now() GROUP BY scope_key - ) tpm ON tpm.scope_key = m.id::text + ) tpm ON tpm.scope_key = $5::text || m.id::text WHERE p.status = 'enabled' AND p.deleted_at IS NULL AND m.enabled = true @@ -170,7 +171,7 @@ WHERE p.status = 'enabled' COALESCE(s.running_count, 0) ASC, COALESCE(s.waiting_count, 0) ASC, COALESCE(s.last_assigned_at, to_timestamp(0)) ASC, - m.created_at ASC`, exactModel, modelType, listOptions.CacheAffinityKeys, cacheAffinityStaleAfterSeconds(listOptions.CacheAffinityPolicy)) + m.created_at ASC`, exactModel, modelType, listOptions.CacheAffinityKeys, cacheAffinityStaleAfterSeconds(listOptions.CacheAffinityPolicy), listOptions.QuotaScopePrefix) if err != nil { return nil, err } @@ -472,6 +473,7 @@ func normalizeListModelCandidatesOptions(modelType string, options ...ListModelC } out := options[0] out.CacheAffinityKey = strings.TrimSpace(out.CacheAffinityKey) + out.QuotaScopePrefix = strings.TrimSpace(out.QuotaScopePrefix) out.CacheAffinityKeys = normalizedCacheAffinityKeys(out.CacheAffinityKey, out.CacheAffinityKeys) if len(out.CacheAffinityKeys) > 0 { out.CacheAffinityKey = out.CacheAffinityKeys[0] diff --git a/deploy/kubernetes/local-acceptance/capacity-controller-rbac.yaml b/deploy/kubernetes/local-acceptance/capacity-controller-rbac.yaml new file mode 100644 index 0000000..91a4066 --- /dev/null +++ b/deploy/kubernetes/local-acceptance/capacity-controller-rbac.yaml @@ -0,0 +1,17 @@ +apiVersion: rbac.authorization.k8s.io/v1 +kind: Role +metadata: + name: easyai-capacity-controller + namespace: easyai +rules: + - apiGroups: ["apps"] + resources: ["deployments"] + resourceNames: ["easyai-worker-ningbo", "easyai-worker-hongkong"] + verbs: ["get"] + - apiGroups: ["apps"] + resources: ["deployments/scale"] + resourceNames: ["easyai-worker-ningbo", "easyai-worker-hongkong"] + verbs: ["get", "update", "patch"] + - apiGroups: [""] + resources: ["pods"] + verbs: ["get", "list", "patch"] diff --git a/scripts/acceptance/local-cluster.sh b/scripts/acceptance/local-cluster.sh index 709a169..9ed5866 100755 --- a/scripts/acceptance/local-cluster.sh +++ b/scripts/acceptance/local-cluster.sh @@ -599,12 +599,20 @@ render_and_apply_application() { chmod 0600 "$rendered" kubectl --context "$context" -n "$namespace" apply \ -f "$repository_root/deploy/kubernetes/production/service-account-rbac.yaml" >/dev/null + kubectl --context "$context" -n "$namespace" apply \ + -f "$manifest_root/capacity-controller-rbac.yaml" >/dev/null kubectl --context "$context" -n "$namespace" apply -f "$manifest_root/local-config.yaml" >/dev/null kubectl --context "$context" -n "$namespace" apply -f "$rendered" >/dev/null + kubectl --context "$context" -n "$namespace" scale \ + deployment/easyai-worker-hongkong --replicas=0 >/dev/null + kubectl --context "$context" -n "$namespace" rollout status \ + deployment/easyai-worker-hongkong --timeout=5m >/dev/null for workload in easyai-worker-ningbo easyai-worker-hongkong; do kubectl --context "$context" -n "$namespace" patch deployment "$workload" \ --type=merge -p='{"spec":{"template":{"spec":{"affinity":null}}}}' >/dev/null done + kubectl --context "$context" -n "$namespace" scale \ + deployment/easyai-worker-hongkong --replicas=2 >/dev/null for workload in easyai-api-ningbo easyai-worker-ningbo; do kubectl --context "$context" -n "$namespace" set env deployment/"$workload" \ diff --git a/scripts/acceptance/provider-burst.sh b/scripts/acceptance/provider-burst.sh index 37ecfb2..33186d4 100644 --- a/scripts/acceptance/provider-burst.sh +++ b/scripts/acceptance/provider-burst.sh @@ -543,6 +543,9 @@ provider_burst() { sample_resources "$resources" "$worker_sample_stop" & active_resource_sampler_pid=$! current_phase=provider_burst_load + while ((10#$(date -u '+%S') > 5)); do + sleep 1 + done run_load video-provider-quota "$load_report" -requests "$requests" sleep 2 touch "$worker_sample_stop"