From daad121f12650d8e385c7231dac9448a12ad930b Mon Sep 17 00:00:00 2001 From: chiyi Date: Sat, 4 Jul 2026 15:23:43 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E8=B0=83=E5=BA=A6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/internal/repo/model_repo.go | 13 +++++++ backend/internal/service/v1.go | 59 +++++++++++++++++++++++++++-- 2 files changed, 68 insertions(+), 4 deletions(-) diff --git a/backend/internal/repo/model_repo.go b/backend/internal/repo/model_repo.go index 367cf8b..3a5866f 100644 --- a/backend/internal/repo/model_repo.go +++ b/backend/internal/repo/model_repo.go @@ -44,6 +44,19 @@ func (r *ModelRepository) Get(ctx context.Context, modelID string) (*model.Model return &item, nil } +// GetAllMatching returns all models that match the given modelID (by alias or id). +// Used for multi-model load balancing when multiple models share the same alias. +func (r *ModelRepository) GetAllMatching(ctx context.Context, modelID string) ([]model.ModelConfig, error) { + var items []model.ModelConfig + if err := r.db.WithContext(ctx). + Where("(alias <> '' AND alias = ?) OR (alias = '' AND id = ?)", modelID, modelID). + Order("weight desc, created_at desc"). + Find(&items).Error; err != nil { + return nil, err + } + return items, nil +} + func (r *ModelRepository) NameMap(ctx context.Context) (map[string]string, error) { items, err := r.List(ctx) if err != nil { diff --git a/backend/internal/service/v1.go b/backend/internal/service/v1.go index efdd0b2..1dfeb23 100644 --- a/backend/internal/service/v1.go +++ b/backend/internal/service/v1.go @@ -91,6 +91,11 @@ type V1Service struct { // two simultaneous requests never start on the same account. tokenCursors sync.Map + // modelCursors holds one strict round-robin cursor per model alias/id (key: modelID, + // value: *uint64). Each pick advances the model's cursor by one so models + // are used in a fixed, even rotation when multiple models share the same alias. + modelCursors sync.Map + // inflight maps an in-progress event ID → the cancel func of its generation // work context, so the maintenance sweep can stop a stuck generation the // moment it abandons the row (instead of letting an orphaned goroutine run on @@ -961,16 +966,29 @@ func (s *V1Service) prepareImage(ctx context.Context, principal *APIPrincipal, i if modelID == "" || prompt == "" { return nil, "", "", 0, errors.New("model and prompt required") } - modelItem, err := s.models.Get(ctx, modelID) + // Try to get all matching models for load balancing + modelItems, err := s.models.GetAllMatching(ctx, modelID) if err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { return nil, "", "", 0, ErrUnknownModel } return nil, "", "", 0, err } - if !modelItem.Enabled || modelItem.Type != "image" { + if len(modelItems) == 0 { return nil, "", "", 0, ErrUnknownModel } + // Filter enabled image models + var validModels []*model.ModelConfig + for i := range modelItems { + if modelItems[i].Enabled && modelItems[i].Type == "image" { + validModels = append(validModels, &modelItems[i]) + } + } + if len(validModels) == 0 { + return nil, "", "", 0, ErrUnknownModel + } + // Select model using round-robin if multiple valid models + modelItem := s.selectModelByRoundRobin(validModels, modelID) // Fail fast before charging if the provider has no usable account. Use the // effective provider: a custom upstream serving this model id routes to // "custom" (effectiveProvider only returns it when such an account exists, so @@ -1031,16 +1049,29 @@ func (s *V1Service) prepareVideo(ctx context.Context, principal *APIPrincipal, i if duration == "" { return nil, "", "", "", 0, errors.New("duration required") } - modelItem, err := s.models.Get(ctx, modelID) + // Try to get all matching models for load balancing + modelItems, err := s.models.GetAllMatching(ctx, modelID) if err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { return nil, "", "", "", 0, ErrUnknownModel } return nil, "", "", "", 0, err } - if !modelItem.Enabled || modelItem.Type != "video" { + if len(modelItems) == 0 { return nil, "", "", "", 0, ErrUnknownModel } + // Filter enabled video models + var validModels []*model.ModelConfig + for i := range modelItems { + if modelItems[i].Enabled && modelItems[i].Type == "video" { + validModels = append(validModels, &modelItems[i]) + } + } + if len(validModels) == 0 { + return nil, "", "", "", 0, ErrUnknownModel + } + // Select model using round-robin if multiple valid models + modelItem := s.selectModelByRoundRobin(validModels, modelID) // Fail fast before charging — effective provider (custom upstream by id, else native). if eff := s.effectiveProvider(ctx, modelItem); eff == "custom" { // custom serves this id (effectiveProvider guaranteed it) — precheck ok @@ -2987,6 +3018,26 @@ func (s *V1Service) nextCursor(pool string) uint64 { return atomic.AddUint64(v.(*uint64), 1) - 1 } +// selectModelByRoundRobin selects a model from a list using round-robin. +// The cursor is per modelID (alias), so different aliases rotate independently. +func (s *V1Service) selectModelByRoundRobin(models []*model.ModelConfig, modelID string) *model.ModelConfig { + if len(models) == 0 { + return nil + } + if len(models) == 1 { + return models[0] + } + cursor := s.nextModelCursor(modelID) + index := cursor % uint64(len(models)) + return models[index] +} + +// nextModelCursor returns the next round-robin index for a model alias/id. +func (s *V1Service) nextModelCursor(modelID string) uint64 { + v, _ := s.modelCursors.LoadOrStore(modelID, new(uint64)) + return atomic.AddUint64(v.(*uint64), 1) - 1 +} + // rotateRoundRobin orders the active accounts by a stable key (ID) and rotates // the slice in place so iteration begins at the pool's current cursor position, // then advances the cursor. This is strict round-robin: account selection