feat(queue): 增加非文本模型分布式准入队列

使用 PostgreSQL 统一同步与异步非文本任务的并发准入、持久化等待和 Worker 容量分配，并将生产 API 与独立 Worker 角色拆分。补充策略管理、共享契约、OpenAPI、Kubernetes 双节点 Worker 清单及跨节点验收脚本；未默认启用任何生产 queue_size 策略。已在原基线完成 Go、前端、迁移、Shell、Kustomize 与长任务容量验收；合入最新主干后将重新执行发布门禁。
2026-07-29 16:15:43 +08:00
parent 2e5a90731b
commit 9e4fc7362d
55 changed files with 4565 additions and 237 deletions
@@ -36,7 +36,11 @@ type Service struct {
 	riverExecutionClient asyncExecutionClient
 	riverDrainingClients map[asyncExecutionClient]struct{}
 	riverWorkerCapacity  int
+	workerInstanceID     string
 	asyncCapacityLoader  func(context.Context, int) (store.AsyncWorkerCapacitySnapshot, error)
+	admissionWakeMu      sync.Mutex
+	admissionWake        chan struct{}
+	admissionListener    sync.Once
 	asyncClientFactory   func(int) (asyncExecutionClient, error)
 	billingMetrics       billingMetricsObserver
 }
@@ -135,7 +139,9 @@ func New(cfg config.Config, db *store.Store, logger *slog.Logger, observers ...b
 			"universal":             clients.UniversalClient{HTTPClient: httpClients.none, ScriptExecutor: scriptExecutor},
 			"simulation":            clients.SimulationClient{},
 		},
