llama-swap/internal/router/scheduler/serial.go

package scheduler

import (
	"fmt"
	"sort"
	"strconv"
	"sync"
	"time"

	"github.com/mostlygeek/llama-swap/internal/config"
	"github.com/mostlygeek/llama-swap/internal/logmon"
	"github.com/mostlygeek/llama-swap/internal/process"
	"github.com/mostlygeek/llama-swap/internal/shared"
)

// Serial is a strict one-model-at-a-time scheduler for a size-1 resource: at
// most one request runs at any instant, and when the next request targets a
// model other than the one loaded, every other running model is evicted and the
// target is loaded before it runs. A single model occupies memory at a time, at
// the cost of throughput.
//
// Serial ignores group/eviction policy entirely: it always evicts every other
// running model, regardless of how groups are configured. That is what makes
// the single-model guarantee a property of the scheduler rather than of the
// config.
//
// # Dispatch order
//
// Scheduling is non-preemptive — a running job is never interrupted, because a
// generation cannot be stopped mid-sample without discarding the work. Priority
// applies at dispatch time: whenever the slot frees, the highest scoring
// waiting request goes next, where
//
//		score = request priority + swap affinity + aging
//
//	  - request priority is the caller's X-LlamaSwap-Priority (0 by default).
//	    Bands sit 100 apart: +100 interactive, 0 normal, -100 batch.
//	  - swap affinity is a bounded bonus for a request that can run without a
//	    model swap. Bounded well below the band spacing, so it only breaks
//	    near-ties and can never promote batch work above interactive.
//	  - aging is waited_seconds / agingDivisor, and is unbounded on purpose: it
//	    is the only term allowed to cross a band, which is what stops a batch job
//	    from starving behind an endless stream of normal traffic.
//
// Equal scores keep arrival order, so with default priorities and swap affinity
// disabled this degrades to exact FIFO.
//
// The consequence of being non-preemptive is worth stating plainly: an
// interactive request can still wait up to one full job duration. Priority
// alone does not fix that — the complementary lever is client-side, where a
// batch producer submits one unit at a time and re-queues so gaps are frequent.
//
// Like FIFO, every Scheduler method runs on the router's single run-loop
// goroutine, so no internal locking is needed for queue state. QueueStats is
// the exception: it is read from the /metrics handler, so the stats mirror it
// reads is guarded by a mutex.
type Serial struct {
	name    string
	logger  *logmon.Monitor
	effects Effects

	// agingDivisor is how long a request must wait to gain one priority
	// point. Zero disables aging.
	agingDivisor time.Duration
	// swapAffinityBonus is added to a request that needs no model swap. Zero
	// disables the term.
	swapAffinityBonus int

	// now is the clock, injectable so tests can drive aging deterministically.
	now func() time.Time

	// queued holds waiting requests in arrival order. Dispatch picks by score
	// rather than popping the head, but the slice itself is never reordered so
	// index order remains arrival order — which is what makes equal scores
	// break in favour of the earlier request.
	queued []queuedReq

	// active is the one request currently being processed (loading or serving),
	// or nil when idle. phase is meaningful only while active != nil.
	active *HandlerReq
	phase  serialPhase

	mu    sync.Mutex
	stats serialStats
}

// queuedReq is one waiting request plus the arrival time aging is measured from.
type queuedReq struct {
	req      HandlerReq
	enqueued time.Time
}

// serialStats mirrors queue state and dispatch counters for QueueStats. It is
// written on the run loop and read by the metrics handler, both under mu.
type serialStats struct {
	depth            map[string]int
	oldest           map[string]time.Time
	dispatched       map[string]uint64
	agingReorders    uint64
	affinityReorders uint64
}

// serialPhase is the lifecycle stage of the active request.
type serialPhase int

const (
	phaseIdle     serialPhase = iota
	phaseSwapping             // waiting for OnSwapDone for active.Model
	phaseServing              // waiting for OnServeDone for active.Model
)

// NewSerial builds a Serial scheduler. It takes no Swapper: eviction is always
// "stop every other running model", so the group planner is not consulted.
func NewSerial(name string, logger *logmon.Monitor, cfg config.SerialConfig, eff Effects) *Serial {
	return &Serial{
		name:              name,
		logger:            logger,
		effects:           eff,
		agingDivisor:      time.Duration(cfg.GetAgingDivisor()) * time.Second,
		swapAffinityBonus: cfg.GetSwapAffinityBonus(),
		now:               time.Now,
		stats: serialStats{
			depth:      make(map[string]int),
			oldest:     make(map[string]time.Time),
			dispatched: make(map[string]uint64),
		},
	}
}

