llmgw/internal/router/lora_test.go

package router

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"slices"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/prometheus/client_golang/prometheus/testutil"

	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
)

// fakeVLLM answers as vLLM does when started with --enable-lora and
// VLLM_ALLOW_RUNTIME_LORA_UPDATING=True: adapters are loaded and unloaded
// through POST /v1/load_lora_adapter and /v1/unload_lora_adapter, listed
// in /v1/models with the base model as their parent, and asked for by name
// as the model of a request. The shapes follow vLLM's OpenAI server: a line
// of text on success, an error object otherwise; "already been loaded" is
// a 400, an adapter that is not there a 404, and so is a request for a
// model the engine does not have.
type fakeVLLM struct {
	srv  *httptest.Server
	base string
	key  string

	waiting atomic.Int64

	mu      sync.Mutex
	loaded  map[string]string
	loads   []string
	unloads []string
	served  []map[string]any
	// gate, when set, holds every load until it is closed; arrived hears of
	// each load as it comes in.
	gate    chan struct{}
	arrived chan string
	// refuse, when set, is the reason every load is refused with.
	refuse string
}

func newFakeVLLM(t *testing.T, base string, preloaded ...string) *fakeVLLM {
	t.Helper()
	f := &fakeVLLM{base: base, key: "engine-key", loaded: map[string]string{}}
	for _, n := range preloaded {
		f.loaded[n] = "/somewhere/" + n
	}
	authed := func(w http.ResponseWriter, r *http.Request) bool {
		if r.Header.Get("Authorization") != "Bearer "+f.key {
			http.Error(w, `{"error":"Unauthorized"}`, http.StatusUnauthorized)
			return false
		}
		return true
	}
	fail := func(w http.ResponseWriter, code int, typ, msg string) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(code)
		fmt.Fprintf(w, `{"error":{"message":%q,"type":%q,"param":null,"code":%d}}`, msg, typ, code)
	}
	mux := http.NewServeMux()
	mux.HandleFunc("GET /metrics", func(w http.ResponseWriter, _ *http.Request) {
		fmt.Fprintf(w, "vllm:num_requests_running 0\nvllm:num_requests_waiting %d\nvllm:kv_cache_usage_perc 0.1\n", f.waiting.Load())
	})
	mux.HandleFunc("GET /v1/models", func(w http.ResponseWriter, r *http.Request) {
		if !authed(w, r) {
			return
		}
		f.mu.Lock()
		data := []map[string]any{{"id": base, "object": "model", "root": "/models/" + base, "parent": nil, "max_model_len": 4096}}
		for n, p := range f.loaded {
			data = append(data, map[string]any{"id": n, "object": "model", "root": p, "parent": base})
		}
		f.mu.Unlock()
		_ = json.NewEncoder(w).Encode(map[string]any{"object": "list", "data": data})
	})
	mux.HandleFunc("POST /v1/load_lora_adapter", func(w http.ResponseWriter, r *http.Request) {
		if !authed(w, r) {
			return
		}
		var req struct {
			Name string `json:"lora_name"`
			Path string `json:"lora_path"`
		}
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil || req.Name == "" || req.Path == "" {
			fail(w, 400, "BadRequestError", "Both 'lora_name' and 'lora_path' must be provided.")
			return
		}
		f.mu.Lock()
		gate, arrived, refuse := f.gate, f.arrived, f.refuse
		f.loads = append(f.loads, req.Name+"="+req.Path)
		f.mu.Unlock()
		if arrived != nil {
			arrived <- req.Name
		}
		if gate != nil {
			<-gate
		}
		if refuse != "" {
			fail(w, 400, "BadRequestError", refuse)
			return
		}
		f.mu.Lock()
		defer f.mu.Unlock()
		if _, ok := f.loaded[req.Name]; ok {
			fail(w, 400, "BadRequestError", fmt.Sprintf("The lora adapter '%s' has already been loaded.", req.Name))
			return
		}
		f.loaded[req.Name] = req.Path
		fmt.Fprintf(w, "Success: LoRA adapter '%s' added successfully.", req.Name)
	})
	mux.HandleFunc("POST /v1/unload_lora_adapter", func(w http.ResponseWriter, r *http.Request) {
		if !authed(w, r) {
			return
		}
		var req struct {
			Name string `json:"lora_name"`
		}
		_ = json.NewDecoder(r.Body).Decode(&req)
		f.mu.Lock()
		defer f.mu.Unlock()
		f.unloads = append(f.unloads, req.Name)
		if _, ok := f.loaded[req.Name]; !ok {
			fail(w, 404, "NotFoundError", fmt.Sprintf("The lora adapter '%s' cannot be found.", req.Name))
			return
		}
		delete(f.loaded, req.Name)
		fmt.Fprintf(w, "Success: LoRA adapter '%s' removed successfully.", req.Name)
	})
	mux.HandleFunc("POST /v1/chat/completions", func(w http.ResponseWriter, r *http.Request) {
		if !authed(w, r) {
			return
		}
		var body map[string]any
		_ = json.NewDecoder(r.Body).Decode(&body)
		model, _ := body["model"].(string)
		f.mu.Lock()
		f.served = append(f.served, body)
		_, have := f.loaded[model]
		f.mu.Unlock()
		if model != base && !have {
			fail(w, 404, "NotFoundError", fmt.Sprintf("The model `%s` does not exist.", model))
			return
		}
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprintf(w, `{"model":%q,"choices":[{"message":{"content":"hi"}}],"usage":{"prompt_tokens":4,"completion_tokens":2}}`, model)
	})
	f.srv = httptest.NewServer(mux)
	t.Cleanup(f.srv.Close)
	return f
}

