mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-09-20 00:48:00 +00:00
feat(scheduler): separate image-generation AUTO chain with UI toggle
Image models previously could not be scheduled through a priority chain: the chat AUTO chain explicitly skips image-kind slots, and AUTO image requests fell back to unordered registry discovery. - config: add auto_image rules (auto_image yaml / image_rules json); legacy auto rules keep their meaning as the chat chain - core: buildAutoImageChain mirrors buildAutoChain with inverted kind filter (image-only); SaveAutoImageRules + AutoImageRules/AutoImageChain - scheduler: ChainImage walks the chain tier-by-tier with round-robin and preference ordering, skipping cooling slots - gateway: handleImage AUTO now runs down AutoImageChain when one is configured (falls back to legacy discovery otherwise) and records the actual served model; handleAutoAPI GET returns image_rules and PUT accepts image_rules independently of rules - webui: priority page gains a chat/image toggle editing two independent lane sets; add-slot picker filters by active kind; persistAuto writes only the active chain's field
This commit is contained in:
@ -32,7 +32,8 @@ type Core struct {
|
||||
store *config.Store
|
||||
scheduler *scheduler.Scheduler
|
||||
registry *provider.Registry
|
||||
autoChain atomic.Pointer[scheduler.Chain]
|
||||
autoChain atomic.Pointer[scheduler.Chain] // chat AUTO chain
|
||||
autoImageChain atomic.Pointer[scheduler.Chain] // image-generation AUTO chain
|
||||
}
|
||||
|
||||
// New builds the core from a config file plus runtime overlay.
|
||||
@ -181,6 +182,9 @@ func (c *Core) Registry() *provider.Registry { return c.registry }
|
||||
// AutoChain returns the current AUTO scheduling chain.
|
||||
func (c *Core) AutoChain() *scheduler.Chain { return c.autoChain.Load() }
|
||||
|
||||
// AutoImageChain returns the persisted image-generation AUTO chain snapshot.
|
||||
func (c *Core) AutoImageChain() *scheduler.Chain { return c.autoImageChain.Load() }
|
||||
|
||||
func (c *Core) DefaultModel() string { return c.cfg.DefaultModel }
|
||||
|
||||
func (c *Core) GatewayKeys() []string { return c.cfg.GatewayKeys }
|
||||
@ -294,6 +298,15 @@ func (c *Core) AutoRules() []config.ModelScope {
|
||||
return out
|
||||
}
|
||||
|
||||
// AutoImageRules returns the configured image-generation AUTO chain rules.
|
||||
func (c *Core) AutoImageRules() []config.ModelScope {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
out := make([]config.ModelScope, len(c.cfg.AutoImage))
|
||||
copy(out, c.cfg.AutoImage)
|
||||
return out
|
||||
}
|
||||
|
||||
// cleanScopes drops empty model entries and normalizes placeholder source
|
||||
// names; it returns nil when no entries survive so an empty scope means
|
||||
// "unrestricted" (nil) instead of a restrictive-but-empty list — a non-nil
|
||||
@ -325,16 +338,35 @@ func (c *Core) SaveAutoRules(entries []config.ModelScope) error {
|
||||
return err
|
||||
}
|
||||
c.buildAutoChain()
|
||||
if ch := c.autoChain.Load(); ch != nil {
|
||||
for _, tn := range ch.Tiers {
|
||||
for _, sl := range tn.Slots {
|
||||
if p := c.registry.ProviderForSlot(sl.Model, sl.Source); p != nil {
|
||||
p.ResetModelCooldown(sl.Model)
|
||||
}
|
||||
c.resetChainCooldowns(c.autoChain.Load())
|
||||
return nil
|
||||
}
|
||||
|
||||
// SaveAutoImageRules persists the image-generation AUTO chain. Image models
|
||||
// are kept as-is (unlike buildAutoChain, which skips them).
|
||||
func (c *Core) SaveAutoImageRules(entries []config.ModelScope) error {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.cfg.AutoImage = cleanScopes(entries)
|
||||
if err := c.saveConfig(); err != nil {
|
||||
return err
|
||||
}
|
||||
c.buildAutoImageChain()
|
||||
c.resetChainCooldowns(c.autoImageChain.Load())
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Core) resetChainCooldowns(ch *scheduler.Chain) {
|
||||
if ch == nil {
|
||||
return
|
||||
}
|
||||
for _, tn := range ch.Tiers {
|
||||
for _, sl := range tn.Slots {
|
||||
if p := c.registry.ProviderForSlot(sl.Model, sl.Source); p != nil {
|
||||
p.ResetModelCooldown(sl.Model)
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ResetHealth clears the scheduling backoff state of every provider.
|
||||
@ -420,6 +452,7 @@ func (c *Core) rebuildRegistry() error {
|
||||
c.registry.Replace(providers)
|
||||
}
|
||||
c.buildAutoChain()
|
||||
c.buildAutoImageChain()
|
||||
return nil
|
||||
}
|
||||
|
||||
@ -466,6 +499,47 @@ func (c *Core) buildAutoChain() {
|
||||
c.autoChain.Store(scheduler.BuildChain(sr, prov))
|
||||
}
|
||||
|
||||
// buildAutoImageChain rebuilds the image-generation AUTO chain snapshot.
|
||||
// Unlike buildAutoChain, only image-kind models participate: chat models in
|
||||
// the rules are skipped so a stale chat slot can't receive image traffic.
|
||||
func (c *Core) buildAutoImageChain() {
|
||||
prov := func(model, source string) scheduler.Provider {
|
||||
p := c.registry.ProviderForSlot(model, source)
|
||||
if p == nil {
|
||||
return nil
|
||||
}
|
||||
if m := p.ModelByID(model); m != nil && m.Kind != "image" {
|
||||
return nil
|
||||
}
|
||||
return p
|
||||
}
|
||||
rules := c.cfg.AutoImage
|
||||
sr := make([]scheduler.Rule, 0, len(rules))
|
||||
for _, e := range rules {
|
||||
model, source := e.Model, e.Source
|
||||
if p := c.registry.ProviderForSlot(e.Model, e.Source); p != nil {
|
||||
if exact := p.ModelIDFold(e.Model); exact != "" {
|
||||
model = exact
|
||||
}
|
||||
if m := p.ModelByID(model); m != nil && m.Kind != "image" {
|
||||
continue // chat-kind models never join the image AUTO chain
|
||||
}
|
||||
if source == "" {
|
||||
source = p.Name()
|
||||
}
|
||||
}
|
||||
sr = append(sr, scheduler.Rule{
|
||||
Model: model,
|
||||
Source: source,
|
||||
Tier: e.Tier,
|
||||
Quota: e.TokenQuota,
|
||||
Period: e.Period,
|
||||
Hours: e.Hours,
|
||||
})
|
||||
}
|
||||
c.autoImageChain.Store(scheduler.BuildChain(sr, prov))
|
||||
}
|
||||
|
||||
// AutoSlotState is the UI-facing health snapshot of one AUTO chain slot.
|
||||
type AutoSlotState struct {
|
||||
Model string `json:"model"`
|
||||
|
||||
Reference in New Issue
Block a user