// OnRequest validates the model and appends the request to the queue, stamping
// its arrival time so aging can be measured, then tries to start the next job.
// Unknown models fail immediately.
func (s *Serial) OnRequest(req HandlerReq) {
	if _, ok := s.effects.ModelState(req.Model); !ok {
		s.logger.Debugf("%s: model %s not handled by this router", s.name, req.Model)
		s.effects.GrantError(req, ErrModelNotFound)
		return
	}
	s.queued = append(s.queued, queuedReq{req: req, enqueued: s.now()})
	s.queueChanged()
	s.startNext()
}

// startNext begins processing the highest scoring waiting request when nothing
// is active. It fast-paths a request whose model is already the sole
// loaded-and-ready process; otherwise it launches a swap that evicts every other
// running model first. The loop skips over requests for models that vanished
// (e.g. a config reload) and requests whose caller disconnected before they
// could be served.
func (s *Serial) startNext() {
	if s.active != nil {
		return // a job is already loading or serving
	}
	for len(s.queued) > 0 {
		idx, score := s.pick()
		q := s.queued[idx]
		s.queued = append(s.queued[:idx], s.queued[idx+1:]...)
		s.queueChanged()

		req := q.req
		state, ok := s.effects.ModelState(req.Model)
		if !ok {
			s.effects.GrantError(req, ErrModelNotFound)
			continue
		}

		waited := s.now().Sub(q.enqueued)
		s.annotate(req, score, waited)

		r := req
		s.active = &r

		evict := s.otherRunning(req.Model)
		if state == process.StateReady && len(evict) == 0 {
			// Already loaded and the only model running — serve immediately.
			s.logger.Debugf("%s: serving model %s (already loaded, priority %d, score %d, waited %s)",
				s.name, req.Model, req.Priority, score, waited.Round(time.Millisecond))
			if s.serve() {
				s.recordDispatch(req.Priority)
				return
			}
			continue // caller gone; pick the next request
		}

		s.logger.Debugf("%s: swapping to model %s (priority %d, score %d, waited %s), evicting %v",
			s.name, req.Model, req.Priority, score, waited.Round(time.Millisecond), evict)
		s.phase = phaseSwapping
		s.recordDispatch(req.Priority)
		s.effects.StartSwap(req.Model, evict)
		return
	}
}

// pick returns the index of the highest scoring queued request and its score.
// It also accounts for whether the aging and swap-affinity terms changed the
// outcome: each is recomputed with that term removed, and a different winner
// means the term reordered this dispatch. The queue is never empty here.
func (s *Serial) pick() (int, int) {
	now := s.now()
	resident, _ := s.noSwapModel()

	idx, score := s.best(now, resident, true, true)
	if len(s.queued) > 1 {
		if other, _ := s.best(now, resident, false, true); other != idx {
			s.countReorder(true)
		}
		if other, _ := s.best(now, resident, true, false); other != idx {
			s.countReorder(false)
		}
	}
	return idx, score
}

// best returns the index and score of the highest scoring queued request, with
// the aging and swap-affinity terms individually switchable so pick can measure
// their effect. Ties keep the earliest arrival because the queue is in arrival
// order and the comparison is strictly greater-than.
func (s *Serial) best(now time.Time, resident string, withAging, withAffinity bool) (int, int) {
	bestIdx, bestScore := 0, 0
	for i, q := range s.queued {
		score := q.req.Priority
		if withAffinity && resident != "" && q.req.Model == resident {
			score += s.swapAffinityBonus
		}
		if withAging {
			score += s.aging(q, now)
		}
		if i == 0 || score > bestScore {
			bestIdx, bestScore = i, score
		}
	}
	return bestIdx, bestScore
}

// aging converts how long a request has waited into priority points. It is
// unbounded: a request that waits long enough eventually outranks anything.
func (s *Serial) aging(q queuedReq, now time.Time) int {
	if s.agingDivisor <= 0 {
		return 0
	}
	waited := now.Sub(q.enqueued)
	if waited <= 0 {
		return 0
	}
	return int(waited / s.agingDivisor)
}

// noSwapModel returns the model that can be served without any swap: the sole
// running process, when it is ready. Anything else — nothing loaded, a process
// still loading, or more than one running — means every request costs a swap,
// so no request earns the affinity bonus.
func (s *Serial) noSwapModel() (string, bool) {
	if s.swapAffinityBonus == 0 {
		return "", false
	}
	running := s.effects.RunningModels()
	if len(running) != 1 {
		return "", false
	}
	for id, state := range running {
		if state == process.StateReady {
			return id, true
		}
	}
	return "", false
}