func (f *fakeVLLM) endpoint(id string) config.Endpoint {
	return config.Endpoint{ID: id, URL: f.srv.URL, Models: []string{f.base}, Engine: "vllm", APIKey: f.key}
}

func (f *fakeVLLM) snapshot() (loads, unloads []string, loaded []string, served int) {
	f.mu.Lock()
	defer f.mu.Unlock()
	for n := range f.loaded {
		loaded = append(loaded, n)
	}
	slices.Sort(loaded)
	return slices.Clone(f.loads), slices.Clone(f.unloads), loaded, len(f.served)
}

func (f *fakeVLLM) lastServed() map[string]any {
	f.mu.Lock()
	defer f.mu.Unlock()
	if len(f.served) == 0 {
		return nil
	}
	return f.served[len(f.served)-1]
}

// onDemandModel is a vLLM model whose replicas hold at most limit of its
// three adapters.
func onDemandModel(limit int) catalog.Model {
	return catalog.Model{
		Name: "coder", Source: "/models/coder",
		Profiles: []catalog.Profile{{Name: "gpu", Engine: "vllm", GPUClass: "h100-80g", TP: 1}},
		Adapters: []catalog.Adapter{
			{Name: "coder-sql", Source: "http://mirror/coder-sql/v1/"},
			{Name: "coder-docs", Source: "http://mirror/coder-docs/v1/"},
			{Name: "coder-hub", Source: "someone/coder-hub-lora"},
		},
		MaxLoadedAdapters: limit,
	}
}

func loraRouter(t *testing.T, limit int, engines ...*fakeVLLM) *Router {
	t.Helper()
	var eps []config.Endpoint
	for i, e := range engines {
		eps = append(eps, e.endpoint(fmt.Sprintf("r%d", i)))
	}
	cfg := baseConfig(eps...)
	cfg.Models = []catalog.Model{onDemandModel(limit)}
	if _, err := catalog.New(cfg.Models); err != nil {
		t.Fatalf("the test's model is not a valid one: %v", err)
	}
	return newTestRouter(t, cfg)
}

func askFor(t *testing.T, rt *Router, model string) int {
	t.Helper()
	rec := post(t, rt.Handler(), map[string]any{"model": model, "messages": []any{}}, nil)
	if rec.Code != http.StatusOK {
		t.Logf("%s: %d %s", model, rec.Code, rec.Body.String())
	}
	return rec.Code
}

