llmgw/internal/router/lora.go

package router

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"slices"
	"strings"
	"sync"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/scheduler"
)

// LoRA adapters that load on demand (a model's max_loaded_adapters).
//
// The router decides which replica holds which adapter, and asks vLLM to
// load and unload them itself, with the engine key it already sends
// requests with. Not the control plane through the agent: a load is on the
// path of the request that wants it, and a round trip through the control
// plane's stream to a node and back would add its latency to every cold
// adapter for nothing — the router reaches the engine directly, the agent
// has already fetched every adapter and linked it where vLLM can find it,
// and the router knows the load on every replica, which is what choosing
// one is about.
//
// What each replica holds is learned from two places: the loads and unloads
// this router made, and the engine's own /v1/models, read when the router
// first meets a replica and every so often after — which corrects for a
// router restart, a replica restart and another router doing the same.

const (
	// adapterLoadTimeout bounds one load. An adapter is tens to hundreds of
	// megabytes read from local disk; one that takes minutes is a replica
	// in trouble, and its requests are better sent elsewhere.
	adapterLoadTimeout = 2 * time.Minute
	// adapterSyncTimeout bounds one read of a replica's /v1/models.
	adapterSyncTimeout = 5 * time.Second
	// adapterSyncEvery is how often what the replicas hold is read back.
	adapterSyncEvery = 30 * time.Second
)

// loraTracker knows which adapters each replica holds and loads more.
type loraTracker struct {
	client *http.Client
	log    *slog.Logger
	now    func() time.Time
	// observe counts loads and evictions, by model; nil counts nothing.
	observe func(model, what, outcome string)

	mu   sync.Mutex
	reps map[string]*replicaAdapters
}

// replicaAdapters is one replica's adapters.
type replicaAdapters struct {
	info  endpoint.Info
	model string
	// synced is true once the engine's own list has been read.
	synced bool
	held   map[string]*heldAdapter
	// loading and unloading are in progress; a request for an adapter being
	// unloaded waits for that before loading it again, so the two calls do
	// not cross on the way to the engine.
	loading   map[string]*adapterLoad
	unloading map[string]chan struct{}
}

type heldAdapter struct {
	lastUsed time.Time
	loadedAt time.Time
	// inUse counts requests running on it here, which make it no candidate
	// for eviction.
	inUse int
}

type adapterLoad struct {
	done chan struct{}
	err  error
}

func newLoraTracker(log *slog.Logger, observe func(model, what, outcome string)) *loraTracker {
	return &loraTracker{
		client:  &http.Client{Timeout: adapterLoadTimeout},
		log:     log,
		now:     time.Now,
		observe: observe,
		reps:    map[string]*replicaAdapters{},
	}
}

// rep is a replica's state, made on first sight. Caller holds mu.
func (t *loraTracker) rep(info endpoint.Info, model string) *replicaAdapters {
	st, ok := t.reps[info.ID]
	if !ok {
		st = &replicaAdapters{
			held: map[string]*heldAdapter{}, loading: map[string]*adapterLoad{},
			unloading: map[string]chan struct{}{},
		}
		t.reps[info.ID] = st
	}
	st.info, st.model = info, model
	return st
}

// view is where an adapter is among the candidates, for the scheduler.
func (t *loraTracker) view(eps []endpoint.Snapshot, adapter string, limit int) *scheduler.AdapterView {
	t.mu.Lock()
	defer t.mu.Unlock()
	v := &scheduler.AdapterView{Holders: map[string]bool{}, Free: map[string]bool{}}
	for _, e := range eps {
		st, ok := t.reps[e.ID]
		if !ok {
			v.Free[e.ID] = true
			continue
		}
		_, held := st.held[adapter]
		_, loading := st.loading[adapter]
		v.Holders[e.ID] = held || loading
		v.Free[e.ID] = len(st.held)+len(st.loading) < limit
	}
	return v
}

