388 lines
12 KiB
Go
388 lines
12 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"sort"
|
|
"time"
|
|
)
|
|
|
|
// normalizeGroupPolicy applies queue defaults while preserving an explicit
|
|
// zero queue timeout as "no timeout".
|
|
func normalizeGroupPolicy(policy groupPolicy) groupPolicy {
|
|
if policy.maxQueue <= 0 {
|
|
policy.maxQueue = defaultGroupMaxQueue
|
|
}
|
|
if policy.queueTimeout > 0 {
|
|
policy.queueTimeoutSet = true
|
|
}
|
|
if !policy.queueTimeoutSet {
|
|
policy.queueTimeout = defaultQueueTimeout
|
|
policy.queueTimeoutSet = true
|
|
}
|
|
return policy
|
|
}
|
|
|
|
// getOrCreateGroupLocked returns an existing group or creates one with the
|
|
// given policy (applying defaults for unset fields). Policy is only applied
|
|
// at creation; existing groups retain their current policy until an explicit
|
|
// updateGroupPolicyLocked call (used for config-refresh policy propagation).
|
|
func (m *modelQueueManager) getOrCreateGroupLocked(key string, policy groupPolicy) *modelQueueGroup {
|
|
g, ok := m.groups[key]
|
|
if !ok {
|
|
policy = normalizeGroupPolicy(policy)
|
|
g = &modelQueueGroup{
|
|
key: key,
|
|
policy: policy,
|
|
inflight: make(map[string]int),
|
|
longInflight: make(map[string]int),
|
|
}
|
|
m.groups[key] = g
|
|
}
|
|
return g
|
|
}
|
|
|
|
// updateGroupPolicyLocked refreshes an existing group's queue policy so that
|
|
// provider max_queue / queue_timeout_ms changes applied via config refresh take
|
|
// effect on subsequent admissions. Only configured policy fields override the
|
|
// current values; in-flight counts and queued items are preserved untouched.
|
|
func (m *modelQueueManager) updateGroupPolicyLocked(g *modelQueueGroup, policy groupPolicy) {
|
|
if g == nil {
|
|
return
|
|
}
|
|
if policy.maxQueue > 0 {
|
|
g.policy.maxQueue = policy.maxQueue
|
|
}
|
|
if policy.queueTimeout > 0 {
|
|
g.policy.queueTimeout = policy.queueTimeout
|
|
g.policy.queueTimeoutSet = true
|
|
} else if policy.queueTimeoutSet {
|
|
g.policy.queueTimeout = policy.queueTimeout
|
|
g.policy.queueTimeoutSet = true
|
|
}
|
|
}
|
|
|
|
// findAvailableNodeLocked returns the available candidate with the lowest
|
|
// absolute in_flight count. Full candidates are skipped, so a provider that is
|
|
// at capacity is deferred behind providers that still have room. When in_flight
|
|
// is equal, the candidate with the lower priority value wins. When both
|
|
// in_flight and priority are equal, all tied candidates remain in best[] so
|
|
// that the rotation logic in the caller picks the next slot after the last
|
|
// selected one.
|
|
func (m *modelQueueManager) findAvailableNodeLocked(group *modelQueueGroup, candidates []candidateNode, long bool) *candidateNode {
|
|
var best []candidateNode
|
|
bestInflight := 0
|
|
bestPriority := 0
|
|
for i := range candidates {
|
|
c := candidates[i]
|
|
if c.capacity <= 0 {
|
|
continue
|
|
}
|
|
slot := c.slotKey()
|
|
inflight := group.inflight[slot]
|
|
if inflight >= c.capacity {
|
|
continue
|
|
}
|
|
// Long-context gate: a long request is excluded from a provider whose
|
|
// long slots are full. Providers with long_context_capacity == 0 declare
|
|
// no dedicated long-slot limit and are not gated here. Normal requests
|
|
// (long == false) never consult long slot state.
|
|
if long && c.longContextCapacity > 0 && group.longInflight[slot] >= c.longContextCapacity {
|
|
continue
|
|
}
|
|
switch {
|
|
case len(best) == 0 || inflight < bestInflight:
|
|
best = []candidateNode{c}
|
|
bestInflight = inflight
|
|
bestPriority = c.priority
|
|
case inflight == bestInflight:
|
|
if c.priority < bestPriority {
|
|
best = []candidateNode{c}
|
|
bestPriority = c.priority
|
|
} else if c.priority == bestPriority {
|
|
best = append(best, c)
|
|
}
|
|
}
|
|
}
|
|
if len(best) == 0 {
|
|
return nil
|
|
}
|
|
sort.Slice(best, func(i, j int) bool {
|
|
return candidateLess(&best[i], &best[j])
|
|
})
|
|
idx := 0
|
|
if group.lastSelectedSlot != "" {
|
|
for i := range best {
|
|
if best[i].slotKey() == group.lastSelectedSlot {
|
|
idx = (i + 1) % len(best)
|
|
break
|
|
}
|
|
}
|
|
}
|
|
return &best[idx]
|
|
}
|
|
|
|
// candidateLess provides a deterministic ordering for equal-inflight/priority rotation:
|
|
// providerID first, then nodeID.
|
|
func candidateLess(a, b *candidateNode) bool {
|
|
if b == nil {
|
|
return true
|
|
}
|
|
if a.providerID != b.providerID {
|
|
return a.providerID < b.providerID
|
|
}
|
|
return a.entry.NodeID < b.entry.NodeID
|
|
}
|
|
|
|
// admit selects an available candidate for the given model group, or queues
|
|
// the request until a slot opens. Blocks until a candidate is assigned, the
|
|
// queue timeout expires, or ctx is cancelled. Returns the selected candidateNode
|
|
// so callers can rewrite the dispatch target for provider-pool requests.
|
|
// The optional long flag marks the request as long-context so it is gated on
|
|
// provider long slots and reserves a long slot on dispatch.
|
|
func (m *modelQueueManager) admit(ctx context.Context, groupKey, adapter, target string, candidates []candidateNode, policy groupPolicy, longFlag ...bool) (*candidateNode, error) {
|
|
candidate, _, err := m.admitWithReason(ctx, groupKey, adapter, target, candidates, policy, longFlag...)
|
|
return candidate, err
|
|
}
|
|
|
|
func (m *modelQueueManager) admitWithReason(ctx context.Context, groupKey, adapter, target string, candidates []candidateNode, policy groupPolicy, longFlag ...bool) (*candidateNode, string, error) {
|
|
long := len(longFlag) > 0 && longFlag[0]
|
|
m.mu.Lock()
|
|
|
|
group := m.getOrCreateGroupLocked(groupKey, policy)
|
|
// Refresh the policy on an already-existing group so capacity/queue changes
|
|
// applied via config refresh (and recomputed per-request from the live store)
|
|
// take effect on this and subsequent admissions.
|
|
m.updateGroupPolicyLocked(group, policy)
|
|
if group.adapter == "" {
|
|
group.adapter = adapter
|
|
}
|
|
if group.target == "" {
|
|
group.target = target
|
|
}
|
|
|
|
candidate := m.findAvailableNodeLocked(group, candidates, long)
|
|
if candidate != nil {
|
|
slot := candidate.slotKey()
|
|
group.inflight[slot]++
|
|
if long && candidate.longContextCapacity > 0 {
|
|
group.longInflight[slot]++
|
|
}
|
|
group.lastSelectedSlot = slot
|
|
m.mu.Unlock()
|
|
return candidate, "dispatched", nil
|
|
}
|
|
|
|
if len(group.queue) >= group.policy.maxQueue {
|
|
m.mu.Unlock()
|
|
return nil, "", fmt.Errorf("model group %q: %w", groupKey, errQueueFull)
|
|
}
|
|
|
|
// Determine queue reason.
|
|
reason := "capacity_full"
|
|
if long {
|
|
longBlocked := false
|
|
for _, c := range candidates {
|
|
if c.capacity <= 0 {
|
|
continue
|
|
}
|
|
slot := c.slotKey()
|
|
inflight := group.inflight[slot]
|
|
if inflight < c.capacity {
|
|
if c.longContextCapacity > 0 && group.longInflight[slot] >= c.longContextCapacity {
|
|
longBlocked = true
|
|
break
|
|
}
|
|
}
|
|
}
|
|
if longBlocked {
|
|
reason = "long_context_capacity_full"
|
|
}
|
|
}
|
|
|
|
var deadline time.Time
|
|
if group.policy.queueTimeout > 0 {
|
|
deadline = time.Now().Add(group.policy.queueTimeout)
|
|
}
|
|
item := &queueItem{
|
|
candidates: candidates,
|
|
waitCh: make(chan admitResult, 1),
|
|
deadline: deadline,
|
|
long: long,
|
|
reason: reason,
|
|
}
|
|
group.queue = append(group.queue, item)
|
|
timeout := group.policy.queueTimeout
|
|
m.mu.Unlock()
|
|
|
|
if timeout <= 0 {
|
|
select {
|
|
case res := <-item.waitCh:
|
|
return res.candidate, item.reason, res.err
|
|
case <-ctx.Done():
|
|
m.mu.Lock()
|
|
m.removeItemLocked(groupKey, item)
|
|
m.mu.Unlock()
|
|
// Handle race: dispatch may have sent to waitCh just before cancellation.
|
|
select {
|
|
case res := <-item.waitCh:
|
|
if res.candidate != nil {
|
|
slot := res.candidate.slotKey()
|
|
var nodeID, providerID string
|
|
if colonIdx := findLastColon(slot); colonIdx > 0 {
|
|
nodeID = slot[:colonIdx]
|
|
providerID = slot[colonIdx+1:]
|
|
} else {
|
|
nodeID = slot
|
|
}
|
|
longSlot := item.long && res.candidate.longContextCapacity > 0
|
|
m.releaseSlotWithLong(groupKey, nodeID, providerID, longSlot)
|
|
}
|
|
default:
|
|
}
|
|
return nil, "", ctx.Err()
|
|
}
|
|
}
|
|
|
|
timer := time.NewTimer(timeout)
|
|
defer timer.Stop()
|
|
|
|
select {
|
|
case res := <-item.waitCh:
|
|
return res.candidate, item.reason, res.err
|
|
case <-timer.C:
|
|
m.mu.Lock()
|
|
m.removeItemLocked(groupKey, item)
|
|
m.mu.Unlock()
|
|
// Handle race: dispatch may have sent to waitCh just before timeout fired.
|
|
select {
|
|
case res := <-item.waitCh:
|
|
if res.candidate != nil {
|
|
// Use public releaseSlot (which acquires m.mu) because we are
|
|
// outside the lock here. This avoids shared-state mutation
|
|
// without mutex that the original releaseSlotByCandidate call
|
|
// would cause.
|
|
slot := res.candidate.slotKey()
|
|
var nodeID, providerID string
|
|
if colonIdx := findLastColon(slot); colonIdx > 0 {
|
|
nodeID = slot[:colonIdx]
|
|
providerID = slot[colonIdx+1:]
|
|
} else {
|
|
nodeID = slot
|
|
}
|
|
longSlot := item.long && res.candidate.longContextCapacity > 0
|
|
m.releaseSlotWithLong(groupKey, nodeID, providerID, longSlot)
|
|
}
|
|
default:
|
|
}
|
|
return nil, "", fmt.Errorf("model group %q: %w", groupKey, errQueueTimeout)
|
|
case <-ctx.Done():
|
|
m.mu.Lock()
|
|
m.removeItemLocked(groupKey, item)
|
|
m.mu.Unlock()
|
|
// Handle race: dispatch may have sent to waitCh just before cancellation.
|
|
select {
|
|
case res := <-item.waitCh:
|
|
if res.candidate != nil {
|
|
// Same fix as timeout path: use public releaseSlot with
|
|
// explicit mutex to avoid concurrent map/slice mutation.
|
|
slot := res.candidate.slotKey()
|
|
var nodeID, providerID string
|
|
if colonIdx := findLastColon(slot); colonIdx > 0 {
|
|
nodeID = slot[:colonIdx]
|
|
providerID = slot[colonIdx+1:]
|
|
} else {
|
|
nodeID = slot
|
|
}
|
|
longSlot := item.long && res.candidate.longContextCapacity > 0
|
|
m.releaseSlotWithLong(groupKey, nodeID, providerID, longSlot)
|
|
}
|
|
default:
|
|
}
|
|
return nil, "", ctx.Err()
|
|
}
|
|
}
|
|
|
|
// trackInflight records a dispatched run so release events can find it.
|
|
// providerID is non-empty for provider-pool dispatches and is used to
|
|
// attribute in-flight counts to the correct provider snapshot. The optional
|
|
// longSlot flag records that a long-context slot was reserved so the terminal
|
|
// release also frees the long slot.
|
|
func (m *modelQueueManager) trackInflight(groupKey, runID, nodeID, providerID string, longSlot ...bool) {
|
|
long := len(longSlot) > 0 && longSlot[0]
|
|
m.mu.Lock()
|
|
m.inflightByRun[runID] = inflightRec{groupKey: groupKey, nodeID: nodeID, providerID: providerID, long: long}
|
|
m.mu.Unlock()
|
|
}
|
|
|
|
// tryDispatchLocked scans the queue from the front and dispatches every item
|
|
// that currently has an available candidate. Unlike a head-only pump, a blocked
|
|
// head does not stop dispatchable items behind it: when the head is a long
|
|
// request whose provider long slots are full, later normal requests that fit
|
|
// ordinary capacity are still dispatched, avoiding head-of-line blocking (SDD
|
|
// D05/D06). FIFO order is preserved among items with the same dispatchability
|
|
// because the scan always picks the earliest dispatchable item first. Expired
|
|
// items are timed out in place regardless of position. Must be called with
|
|
// m.mu held.
|
|
func (m *modelQueueManager) tryDispatchLocked(group *modelQueueGroup) {
|
|
for {
|
|
progressed := false
|
|
for i := 0; i < len(group.queue); i++ {
|
|
item := group.queue[i]
|
|
|
|
if !item.deadline.IsZero() && time.Now().After(item.deadline) {
|
|
group.queue = append(group.queue[:i], group.queue[i+1:]...)
|
|
select {
|
|
case item.waitCh <- admitResult{err: fmt.Errorf("model group %q: %w", group.key, errQueueTimeout)}:
|
|
default:
|
|
}
|
|
progressed = true
|
|
break
|
|
}
|
|
|
|
candidate := m.findAvailableNodeLocked(group, item.candidates, item.long)
|
|
if candidate == nil {
|
|
// No slot for this item under current provider state; skip it so a
|
|
// later dispatchable item is not blocked behind it.
|
|
continue
|
|
}
|
|
|
|
group.queue = append(group.queue[:i], group.queue[i+1:]...)
|
|
slot := candidate.slotKey()
|
|
reserveLong := item.long && candidate.longContextCapacity > 0
|
|
group.inflight[slot]++
|
|
if reserveLong {
|
|
group.longInflight[slot]++
|
|
}
|
|
|
|
select {
|
|
case item.waitCh <- admitResult{candidate: candidate}:
|
|
group.lastSelectedSlot = slot
|
|
default:
|
|
// Caller already timed out or cancelled; release the reserved slot.
|
|
group.inflight[slot]--
|
|
if reserveLong {
|
|
group.longInflight[slot]--
|
|
}
|
|
}
|
|
progressed = true
|
|
break
|
|
}
|
|
if !progressed {
|
|
break
|
|
}
|
|
}
|
|
}
|
|
|
|
func (m *modelQueueManager) removeItemLocked(groupKey string, item *queueItem) {
|
|
group, ok := m.groups[groupKey]
|
|
if !ok {
|
|
return
|
|
}
|
|
for i, it := range group.queue {
|
|
if it == item {
|
|
group.queue = append(group.queue[:i], group.queue[i+1:]...)
|
|
return
|
|
}
|
|
}
|
|
}
|