func TestAnAdapterThatLoadsOnDemandIsLoadedByTheFirstRequestAndKeptForTheNext(t *testing.T) {
	e := newFakeVLLM(t, "coder")
	rt := loraRouter(t, 2, e)

	if askFor(t, rt, "coder-sql") != http.StatusOK {
		t.Fatal("the first request for the adapter failed")
	}
	if askFor(t, rt, "coder-sql") != http.StatusOK {
		t.Fatal("the second request for the adapter failed")
	}
	loads, _, loaded, _ := e.snapshot()
	if !slices.Equal(loads, []string{"coder-sql=llmgw-lora/coder-sql"}) {
		t.Errorf("loads %v, want one, by the link the agent makes", loads)
	}
	if !slices.Equal(loaded, []string{"coder-sql"}) {
		t.Errorf("the engine holds %v", loaded)
	}
	// vLLM serves the adapter as a model of its own name; llama.cpp's
	// "lora" field would mean nothing to it.
	body := e.lastServed()
	if body["model"] != "coder-sql" || body["lora"] != nil {
		t.Errorf("the engine was asked for %v with lora %v", body["model"], body["lora"])
	}
	if got := testutil.ToFloat64(rt.m.adapterAffinity.WithLabelValues("coder", adapterLoaded)); got != 1 {
		t.Errorf("loaded: %v", got)
	}
	if got := testutil.ToFloat64(rt.m.adapterAffinity.WithLabelValues("coder", adapterHit)); got != 1 {
		t.Errorf("hits: %v", got)
	}
	if got := testutil.ToFloat64(rt.m.adapterLoads.WithLabelValues("coder", "ok")); got != 1 {
		t.Errorf("load metric: %v", got)
	}

	// A repository id is handed to vLLM as it is, which downloads it.
	if askFor(t, rt, "coder-hub") != http.StatusOK {
		t.Fatal("the adapter from a repository failed")
	}
	if loads, _, _, _ := e.snapshot(); loads[len(loads)-1] != "coder-hub=someone/coder-hub-lora" {
		t.Errorf("loads %v", loads)
	}
	// And the model itself is still the model.
	if askFor(t, rt, "coder") != http.StatusOK || e.lastServed()["model"] != "coder" {
		t.Error("the base model was not served as itself")
	}
}

func TestRequestsForAnAdapterGoToTheReplicaThatHoldsIt(t *testing.T) {
	es := []*fakeVLLM{newFakeVLLM(t, "coder"), newFakeVLLM(t, "coder"), newFakeVLLM(t, "coder")}
	rt := loraRouter(t, 2, es...)

	for range 12 {
		if askFor(t, rt, "coder-docs") != http.StatusOK {
			t.Fatal("a request failed")
		}
	}
	total, holders := 0, 0
	for _, e := range es {
		loads, _, _, served := e.snapshot()
		total += len(loads)
		if served > 0 {
			holders++
			if served != 12 {
				t.Errorf("the holder served %d of 12", served)
			}
		}
	}
	if total != 1 || holders != 1 {
		t.Errorf("%d loads across %d replicas; want one load, and every request on that replica", total, holders)
	}
}

func TestTwoRequestsForAColdAdapterLoadItOnce(t *testing.T) {
	es := []*fakeVLLM{newFakeVLLM(t, "coder"), newFakeVLLM(t, "coder")}
	gate, arrived := make(chan struct{}), make(chan string, 8)
	for _, e := range es {
		e.gate, e.arrived = gate, arrived
	}
	rt := loraRouter(t, 2, es...)

	codes := make(chan int, 3)
	go func() { codes <- askFor(t, rt, "coder-sql") }()
	<-arrived // the first load is under way, and slow
	for range 2 {
		go func() { codes <- askFor(t, rt, "coder-sql") }()
	}
	// Long enough for the others to have scheduled and found the load.
	time.Sleep(100 * time.Millisecond)
	close(gate)
	for range 3 {
		if c := <-codes; c != http.StatusOK {
			t.Errorf("status %d", c)
		}
	}
	total := 0
	for _, e := range es {
		loads, _, _, _ := e.snapshot()
		total += len(loads)
	}
	if total != 1 {
		t.Errorf("%d loads for one cold adapter", total)
	}
	if got := testutil.ToFloat64(rt.m.adapterAffinity.WithLabelValues("coder", adapterJoined)); got != 2 {
		t.Errorf("%v requests joined the load, want 2", got)
	}
}

func TestTheLeastRecentlyUsedAdapterIsUnloadedToMakeRoom(t *testing.T) {
	e := newFakeVLLM(t, "coder")
	rt := loraRouter(t, 2, e)

	for _, m := range []string{"coder-sql", "coder-docs", "coder-sql", "coder-hub"} {
		if askFor(t, rt, m) != http.StatusOK {
			t.Fatalf("%s failed", m)
		}
	}
	_, unloads, loaded, _ := e.snapshot()
	if !slices.Equal(unloads, []string{"coder-docs"}) {
		t.Errorf("unloaded %v; coder-docs was the least recently used", unloads)
	}
	if !slices.Equal(loaded, []string{"coder-hub", "coder-sql"}) {
		t.Errorf("the engine holds %v", loaded)
	}
	if got := testutil.ToFloat64(rt.m.adapterEvictions.WithLabelValues("coder", "ok")); got != 1 {
		t.Errorf("evictions: %v", got)
	}
}