// annotate records the dispatch decision on the request's context so it reaches
// the activity log. This is how "why did my request take 20 minutes" gets
// answered after the fact.
func (s *Serial) annotate(req HandlerReq, score int, waited time.Duration) {
	if req.Ctx == nil {
		return
	}
	fields := [...]struct{ key, value string }{
		{"serial_priority", strconv.Itoa(req.Priority)},
		{"serial_band", shared.PriorityBand(req.Priority)},
		{"serial_score", strconv.Itoa(score)},
		{"serial_queue_wait_ms", strconv.FormatInt(waited.Milliseconds(), 10)},
	}
	for _, f := range fields {
		if err := shared.SetReqData(req.Ctx, f.key, f.value); err != nil {
			// Every key shares one context, so the first failure means none
			// of them will land.
			s.logger.Debugf("%s: failed to set %s metadata: %v", s.name, f.key, err)
			return
		}
	}
}

// serve hands the active request its tracked handler. It returns true when the
// request is now serving (await OnServeDone); false when the caller had already
// disconnected, in which case active is cleared so the next job can start.
func (s *Serial) serve() bool {
	if s.effects.GrantServe(*s.active, s.active.Model) {
		s.phase = phaseServing
		return true
	}
	s.logger.Debugf("%s: caller for model %s gone before serve", s.name, s.active.Model)
	s.active = nil
	s.phase = phaseIdle
	return false
}

// OnSwapDone fires when the load for the active request completes. On success the
// request is served; on failure its caller receives the error and the queue
// advances. A SwapDone that does not match the active load (e.g. its request was
// unloaded or cancelled mid-load) is ignored.
func (s *Serial) OnSwapDone(ev SwapDone) {
	if s.active == nil || s.phase != phaseSwapping || s.active.Model != ev.ModelID {
		return
	}
	if ev.Err != nil {
		s.logger.Debugf("%s: swap for model %s failed: %v", s.name, ev.ModelID, ev.Err)
		s.effects.GrantError(*s.active, ev.Err)
		s.active = nil
		s.phase = phaseIdle
		s.startNext()
		return
	}
	if !s.serve() {
		s.startNext() // caller vanished while the model loaded; move on
	}
}

// OnServeDone fires when the active request's handler returns. The slot is freed
// and the next queued request begins.
func (s *Serial) OnServeDone(ev ServeDoneEvent) {
	if s.active == nil || s.phase != phaseServing {
		return
	}
	s.active = nil
	s.phase = phaseIdle
	s.startNext()
}

// OnCancel removes a disconnected client's request from the queue. A request that
// is already active is left to finish: if it was loading, OnSwapDone's serve()
// will find the caller gone (GrantServe false) and advance; if it was serving,
// its handler returns normally and reaches OnServeDone.
func (s *Serial) OnCancel(req HandlerReq) {
	if len(s.queued) == 0 {
		return
	}
	kept := s.queued[:0]
	removed := false
	for _, q := range s.queued {
		if q.req.Respond == req.Respond {
			removed = true
			continue
		}
		kept = append(kept, q)
	}
	s.queued = kept
	if removed {
		s.logger.Debugf("%s: cancelled request for model %s pruned from queue", s.name, req.Model)
		s.queueChanged()
	}
}

// OnUnload reconciles state for an unload, stops the targeted processes, and
// advances the queue. It mirrors the FIFO contract: queued requests for unloaded
// models are failed; an active *loading* request for an unloaded model is failed
// (its swap goroutine is left to finish and its SwapDone is then ignored); an
// active *serving* request is left for its handler to end when StopProcesses
// kills the upstream. The Stop is synchronous so callers of Unload can rely on
// the processes being stopped on return.
func (s *Serial) OnUnload(targets []string, timeout time.Duration) {
	unloadErr := fmt.Errorf("%s: model unloaded", s.name)

	targetSet := make(map[string]bool, len(targets))
	for _, id := range targets {
		targetSet[id] = true
	}

	if s.active != nil && s.phase == phaseSwapping && targetSet[s.active.Model] {
		s.effects.GrantError(*s.active, unloadErr)
		s.active = nil
		s.phase = phaseIdle
	}

	if len(s.queued) > 0 {
		kept := s.queued[:0]
		for _, q := range s.queued {
			if targetSet[q.req.Model] {
				s.effects.GrantError(q.req, unloadErr)
				continue
			}
			kept = append(kept, q)
		}
		s.queued = kept
		s.queueChanged()
	}

	s.effects.StopProcesses(timeout, targets)

	// A still-serving active request advances via OnServeDone when its killed
	// handler returns; only start the next job when nothing is active now.
	if s.active == nil {
		s.startNext()
	}
}