// syncUnknown reads the adapters of the candidates the router has not met
// yet, all at once, so the first decision after a restart already knows
// what each replica holds.
func (t *loraTracker) syncUnknown(ctx context.Context, eps []endpoint.Snapshot, model string) {
	var wg sync.WaitGroup
	t.mu.Lock()
	for _, e := range eps {
		if st, ok := t.reps[e.ID]; ok && st.synced {
			continue
		}
		t.rep(e.Info, model)
		wg.Add(1)
		go func(info endpoint.Info) {
			defer wg.Done()
			t.sync(ctx, info)
		}(e.Info)
	}
	t.mu.Unlock()
	wg.Wait()
}

// sync reads a replica's /v1/models and takes it as what the replica holds.
//
// vLLM lists each adapter it has registered as a model whose parent is the
// base model. An adapter loaded while the list was on its way is kept: the
// list is older than the load.
func (t *loraTracker) sync(ctx context.Context, info endpoint.Info) {
	started := t.now()
	names, err := t.list(ctx, info)
	t.mu.Lock()
	defer t.mu.Unlock()
	st, ok := t.reps[info.ID]
	if !ok {
		return
	}
	// Marked even on a failure: the replica is then taken to hold what this
	// router loaded on it, and the next sync tries again. Blocking every
	// request on a replica whose list will not come would be worse.
	st.synced = true
	if err != nil {
		t.log.Debug("could not read a replica's adapters", "endpoint", info.ID, "err", err)
		return
	}
	listed := map[string]bool{}
	for _, n := range names {
		listed[n] = true
		if _, ok := st.held[n]; !ok {
			// Loaded by someone else, or before this router started: least
			// recently used as far as this router knows, so first to go.
			st.held[n] = &heldAdapter{}
		}
	}
	for n, h := range st.held {
		if !listed[n] && h.loadedAt.Before(started) && h.inUse == 0 {
			delete(st.held, n)
		}
	}
}

// list asks a replica which adapters it has registered.
func (t *loraTracker) list(ctx context.Context, info endpoint.Info) ([]string, error) {
	ctx, cancel := context.WithTimeout(ctx, adapterSyncTimeout)
	defer cancel()
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, info.URL+"/v1/models", nil)
	if err != nil {
		return nil, err
	}
	if info.APIKey != "" {
		req.Header.Set("Authorization", "Bearer "+info.APIKey)
	}
	resp, err := t.client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("models: %s", resp.Status)
	}
	var body struct {
		Data []struct {
			ID     string  `json:"id"`
			Parent *string `json:"parent"`
		} `json:"data"`
	}
	if err := json.NewDecoder(io.LimitReader(resp.Body, 4<<20)).Decode(&body); err != nil {
		return nil, fmt.Errorf("models: %w", err)
	}
	var out []string
	for _, m := range body.Data {
		if m.Parent != nil && *m.Parent != "" {
			out = append(out, m.ID)
		}
	}
	return out, nil
}

// resyncAll reads back what every replica the router has met holds, and
// forgets the ones that have gone.
func (t *loraTracker) resyncAll(ctx context.Context) {
	t.mu.Lock()
	infos := make([]endpoint.Info, 0, len(t.reps))
	for _, st := range t.reps {
		infos = append(infos, st.info)
	}
	t.mu.Unlock()
	for _, info := range infos {
		if ctx.Err() != nil {
			return
		}
		t.sync(ctx, info)
	}
}

// run resyncs until ctx ends.
func (t *loraTracker) run(ctx context.Context) {
	tick := time.NewTicker(adapterSyncEvery)
	defer tick.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-tick.C:
			t.resyncAll(ctx)
		}
	}
}

// keep forgets every replica not among ids.
func (t *loraTracker) keep(ids map[string]struct{}) {
	t.mu.Lock()
	defer t.mu.Unlock()
	for id := range t.reps {
		if _, ok := ids[id]; !ok {
			delete(t.reps, id)
		}
	}
}