func TestAnAdapterLoadedBeforeTheRouterStartedIsFoundThere(t *testing.T) {
	cold := newFakeVLLM(t, "coder")
	warm := newFakeVLLM(t, "coder", "coder-docs")
	rt := loraRouter(t, 2, cold, warm)

	for range 4 {
		if askFor(t, rt, "coder-docs") != http.StatusOK {
			t.Fatal("a request failed")
		}
	}
	coldLoads, _, _, coldServed := cold.snapshot()
	warmLoads, _, _, warmServed := warm.snapshot()
	if len(coldLoads)+len(warmLoads) != 0 || coldServed != 0 || warmServed != 4 {
		t.Errorf("loads %v %v, served %d %d; want everything on the replica that had it", coldLoads, warmLoads, coldServed, warmServed)
	}
}

func TestAnAdapterUnloadedBehindTheRoutersBackIsLoadedAgain(t *testing.T) {
	e := newFakeVLLM(t, "coder")
	rt := loraRouter(t, 2, e)
	if askFor(t, rt, "coder-sql") != http.StatusOK {
		t.Fatal("the first request failed")
	}
	// Another router, or a restart, takes it away.
	e.mu.Lock()
	delete(e.loaded, "coder-sql")
	e.mu.Unlock()

	if askFor(t, rt, "coder-sql") != http.StatusOK {
		t.Fatal("the request after the adapter went was not served")
	}
	loads, _, loaded, _ := e.snapshot()
	if len(loads) != 2 || !slices.Equal(loaded, []string{"coder-sql"}) {
		t.Errorf("loads %v, holds %v; want it loaded again", loads, loaded)
	}
}

func TestAnAdapterTheEngineRefusesFailsTheRequestWithTheEnginesReason(t *testing.T) {
	es := []*fakeVLLM{newFakeVLLM(t, "coder"), newFakeVLLM(t, "coder")}
	for _, e := range es {
		e.refuse = "LoRA rank 64 is greater than max_lora_rank 16."
	}
	rt := loraRouter(t, 2, es...)

	rec := post(t, rt.Handler(), map[string]any{"model": "coder-sql", "messages": []any{}}, nil)
	if rec.Code != http.StatusBadGateway || !strings.Contains(rec.Body.String(), "max_lora_rank") {
		t.Errorf("%d %s", rec.Code, rec.Body.String())
	}
	total := 0
	for _, e := range es {
		loads, _, _, _ := e.snapshot()
		total += len(loads)
	}
	if total != 1 {
		t.Errorf("%d loads; an adapter refused for what it is is not tried on every replica", total)
	}
	if got := testutil.ToFloat64(rt.m.adapterLoads.WithLabelValues("coder", "failed")); got != 1 {
		t.Errorf("failed loads: %v", got)
	}
}

func TestABusyHolderGivesTheAdapterToAnotherReplica(t *testing.T) {
	busy := newFakeVLLM(t, "coder", "coder-sql")
	idle := newFakeVLLM(t, "coder")
	busy.waiting.Store(40)
	rt := loraRouter(t, 2, busy, idle)
	// The busy one's queue has been read.
	deadline := time.Now().Add(5 * time.Second)
	for {
		s := rt.reg.Snapshots("coder")
		if i := slices.IndexFunc(s, func(s endpoint.Snapshot) bool { return s.ID == "r0" }); i >= 0 && s[i].Stats.Waiting == 40 {
			break
		}
		if time.Now().After(deadline) {
			t.Fatal("the busy replica's metrics were never read")
		}
		time.Sleep(20 * time.Millisecond)
	}

	if askFor(t, rt, "coder-sql") != http.StatusOK {
		t.Fatal("the request failed")
	}
	loads, _, _, served := idle.snapshot()
	if len(loads) != 1 || served != 1 {
		t.Errorf("the idle replica loaded %v and served %d; the busy holder should have given it up", loads, served)
	}
}

func TestAdaptersLoadedAtStartOnVLLMAreAskedForByName(t *testing.T) {
	e := newFakeVLLM(t, "coder", "coder-sql")
	m := onDemandModel(0)
	cfg := baseConfig(e.endpoint("r0"))
	cfg.Models = []catalog.Model{m}
	rt := newTestRouter(t, cfg)

	if askFor(t, rt, "coder-sql") != http.StatusOK {
		t.Fatal("the request failed")
	}
	if body := e.lastServed(); body["model"] != "coder-sql" || body["lora"] != nil {
		t.Errorf("asked for %v with lora %v", body["model"], body["lora"])
	}
	if loads, _, _, _ := e.snapshot(); len(loads) != 0 {
		t.Errorf("loaded %v; these were loaded at start", loads)
	}
}