// OnShutdown grants err to every request the scheduler still holds: an active
// loading request and all queued requests. A serving request is torn down with
// its process by the baseRouter.
func (s *Serial) OnShutdown(err error) {
	if s.active != nil && s.phase == phaseSwapping {
		s.effects.GrantError(*s.active, err)
		s.active = nil
		s.phase = phaseIdle
	}
	for _, q := range s.queued {
		s.effects.GrantError(q.req, err)
	}
	s.queued = nil
	s.queueChanged()
}

// otherRunning returns every running model except target, sorted for
// deterministic eviction.
func (s *Serial) otherRunning(target string) []string {
	var out []string
	for id := range s.effects.RunningModels() {
		if id != target {
			out = append(out, id)
		}
	}
	sort.Strings(out)
	return out
}

// queueChanged is called after every mutation of s.queued. It refreshes the
// stats mirror and tells each waiter its new position.
func (s *Serial) queueChanged() {
	s.syncStats()
	s.broadcastPositions()
}

// broadcastPositions sends each waiter its 1-indexed place in dispatch order
// rather than arrival order, so a queued caller sees the position priority
// actually earned it. The ranking is a snapshot: aging keeps moving, so a
// position only holds until the next queue change.
func (s *Serial) broadcastPositions() {
	if len(s.queued) == 0 {
		return
	}
	now := s.now()
	resident, _ := s.noSwapModel()

	scores := make([]int, len(s.queued))
	order := make([]int, len(s.queued))
	for i, q := range s.queued {
		score := q.req.Priority
		if resident != "" && q.req.Model == resident {
			score += s.swapAffinityBonus
		}
		scores[i] = score + s.aging(q, now)
		order[i] = i
	}
	// Stable so equal scores keep arrival order, matching best().
	sort.SliceStable(order, func(a, b int) bool { return scores[order[a]] > scores[order[b]] })

	ranked := make([]HandlerReq, len(order))
	for rank, i := range order {
		ranked[rank] = s.queued[i].req
	}
	broadcastQueuePositions(ranked)
}

// syncStats refreshes the per-band depth and oldest-arrival mirror that
// QueueStats reads. Enqueue times are stored rather than durations so the wait
// is computed fresh at read time instead of going stale between queue changes.
func (s *Serial) syncStats() {
	depth := make(map[string]int, len(shared.PriorityBands))
	oldest := make(map[string]time.Time, len(shared.PriorityBands))
	for _, q := range s.queued {
		band := shared.PriorityBand(q.req.Priority)
		depth[band]++
		if t, ok := oldest[band]; !ok || q.enqueued.Before(t) {
			oldest[band] = q.enqueued
		}
	}

	s.mu.Lock()
	defer s.mu.Unlock()
	s.stats.depth = depth
	s.stats.oldest = oldest
}

// recordDispatch counts one request handed off, bucketed by band.
func (s *Serial) recordDispatch(priority int) {
	band := shared.PriorityBand(priority)
	s.mu.Lock()
	defer s.mu.Unlock()
	s.stats.dispatched[band]++
}

// countReorder records that aging (or swap affinity) changed which request a
// dispatch picked.
func (s *Serial) countReorder(aging bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if aging {
		s.stats.agingReorders++
	} else {
		s.stats.affinityReorders++
	}
}

// QueueStats implements StatsReporter. It is safe to call from any goroutine.
func (s *Serial) QueueStats() QueueStats {
	now := s.now()

	s.mu.Lock()
	defer s.mu.Unlock()

	bands := make(map[string]BandStats, len(shared.PriorityBands))
	for _, band := range shared.PriorityBands {
		bs := BandStats{
			Depth:      s.stats.depth[band],
			Dispatched: s.stats.dispatched[band],
		}
		if t, ok := s.stats.oldest[band]; ok {
			if wait := now.Sub(t); wait > 0 {
				bs.OldestWait = wait
			}
		}
		bands[band] = bs
	}
	return QueueStats{
		Bands:            bands,
		AgingReorders:    s.stats.agingReorders,
		AffinityReorders: s.stats.affinityReorders,
	}
}