-		httpClients: httpClients,
+		httpClients:      httpClients,
+		workerInstanceID: asyncWorkerID(),
+		admissionWake:    make(chan struct{}),
 	}
 	if len(observers) > 0 {
 		service.billingMetrics = observers[0]
@@ -172,7 +178,15 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut

 func (s *Service) executeWithToken(ctx context.Context, task store.GatewayTask, user *auth.User, onDelta clients.StreamDelta, executionToken string) (Result, error) {
 	wasRunning := task.Status == "running"
-	claimed, err := s.store.ClaimTaskExecution(ctx, task.ID, executionToken, taskExecutionLeaseTTL)
+	preliminaryModelType := modelTypeFromKind(task.Kind, normalizeRequest(task.Kind, task.Request))
+	distributedAdmission := distributedAdmissionModelType(preliminaryModelType)
+	var claimed store.GatewayTask
+	var err error
+	if distributedAdmission && !wasRunning {
+		claimed, err = s.store.ClaimTaskPreparation(ctx, task.ID, executionToken, taskExecutionLeaseTTL)
+	} else {
+		claimed, err = s.store.ClaimTaskExecution(ctx, task.ID, executionToken, taskExecutionLeaseTTL)
+	}
 	if err != nil {
 		return Result{}, err
 	}
@@ -181,6 +195,21 @@ func (s *Service) executeWithToken(ctx context.Context, task store.GatewayTask,
 	defer stopExecution()
 	go s.renewTaskExecutionLease(executionCtx, stopExecution, task.ID, task.ExecutionToken)
 	ctx = executionCtx
+	defer func() {
+		if !requestCancelledBeforeSubmission(ctx) {
+			return
+		}
+		cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 10*time.Second)
+		defer cancel()
+		if _, _, cleanupErr := s.store.CancelTaskBeforeUpstreamSubmission(
+			cleanupCtx,
+			task.ID,
+			task.ExecutionToken,
+			"client disconnected before upstream submission",
+		); cleanupErr != nil && s.logger != nil {
+			s.logger.Warn("cancel disconnected task before upstream submission failed", "taskID", task.ID, "error", cleanupErr)
+		}
+	}()
 	executeStartedAt := time.Now()
 	restoredRequest, err := s.restoreTaskRequestReferences(ctx, task)
 	if err != nil {
@@ -189,13 +218,16 @@ func (s *Service) executeWithToken(ctx context.Context, task store.GatewayTask,
 	body := normalizeRequest(task.Kind, restoredRequest)
 	responseExecution := responseExecutionContext{}
 	modelType := modelTypeFromKind(task.Kind, body)
-	if err := s.store.MarkTaskRunning(ctx, task.ID, task.ExecutionToken, modelType, s.slimTaskRequestSnapshot(task, body)); err != nil {
-		return Result{}, err
-	}
-	if !wasRunning {
-		if err := s.emit(ctx, task.ID, "task.running", "running", "starting", 0.12, "task pulled from queue and started", map[string]any{"modelType": modelType}, task.RunMode == "simulation"); err != nil {
+	distributedAdmission = distributedAdmissionModelType(modelType)
+	if !distributedAdmission {
+		if err := s.store.MarkTaskRunning(ctx, task.ID, task.ExecutionToken, modelType, s.slimTaskRequestSnapshot(task, body)); err != nil {
 			return Result{}, err
 		}
+		if !wasRunning {
+			if err := s.emit(ctx, task.ID, "task.running", "running", "starting", 0.12, "task pulled from queue and started", map[string]any{"modelType": modelType}, task.RunMode == "simulation"); err != nil {
+				return Result{}, err
+			}
+		}
 	}
 	if err := validateRequest(task.Kind, body); err != nil {
 		validationErr := parameterPreprocessClientError(err)
@@ -477,6 +509,90 @@ func (s *Service) executeWithToken(ctx context.Context, task store.GatewayTask,
 			}
 		}
 	}
+	var admittedLeases []store.ConcurrencyLease
+	admittedPlatformModelID := ""
+	if distributedAdmission && len(candidates) > 0 {
+		admissionScopes, groupID := s.admissionScopes(ctx, user, candidates[0])
+		hasConcurrentLimit := false
+		for _, scope := range admissionScopes {
+			if scope.ConcurrentLimit > 0 {
+				hasConcurrentLimit = true
+				break
+			}
+		}
+		if hasConcurrentLimit {
+			if err := s.store.CheckRateLimits(ctx, s.rateLimitReservations(ctx, user, candidates[0], body)); err != nil {
+				failed, finishErr := s.failTask(ctx, task.ID, task.ExecutionToken, clients.ErrorCode(err), err.Error(), task.RunMode == "simulation", err)
+				if finishErr != nil {
+					return Result{}, finishErr
+				}
+				return Result{Task: failed, Output: failed.Result}, err
+			}
+			plan := taskAdmissionPlan{
+				Candidate: candidates[0],
+				Body:      body,
+				ModelType: modelType,
+				GroupID:   groupID,
+				Scopes:    admissionScopes,
+				Eligible:  true,
+			}
+			var admissionResult store.TaskAdmissionResult
+			var admissionErr error
+			if task.AsyncMode {
+				admissionResult, admissionErr = s.tryTaskAdmission(ctx, task, plan, "")
+				if admissionErr == nil && !admissionResult.Admitted {
+					_ = s.store.ReleaseTaskPreparation(context.WithoutCancel(ctx), task.ID, task.ExecutionToken)
+					return Result{Task: task, Output: task.Result}, &TaskQueuedError{Delay: time.Second}
+				}
+			} else {
+				admissionResult, admissionErr = s.waitForSynchronousAdmission(ctx, task, plan)
+			}
+			if admissionErr != nil {
+				if requestCancelledBeforeSubmission(ctx) {
+					cancelled, changed, cancelErr := s.store.CancelTaskBeforeUpstreamSubmission(
+						context.WithoutCancel(ctx),
+						task.ID,
+						task.ExecutionToken,
+						"client disconnected before upstream submission",
+					)
+					if cancelErr != nil {
+						return Result{}, cancelErr
+					}
+					if changed {
+						return Result{Task: cancelled, Output: cancelled.Result}, admissionErr
+					}
+				}
+				code := clients.ErrorCode(admissionErr)
+				if errors.Is(admissionErr, context.Canceled) {
+					code = "client_disconnected"
+				}
+				failed, finishErr := s.failTask(context.WithoutCancel(ctx), task.ID, task.ExecutionToken, code, admissionErr.Error(), task.RunMode == "simulation", admissionErr)
+				if finishErr != nil && !errors.Is(finishErr, store.ErrTaskExecutionLeaseLost) {
+					return Result{}, finishErr
+				}
+				return Result{Task: failed, Output: failed.Result}, admissionErr
+			}
+			admittedLeases = admissionResult.Leases
+			if len(admittedLeases) == 0 {
+				admittedLeases, err = s.store.ActiveTaskAdmissionLeases(ctx, task.ID)
+				if err != nil {
+					return Result{}, err
+				}
+			}
+			admittedPlatformModelID = admissionResult.Admission.PlatformModelID
+		}
+		if err := s.store.MarkTaskRunning(ctx, task.ID, task.ExecutionToken, modelType, s.slimTaskRequestSnapshot(task, body)); err != nil {
+			return Result{}, err
+		}
+		if !wasRunning {
+			if err := s.emit(ctx, task.ID, "task.running", "running", "starting", 0.12, "task admitted and started", map[string]any{"modelType": modelType}, task.RunMode == "simulation"); err != nil {
+				return Result{}, err
+			}
+		}
+	}
+	if distributedAdmission {
+		defer s.store.DeleteTaskAdmission(context.WithoutCancel(ctx), task.ID)
+	}
 	firstCandidateBody := body
 	normalizedModelType := modelType
 	attemptNo := task.AttemptCount
@@ -579,6 +695,33 @@ candidatesLoop:
 		if !available {
 			continue
 		}
+		if distributedAdmission && candidate.PlatformModelID != admittedPlatformModelID {
+			admissionResult, candidateLimited, admissionErr := s.ensureCandidateAdmission(ctx, task, user, body, candidate)
+			if admissionErr != nil {
+				lastErr = admissionErr
+				lastCandidate = candidate
+				if errors.Is(admissionErr, store.ErrQueueTimeout) {
+					break
+				}
+				if candidateLimited && errors.Is(admissionErr, store.ErrRateLimited) {
+					continue
+				}
+				return Result{}, admissionErr
+			}
+			if candidateLimited {
+				admittedPlatformModelID = admissionResult.Admission.PlatformModelID
+				admittedLeases = admissionResult.Leases
+				if len(admittedLeases) == 0 {
+					admittedLeases, err = s.store.ActiveTaskAdmissionLeases(ctx, task.ID)
+					if err != nil {
+						return Result{}, err
+					}
+				}
+			} else {
+				admittedPlatformModelID = ""
+				admittedLeases = nil
+			}
+		}
 		platformsVisited++
 		lastCandidate = candidate
 		clientAttempts := clientAttemptsForCandidate(candidate)
@@ -614,7 +757,21 @@ candidatesLoop:
 			}
 			candidateBody := preprocessing.Body
 			candidatePricing := pricingByCandidate[pricingCandidateKey(candidate)]
-			response, err := s.runCandidate(ctx, task, user, candidateBody, preprocessing.Log, candidate, candidatePricing, nextAttemptNo, candidateDelta, responseExecution, runnerPolicy.CacheAffinityPolicy, cacheAffinityKeys.Record)
+			response, err := s.runCandidate(ctx, task, user, candidateBody, preprocessing.Log, candidate, candidatePricing, nextAttemptNo, candidateDelta, responseExecution, runnerPolicy.CacheAffinityPolicy, cacheAffinityKeys.Record, admittedPlatformModelID, admittedLeases)
+			if err != nil && requestCancelledBeforeSubmission(ctx) {
+				cancelled, changed, cancelErr := s.store.CancelTaskBeforeUpstreamSubmission(
+					context.WithoutCancel(ctx),
+					task.ID,
+					task.ExecutionToken,
+					"client disconnected before upstream submission",
+				)
+				if cancelErr != nil {
+					return Result{}, cancelErr
+				}
+				if changed {
+					return Result{Task: cancelled, Output: cancelled.Result}, err
+				}
+			}
 			if err != nil && isVolcesRemoteTaskCancellation(candidate, err) {
 				cancelled, changed, cancelErr := s.store.CancelSubmittedTask(ctx, task.ID, task.ExecutionToken, "任务已由火山引擎取消")
 				if cancelErr != nil {
@@ -981,7 +1138,22 @@ func billingItemsFixedTotal(items []any) fixedAmount {
 	return total
 }

-func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user *auth.User, body map[string]any, preprocessing parameterPreprocessingLog, candidate store.RuntimeModelCandidate, pricing resolvedPricing, attemptNo int, onDelta clients.StreamDelta, responseExecution responseExecutionContext, cacheAffinityPolicy map[string]any, cacheAffinityRecordKeys []string) (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,
+	pricing resolvedPricing,
+	attemptNo int,
+	onDelta clients.StreamDelta,
+	responseExecution responseExecutionContext,
+	cacheAffinityPolicy map[string]any,
+	cacheAffinityRecordKeys []string,
+	admittedPlatformModelID string,
+	admittedLeases []store.ConcurrencyLease,
+) (clients.Response, error) {
 	var err error
 	candidate, err = candidateWithEnvironmentCredentials(candidate)
 	if err != nil {
@@ -990,12 +1162,24 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
 	simulated := isSimulation(task, candidate)
 	baseAttemptMetrics := mergeMetrics(attemptMetrics(candidate, attemptNo, simulated), parameterPreprocessingMetrics(preprocessing))
 	reservations := s.rateLimitReservations(ctx, user, candidate, body)
+	if admittedPlatformModelID == candidate.PlatformModelID && len(admittedLeases) > 0 {
+		filtered := make([]store.RateLimitReservation, 0, len(reservations))
+		for _, reservation := range reservations {
+			if reservation.Metric != "concurrent" {
+				filtered = append(filtered, reservation)
+			}
+		}
+		reservations = filtered
+	}
 	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)}
 	}
+	if admittedPlatformModelID == candidate.PlatformModelID {
+		limitResult.Leases = append(limitResult.Leases, admittedLeases...)
+	}
 	rateReservationsFinalized := false
 	defer func() {
 		if !rateReservationsFinalized {
@@ -1124,6 +1308,7 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
 	submissionStatus := "not_submitted"
 	if strings.TrimSpace(task.RemoteTaskID) != "" {
 		submissionStatus = "response_received"
+		markUpstreamSubmissionStarted(ctx)
 		if err := s.store.SetAttemptUpstreamSubmissionStatus(ctx, attemptID, submissionStatus); err != nil {
 			return clients.Response{}, fmt.Errorf("restore upstream submission status: %w", err)
 		}
@@ -1181,7 +1366,11 @@ func (s *Service) runCandidate(ctx context.Context, task store.GatewayTask, user
 			return setSubmissionStatus("response_received")
 		},
 		OnUpstreamSubmissionStarted: func() error {
-			return setSubmissionStatus("submitting")
+			if err := setSubmissionStatus("submitting"); err != nil {
+				return err
+			}
+			markUpstreamSubmissionStarted(ctx)
+			return nil
 		},
 		OnUpstreamResponseReceived: func() error {
 			return setSubmissionStatus("response_received")
@@ -1907,7 +2096,10 @@ func canonicalModelType(value string) string {

 func isKnownModelType(value string) bool {
 	switch value {
-	case "text_generate", "text_embedding", "text_rerank", "image_generate", "image_edit", "image_vectorize", "video_generate", "video_enhance", "image_to_video", "text_to_video", "video_edit", "video_reference", "video_first_last_frame", "omni_video", "omni", "audio_generate", "text_to_speech", "voice_clone":
+	case "text_generate", "text_embedding", "text_rerank",
+		"image_generate", "image_edit", "image_analysis", "image_vectorize",
+		"video_generate", "video_enhance", "image_to_video", "text_to_video", "video_edit", "video_reference", "video_first_last_frame", "video_understanding",
+		"omni_video", "omni", "audio_generate", "audio_understanding", "text_to_speech", "voice_clone":
 		return true
 	default:
 		return false