// forget records that a replica turned out not to hold an adapter after
// all: it answered a request for it with "no such model".
func (t *loraTracker) forget(id, adapter string) {
	t.mu.Lock()
	defer t.mu.Unlock()
	if st, ok := t.reps[id]; ok {
		delete(st.held, adapter)
	}
}

// Affinity results, for the metric and the tests: the replica held the
// adapter, was already loading it for another request, or loaded it for
// this one.
const (
	adapterHit    = "hit"
	adapterJoined = "joined"
	adapterLoaded = "loaded"
)

// errAdapterLoad is a load the engine refused: the adapter itself is wrong
// — not found, not a LoRA for this model, a rank above --max-lora-rank —
// and every other replica would refuse it too.
var errAdapterLoad = errors.New("the adapter could not be loaded")

// acquire makes sure a replica holds an adapter, loading it there if it
// does not, and holds it there until release is called: an adapter a
// request is running on is never the one unloaded to make room.
//
// One load per adapter per replica at a time. A request that finds the
// adapter being loaded waits for that load instead of starting another, and
// the scheduler counts a replica that is loading an adapter as holding it,
// so two requests for one cold adapter load it once — on one replica,
// unless that one is too busy to take the second.
func (t *loraTracker) acquire(ctx context.Context, info endpoint.Info, model, adapter, path string, limit int) (string, func(), error) {
	how := adapterHit
	for {
		t.mu.Lock()
		st := t.rep(info, model)
		if h, ok := st.held[adapter]; ok {
			h.inUse++
			h.lastUsed = t.now()
			t.mu.Unlock()
			return how, t.releaser(h), nil
		}
		if l, ok := st.loading[adapter]; ok {
			t.mu.Unlock()
			select {
			case <-l.done:
			case <-ctx.Done():
				return "", nil, ctx.Err()
			}
			if l.err != nil {
				return "", nil, l.err
			}
			if how == adapterHit {
				how = adapterJoined
			}
			continue
		}
		if ch, ok := st.unloading[adapter]; ok {
			t.mu.Unlock()
			select {
			case <-ch:
			case <-ctx.Done():
				return "", nil, ctx.Err()
			}
			continue
		}

		l := &adapterLoad{done: make(chan struct{})}
		st.loading[adapter] = l
		victims := st.victims(adapter, limit)
		for _, v := range victims {
			delete(st.held, v)
			st.unloading[v] = make(chan struct{})
		}
		t.mu.Unlock()
		// On its own, not on the request: the client going away must not
		// abandon a load other requests are waiting for.
		go t.load(st, info, model, adapter, path, victims, l)
		how = adapterLoaded
	}
}

// victims are the adapters to unload so that one more fits under the
// limit: the least recently used of those no request is running on. When
// every one is busy the replica goes over the limit for now, and the next
// load brings it back — vLLM itself only needs the files, and refusing the
// request would serve no one. Caller holds mu.
func (st *replicaAdapters) victims(adding string, limit int) []string {
	var out []string
	count := len(st.held) + len(st.loading)
	for count > limit {
		victim := ""
		var oldest time.Time
		for n, h := range st.held {
			if n == adding || h.inUse > 0 || slices.Contains(out, n) {
				continue
			}
			if victim == "" || h.lastUsed.Before(oldest) || (h.lastUsed.Equal(oldest) && n < victim) {
				victim, oldest = n, h.lastUsed
			}
		}
		if victim == "" {
			break
		}
		out = append(out, victim)
		count--
	}
	return out
}

