llmgw/internal/agent/engines/lora_test.go

package engines

import (
	"os"
	"path/filepath"
	"slices"
	"testing"
)

func adapterSpec() ShardSpec {
	s := singleNodeSpec()
	s.Profile.Engine = "vllm"
	s.Profile.Adapters = []Adapter{
		{Name: "sql", Path: "/cache/ab-sql"},
		{Name: "docs", Path: "/cache/cd-docs"},
	}
	return s
}

func TestVLLMLoadsEveryAdapterAtStartUnderOneFlag(t *testing.T) {
	cmd := buildOne(t, &VLLM{}, adapterSpec())
	if !slices.Contains(cmd.Args, "--enable-lora") {
		t.Errorf("LoRA not enabled: %v", cmd.Args)
	}
	// vLLM keeps only the last of a repeated --lora-modules, so both
	// adapters must follow one flag.
	i := slices.Index(cmd.Args, "--lora-modules")
	if i < 0 || !slices.Equal(cmd.Args[i+1:i+3], []string{"sql=/cache/ab-sql", "docs=/cache/cd-docs"}) {
		t.Errorf("adapters not registered at start: %v", cmd.Args)
	}
	if n := countFlag(cmd.Args, "--lora-modules"); n != 1 {
		t.Errorf("--lora-modules given %d times", n)
	}
	if _, set := envMap(cmd.Env)[RuntimeLoRAEnv]; set {
		t.Error("runtime loading was allowed for adapters loaded at start")
	}
}

func TestVLLMWithAdaptersOnDemandStartsHoldingNone(t *testing.T) {
	s := adapterSpec()
	s.Profile = s.Profile.WithRuntimeLoRA(3)
	cmd := buildOne(t, &VLLM{}, s)
	if slices.Contains(cmd.Args, "--lora-modules") {
		t.Errorf("adapters registered at start: %v", cmd.Args)
	}
	if !slices.Contains(cmd.Args, "--enable-lora") || argValue(cmd.Args, "--max-loras") != "3" {
		t.Errorf("want LoRA enabled with the cap as --max-loras: %v", cmd.Args)
	}
	if envMap(cmd.Env)[RuntimeLoRAEnv] != "True" {
		t.Errorf("the load and unload routes were not turned on: %v", cmd.Env)
	}
}

func TestTheOperatorsMaxLorasWinsOverTheCap(t *testing.T) {
	p := adapterSpec().Profile
	p.Args = []string{"--max-loras", "8"}
	got := p.WithRuntimeLoRA(3)
	if argValue(got.Args, "--max-loras") != "8" || countFlag(got.Args, "--max-loras") != 1 {
		t.Errorf("args %v", got.Args)
	}
	if p.Env != nil {
		t.Error("the profile's own env map was written to")
	}
}

func TestRuntimeLoRALeavesOtherProfilesAlone(t *testing.T) {
	llama := adapterSpec().Profile
	llama.Engine = "llamacpp"
	none := adapterSpec().Profile
	none.Adapters = nil
	for _, p := range []Profile{llama, none} {
		if got := p.WithRuntimeLoRA(2); got.RuntimeLoRA() || len(got.Args) != 0 {
			t.Errorf("%s with %d adapters: %+v", p.Engine, len(p.Adapters), got)
		}
	}
}

func TestTheRouterLoadsAnAdapterByItsLinkOrItsRepository(t *testing.T) {
	for source, want := range map[string]string{
		"http://mirror:8081/sql-lora/v1/":    "llmgw-lora/sql",
		"/srv/adapters/sql":                  "llmgw-lora/sql",
		"hf:org/repo/adapter.gguf":           "llmgw-lora/sql",
		"someone/sql-lora":                   "someone/sql-lora",
		"https://huggingface.co/someone/sql": "llmgw-lora/sql",
	} {
		if got := RuntimeLoRAPath("sql", source); got != want {
			t.Errorf("%s: %s, want %s", source, got, want)
		}
	}
}

func TestAdaptersAreLinkedUnderTheirNamesAndRelinkedWhenTheyChange(t *testing.T) {
	dir := t.TempDir()
	old := filepath.Join(dir, "old-sql")
	cur := filepath.Join(dir, "new-sql")
	for _, d := range []string{old, cur} {
		if err := os.Mkdir(d, 0o755); err != nil {
			t.Fatal(err)
		}
	}
	ads := []Adapter{{Name: "sql", Path: old}, {Name: "hub", Path: "someone/hub-lora"}}
	if err := LinkRuntimeAdapters(dir, ads); err != nil {
		t.Fatal(err)
	}
	ads[0].Path = cur
	if err := LinkRuntimeAdapters(dir, ads); err != nil {
		t.Fatal(err)
	}
	got, err := os.Readlink(filepath.Join(dir, RuntimeLoRADir, "sql"))
	if err != nil || got != cur {
		t.Errorf("link to %q (%v), want %q", got, err, cur)
	}
	if _, err := os.Lstat(filepath.Join(dir, RuntimeLoRADir, "hub")); !os.IsNotExist(err) {
		t.Errorf("a repository id was linked: %v", err)
	}
	if err := LinkRuntimeAdapters(dir, []Adapter{{Name: "../escape", Path: cur}}); err == nil {
		t.Error("an adapter named like a path was linked")
	}
}

func countFlag(args []string, flag string) int {
	n := 0
	for _, a := range args {
		if a == flag {
			n++
		}
	}
	return n
}