Serving Models from Go

// batcher.go — a dynamic batching server: bounded queue, max batch size, max wait, backpressure.
package main

import (
	"context"
	"errors"
	"fmt"
	"sort"
	"sync"
	"sync/atomic"
	"time"
)

var ErrBusy = errors.New("server busy")

type request struct {
	input    string
	enqueued time.Time
	reply    chan result // buffered: the worker never blocks on a caller that gave up
	ctx      context.Context
}

type result struct {
	output string
	waited time.Duration
	batch  int
}

// Model stands in for a GPU forward pass: a fixed cost per call plus a small cost per item.
// That shape — expensive to start, cheap to widen — is why batching pays.
type Model struct{ Fixed, PerItem time.Duration }

func (m Model) Run(inputs []string) []string {
	time.Sleep(m.Fixed + time.Duration(len(inputs))*m.PerItem)
	out := make([]string, len(inputs))
	for i, in := range inputs {
		out[i] = "embedding(" + in + ")"
	}
	return out
}

type Batcher struct {
	queue    chan *request
	model    Model
	maxBatch int
	maxWait  time.Duration
	wg       sync.WaitGroup
	batches  atomic.Int64
	items    atomic.Int64
}

func NewBatcher(m Model, queueSize, maxBatch int, maxWait time.Duration) *Batcher {
	b := &Batcher{queue: make(chan *request, queueSize), model: m, maxBatch: maxBatch, maxWait: maxWait}
	b.wg.Add(1)
	go b.loop()
	return b
}

// Submit is what an HTTP handler calls. It never queues without bound: a full queue is an
// immediate error the handler turns into 429 or 503.
func (b *Batcher) Submit(ctx context.Context, input string) (result, error) {
	r := &request{input: input, enqueued: time.Now(), reply: make(chan result, 1), ctx: ctx}
	select {
	case b.queue <- r:
	default:
		return result{}, ErrBusy
	}
	select {
	case res := <-r.reply:
		return res, nil
	case <-ctx.Done():
		return result{}, ctx.Err()
	}
}

// loop owns the batch. Collect until the batch is full OR the oldest request has waited maxWait.
func (b *Batcher) loop() {
	defer b.wg.Done()
	for first := range b.queue {
		batch := []*request{first}
		deadline := time.NewTimer(b.maxWait)
	collect:
		for len(batch) < b.maxBatch {
			select {
			case r, ok := <-b.queue:
				if !ok {
					break collect
				}
				batch = append(batch, r)
			case <-deadline.C:
				break collect
			}
		}
		deadline.Stop()

		// Skip requests whose callers already gave up: do not spend the model on them.
		live := batch[:0]
		inputs := make([]string, 0, len(batch))
		for _, r := range batch {
			if r.ctx.Err() == nil {
				live = append(live, r)
				inputs = append(inputs, r.input)
			}
		}
		if len(live) == 0 {
			continue
		}
		start := time.Now()
		outputs := b.model.Run(inputs)
		b.batches.Add(1)
		b.items.Add(int64(len(live)))
		for i, r := range live {
			r.reply <- result{outputs[i], start.Sub(r.enqueued), len(live)}
		}
	}
}

func (b *Batcher) Close() { close(b.queue); b.wg.Wait() }

// load sends requests at a fixed rate (open loop) and reports what callers experienced.
func load(name string, b *Batcher, ratePerSec int, d time.Duration) {
	var mu sync.Mutex
	var lat []time.Duration
	var rejected, timedOut atomic.Int64
	var wg sync.WaitGroup
	interval := time.Second / time.Duration(ratePerSec)
	start := time.Now()
	for i := 0; time.Since(start) < d; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			ctx, cancel := context.WithTimeout(context.Background(), 250*time.Millisecond)
			defer cancel()
			t0 := time.Now()
			_, err := b.Submit(ctx, fmt.Sprint("text", i))
			switch {
			case errors.Is(err, ErrBusy):
				rejected.Add(1)
			case err != nil:
				timedOut.Add(1)
			default:
				mu.Lock()
				lat = append(lat, time.Since(t0))
				mu.Unlock()
			}
		}()
		time.Sleep(time.Until(start.Add(time.Duration(i+1) * interval)))
	}
	wg.Wait()
	elapsed := time.Since(start)
	b.Close()
	sort.Slice(lat, func(i, j int) bool { return lat[i] < lat[j] })
	p := func(q float64) float64 {
		if len(lat) == 0 {
			return 0
		}
		return float64(lat[int(q*float64(len(lat)-1))].Microseconds()) / 1000
	}
	avgBatch := 0.0
	if n := b.batches.Load(); n > 0 {
		avgBatch = float64(b.items.Load()) / float64(n)
	}
	fmt.Printf("%-26s %5d ok/s  batch %4.1f  p50 %5.1f ms  p99 %6.1f ms  rejected %4d  timed out %4d\n",
		name, int(float64(len(lat))/elapsed.Seconds()), avgBatch, p(0.5), p(0.99), rejected.Load(), timedOut.Load())
}

func main() {
	model := Model{Fixed: 8 * time.Millisecond, PerItem: 250 * time.Microsecond}
	const rate, dur = 600, 1500 * time.Millisecond
	fmt.Printf("model: %v per call + %v per item. Offered load: %d requests/s\n\n", model.Fixed, model.PerItem, rate)

	load("no batching (batch=1)", NewBatcher(model, 64, 1, 0), rate, dur)
	load("batch<=8,  wait 2 ms", NewBatcher(model, 64, 8, 2*time.Millisecond), rate, dur)
	load("batch<=32, wait 2 ms", NewBatcher(model, 64, 32, 2*time.Millisecond), rate, dur)
	load("batch<=32, wait 20 ms", NewBatcher(model, 64, 32, 20*time.Millisecond), rate, dur)
	fmt.Println("\nUnbatched, the model manages ~120 calls/s against 600 offered: the queue stays full, the")
	fmt.Println("few requests that succeed wait almost their whole deadline, and the rest are rejected or expire.")
	fmt.Println("Batched, the same model serves the whole load. A longer wait fills batches but adds latency.")
}