优化调度
This commit is contained in:
@@ -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 {
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user