优化调度

This commit is contained in:
2026-07-04 15:23:43 +08:00
parent b4e9b233ca
commit daad121f12
2 changed files with 68 additions and 4 deletions
+13
View File
@@ -44,6 +44,19 @@ func (r *ModelRepository) Get(ctx context.Context, modelID string) (*model.Model
return &item, nil 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) { func (r *ModelRepository) NameMap(ctx context.Context) (map[string]string, error) {
items, err := r.List(ctx) items, err := r.List(ctx)
if err != nil { if err != nil {
+55 -4
View File
@@ -91,6 +91,11 @@ type V1Service struct {
// two simultaneous requests never start on the same account. // two simultaneous requests never start on the same account.
tokenCursors sync.Map 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 // 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 // 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 // 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 == "" { if modelID == "" || prompt == "" {
return nil, "", "", 0, errors.New("model and prompt required") 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 err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) { if errors.Is(err, gorm.ErrRecordNotFound) {
return nil, "", "", 0, ErrUnknownModel return nil, "", "", 0, ErrUnknownModel
} }
return nil, "", "", 0, err return nil, "", "", 0, err
} }
if !modelItem.Enabled || modelItem.Type != "image" { if len(modelItems) == 0 {
return nil, "", "", 0, ErrUnknownModel 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 // Fail fast before charging if the provider has no usable account. Use the
// effective provider: a custom upstream serving this model id routes to // effective provider: a custom upstream serving this model id routes to
// "custom" (effectiveProvider only returns it when such an account exists, so // "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 == "" { if duration == "" {
return nil, "", "", "", 0, errors.New("duration required") 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 err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) { if errors.Is(err, gorm.ErrRecordNotFound) {
return nil, "", "", "", 0, ErrUnknownModel return nil, "", "", "", 0, ErrUnknownModel
} }
return nil, "", "", "", 0, err return nil, "", "", "", 0, err
} }
if !modelItem.Enabled || modelItem.Type != "video" { if len(modelItems) == 0 {
return nil, "", "", "", 0, ErrUnknownModel 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). // Fail fast before charging — effective provider (custom upstream by id, else native).
if eff := s.effectiveProvider(ctx, modelItem); eff == "custom" { if eff := s.effectiveProvider(ctx, modelItem); eff == "custom" {
// custom serves this id (effectiveProvider guaranteed it) — precheck ok // 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 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 // 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, // 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 // then advances the cursor. This is strict round-robin: account selection