// load unloads the victims, then loads the adapter, and tells whoever is
// waiting.
func (t *loraTracker) load(st *replicaAdapters, info endpoint.Info, model, adapter, path string, victims []string, l *adapterLoad) {
	ctx, cancel := context.WithTimeout(context.Background(), adapterLoadTimeout)
	defer cancel()
	for _, v := range victims {
		err := t.call(ctx, info, "/v1/unload_lora_adapter", map[string]string{"lora_name": v})
		if err != nil && !notLoaded(err) {
			t.log.Warn("could not unload an adapter to make room; the replica holds one more than its limit",
				"endpoint", info.ID, "adapter", v, "err", err)
			t.count(model, "evict", "failed")
		} else {
			t.count(model, "evict", "ok")
		}
		t.mu.Lock()
		close(st.unloading[v])
		delete(st.unloading, v)
		t.mu.Unlock()
	}

	err := t.call(ctx, info, "/v1/load_lora_adapter", map[string]string{"lora_name": adapter, "lora_path": path})
	if err != nil && alreadyLoaded(err) {
		err = nil
	}
	now := t.now()
	t.mu.Lock()
	delete(st.loading, adapter)
	if err == nil {
		st.held[adapter] = &heldAdapter{lastUsed: now, loadedAt: now}
	}
	l.err = err
	close(l.done)
	t.mu.Unlock()
	if err != nil {
		t.log.Warn("could not load an adapter", "endpoint", info.ID, "adapter", adapter, "path", path, "err", err)
		t.count(model, "load", "failed")
		return
	}
	t.log.Info("adapter loaded", "endpoint", info.ID, "adapter", adapter, "unloaded", victims)
	t.count(model, "load", "ok")
}

func (t *loraTracker) count(model, what, outcome string) {
	if t.observe != nil {
		t.observe(model, what, outcome)
	}
}

// releaser is the release for one acquire, safe to call more than once.
func (t *loraTracker) releaser(h *heldAdapter) func() {
	var once sync.Once
	return func() {
		once.Do(func() {
			t.mu.Lock()
			h.inUse--
			h.lastUsed = t.now()
			t.mu.Unlock()
		})
	}
}

// engineError is a non-success answer from the engine's adapter routes.
type engineError struct {
	status int
	msg    string
}

func (e *engineError) Error() string { return fmt.Sprintf("engine answered %d: %s", e.status, e.msg) }

// Is makes a refusal of the adapter itself an errAdapterLoad. A 4xx is
// about what was asked; a 5xx, or no answer, is about the replica.
func (e *engineError) Is(target error) bool {
	return target == errAdapterLoad && e.status >= 400 && e.status < 500
}

// call posts to one of vLLM's adapter routes. Both answer 200 with a line of
// text on success, and an OpenAI-style error otherwise.
func (t *loraTracker) call(ctx context.Context, info endpoint.Info, path string, body any) error {
	b, err := json.Marshal(body)
	if err != nil {
		return err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, info.URL+path, bytes.NewReader(b))
	if err != nil {
		return err
	}
	req.Header.Set("Content-Type", "application/json")
	if info.APIKey != "" {
		req.Header.Set("Authorization", "Bearer "+info.APIKey)
	}
	resp, err := t.client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	msg, _ := io.ReadAll(io.LimitReader(resp.Body, 4<<10))
	if resp.StatusCode == http.StatusOK {
		return nil
	}
	var e struct {
		Error struct {
			Message string `json:"message"`
		} `json:"error"`
		Message string `json:"message"`
	}
	text := strings.TrimSpace(string(msg))
	if json.Unmarshal(msg, &e) == nil {
		if e.Error.Message != "" {
			text = e.Error.Message
		} else if e.Message != "" {
			text = e.Message
		}
	}
	return &engineError{status: resp.StatusCode, msg: text}
}

// alreadyLoaded is vLLM refusing a load because the adapter is there: what
// the load was for.
func alreadyLoaded(err error) bool {
	var e *engineError
	return errors.As(err, &e) && e.status == http.StatusBadRequest && strings.Contains(e.msg, "already been loaded")
}

// notLoaded is vLLM refusing an unload because the adapter is not there,
// which is what the unload was for. Newer vLLM says so with a 404, older
// with a 400.
func notLoaded(err error) bool {
	var e *engineError
	return errors.As(err, &e) && (e.status == http.StatusNotFound ||
		e.status == http.StatusBadRequest && strings.Contains(e.msg, "cannot be found"))
}