diff --git a/docs/admin-api.openapi.json b/docs/admin-api.openapi.json index 74c449b..feca378 100644 --- a/docs/admin-api.openapi.json +++ b/docs/admin-api.openapi.json @@ -4930,7 +4930,7 @@ }, "source": { "type": "string", - "description": "Source is the adapter file (GGUF for llama.cpp), fetched by each node as the model's weights are." + "description": "Source is the adapter: a GGUF file for llama.cpp; for vLLM a PEFT adapter directory (a mirror version ending in \"/\", a path on the node, or a Hugging Face repository id). Fetched by each node as the model's weights are." }, "scale": { "type": "number", @@ -5077,7 +5077,11 @@ "items": { "$ref": "#/components/schemas/catalog.Adapter" }, - "description": "Adapters are LoRA adapters served on this model's replicas, each asked for by its own name as though it were a model. Every replica loads every adapter, applying one only to a request that names it." + "description": "Adapters are LoRA adapters served on this model's replicas, each asked for by its own name as though it were a model. Every replica loads every adapter, applying one only to a request that names it — unless MaxLoadedAdapters is set." + }, + "max_loaded_adapters": { + "type": "integer", + "description": "MaxLoadedAdapters, above zero, has each replica hold at most this many adapters at once, loaded when a request first asks for one and unloaded least recently used first to make room. The router sends a request for an adapter to a replica that already holds it while that one has room to take it. vLLM profiles only: llama-server loads adapters at start and cannot load one afterwards. Zero loads every adapter on every replica at start." }, "fallbacks": { "type": "array", diff --git a/docs/install.md b/docs/install.md index 114315a..cbd4b05 100644 --- a/docs/install.md +++ b/docs/install.md @@ -1028,6 +1028,55 @@ The adapter file must be a GGUF LoRA made for this model — llama.cpp's `convert_lora_to_gguf.py` makes one from a PEFT adapter — and not a model with the adapter already merged in, which is what most "LoRA GGUF" repositories hold. +#### Adapters on vLLM, and loading them on demand + +A vLLM model carries adapters the same way. There an adapter is a PEFT adapter +directory (`adapter_config.json` and `adapter_model.safetensors`): a mirror +version ending in `/`, a path on the nodes, or a Hugging Face repository id, +which vLLM downloads itself. A request names the adapter as its model, as on +llama.cpp; `scale` applies on llama.cpp only. Without anything more, every +replica registers every adapter at start (`--enable-lora --lora-modules`). +vLLM mixes one adapter per batch by default; raise that with `--max-loras` in +the profile's `args`, and `--max-lora-rank` (16 by default) for adapters of a +higher rank. + +With many adapters, have each replica hold only some of them: + +```yaml +- name: qwen2.5-7b-instruct + source: Qwen/Qwen2.5-7B-Instruct + max_loaded_adapters: 4 # per replica; vllm profiles only + adapters: + - name: support-de + source: http://mirror:8081/support-de/v3/ + - name: support-fr + source: someone/support-fr-lora + profiles: [...] +``` + +Replicas then start holding none (vLLM runs with +`VLLM_ALLOW_RUNTIME_LORA_UPDATING=True` and `--max-loras 4`). The first +request for an adapter has the router load it on a replica through vLLM's +`/v1/load_lora_adapter`, and the requests after it go to that replica while +it is not saturated; when it is, the next request goes where load sends it and +the adapter is loaded there too. Requests that arrive while an adapter is being +loaded wait for that load rather than starting another. A replica that holds +`max_loaded_adapters` already unloads the one least recently used, never one a +request is running on. Each node still fetches every adapter when a replica +starts, and links it as `llmgw-lora/` in the directory engines run in, +which is the path the router has vLLM load it from — so adding an adapter +still needs the replicas restarted, as above. An adapter vLLM refuses (not +found, a rank above `--max-lora-rank`, not for this model) fails the request +with vLLM's reason. + +What each replica holds is what the router loaded there, checked against each +replica's `/v1/models` when the router first meets it and every 30 seconds. +`llmgw_adapter_affinity_total` counts where requests found their adapter (hit, +joined a load, loaded it), `llmgw_adapter_loads_total` and +`llmgw_adapter_evictions_total` what that cost. llama.cpp keeps loading every +adapter at start: llama-server can change the scale of an adapter it started +with, not load a new one. + ### Models served outside the cluster A model can be an OpenAI-compatible API elsewhere — a hosted provider, or diff --git a/docs/reference.md b/docs/reference.md index d8a2484..09e6080 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -114,6 +114,7 @@ rather than repeated. | `owner` | `string` | [Model](#model) | | `external` | `*External` | [Model](#model) | | `adapters` | `[]Adapter` | [Model](#model) | +| `max_loaded_adapters` | `int` | [Model](#model) | | `fallbacks` | `[]string` | [Model](#model) | | `canary` | `*Canary` | [Model](#model) | | `eval` | `*EvalGate` | [Model](#model) | @@ -372,8 +373,10 @@ Name is what clients ask for, like a model name. `string` -Source is the adapter file (GGUF for llama.cpp), fetched by each node -as the model's weights are. +Source is the adapter: a GGUF file for llama.cpp; for vLLM a PEFT +adapter directory (a mirror version ending in "/", a path on the node, +or a Hugging Face repository id). Fetched by each node as the model's +weights are. ### `scale` @@ -1078,7 +1081,20 @@ Takes a [External](#external) block. Adapters are LoRA adapters served on this model's replicas, each asked for by its own name as though it were a model. Every replica loads -every adapter, applying one only to a request that names it. +every adapter, applying one only to a request that names it — unless +MaxLoadedAdapters is set. + +### `max_loaded_adapters` + +`int` + +MaxLoadedAdapters, above zero, has each replica hold at most this many +adapters at once, loaded when a request first asks for one and unloaded +least recently used first to make room. The router sends a request for +an adapter to a replica that already holds it while that one has room +to take it. vLLM profiles only: llama-server loads adapters at start +and cannot load one afterwards. Zero loads every adapter on every +replica at start. ### `fallbacks` diff --git a/internal/agent/agent.go b/internal/agent/agent.go index b792358..4ac7192 100644 --- a/internal/agent/agent.go +++ b/internal/agent/agent.go @@ -1053,6 +1053,15 @@ func (a *Agent) startShard(ctx context.Context, m *llmgwv1.StartShard, out chan< } spec.Profile.Adapters[i].Path = path } + // Adapters that load on demand are loaded by the router, by a path that + // has to mean the same file on every node: a link of the adapter's name + // in the directory engines run in. + if spec.Profile.RuntimeLoRA() { + if err := engines.LinkRuntimeAdapters(a.sup.EngineDir(), spec.Profile.Adapters); err != nil { + report(supervisor.Failed, err.Error(), nil) + return + } + } cmds, err := eng.Commands(spec) if err != nil { diff --git a/internal/agent/engines/lora.go b/internal/agent/engines/lora.go new file mode 100644 index 0000000..000e2bf --- /dev/null +++ b/internal/agent/engines/lora.go @@ -0,0 +1,136 @@ +package engines + +import ( + "fmt" + "maps" + "os" + "path/filepath" + "strconv" + "strings" +) + +// RuntimeLoRAEnv is the variable that lets vLLM load and unload adapters +// while it runs, through POST /v1/load_lora_adapter and +// /v1/unload_lora_adapter. Without it those routes do not exist. +// +// It doubles as the mark that a profile's adapters load on demand: the +// catalog sets it for a model with max_loaded_adapters and refuses it set by +// hand anywhere else, so it cannot disagree with what the router does. +const RuntimeLoRAEnv = "VLLM_ALLOW_RUNTIME_LORA_UPDATING" + +// RuntimeLoRADir is where an agent links each adapter it fetched for a +// replica whose adapters load on demand, under the adapter's name, relative +// to the directory engines run in. +// +// Relative on purpose. The router asks the engine to load an adapter by a +// path, and does not know where any node keeps its cache — each node may +// keep it somewhere else, and a replica spanning nodes reads the adapter on +// every one of them. vLLM resolves a relative path that exists against its +// working directory, which on every node is the one the agent starts +// engines in, and the agent links the adapter there before the engine +// starts. So one path, the same everywhere, names the right file on each. +const RuntimeLoRADir = "llmgw-lora" + +// WithRuntimeLoRA has a vLLM profile with adapters start holding none of +// them, ready to load up to limit at a time on demand: the engine is allowed +// runtime loading, and --max-loras — how many adapters one batch may mix — +// is the cap, unless the operator's arguments already say. Anything else is +// returned as it is. +func (p Profile) WithRuntimeLoRA(limit int) Profile { + if limit <= 0 || p.Engine != "vllm" || len(p.Adapters) == 0 { + return p + } + env := maps.Clone(p.Env) + if env == nil { + env = map[string]string{} + } + env[RuntimeLoRAEnv] = "True" + p.Env = env + if !hasFlag(p.Args, "--max-loras") { + p.Args = append([]string{"--max-loras", strconv.Itoa(limit)}, p.Args...) + } + return p +} + +// RuntimeLoRA reports whether a profile's adapters load on demand rather +// than at start. +func (p Profile) RuntimeLoRA() bool { + v, ok := p.Env[RuntimeLoRAEnv] + return ok && v != "" && v != "0" && !strings.EqualFold(v, "false") +} + +// RuntimeLoRAPath is the lora_path the router gives vLLM to load an adapter +// from: the link the agent made, or, for a Hugging Face repository id, the +// id, which vLLM downloads itself and the agent leaves alone. +func RuntimeLoRAPath(name, source string) string { + if IsHubID(source) { + return source + } + return RuntimeLoRADir + "/" + name +} + +// IsHubID reports whether a source is a bare Hugging Face repository id +// ("org/repo") rather than a URL or a path: what the agent passes to an +// engine untouched. +func IsHubID(source string) bool { + switch { + case source == "", + strings.HasPrefix(source, "/"), strings.HasPrefix(source, "./"), strings.HasPrefix(source, "../"), + strings.HasPrefix(source, "~"), + strings.HasPrefix(source, "http://"), strings.HasPrefix(source, "https://"), + strings.HasPrefix(source, "hf:"): + return false + } + return true +} + +// LinkRuntimeAdapters makes RuntimeLoRADir under dir hold a link to each +// fetched adapter, named as the adapter is. A link already there is +// replaced, so a replica started after an adapter's source changed loads +// the new one. Adapters still named by a repository id have nothing local +// to link. +func LinkRuntimeAdapters(dir string, ads []Adapter) error { + if len(ads) == 0 { + return nil + } + root := filepath.Join(dir, RuntimeLoRADir) + if err := os.MkdirAll(root, 0o755); err != nil { + return fmt.Errorf("making the adapters' directory: %w", err) + } + for _, a := range ads { + if IsHubID(a.Path) { + continue + } + if a.Name == "" || strings.ContainsAny(a.Name, `/\`) || a.Name == "." || a.Name == ".." { + return fmt.Errorf("adapter %q: not a name a file can have", a.Name) + } + target, err := filepath.Abs(a.Path) + if err != nil { + return fmt.Errorf("adapter %s: %w", a.Name, err) + } + link := filepath.Join(root, a.Name) + // Made beside the old one and renamed over it, so an engine loading + // the adapter at this moment sees one or the other and never neither. + tmp := link + ".new" + _ = os.Remove(tmp) + if err := os.Symlink(target, tmp); err != nil { + return fmt.Errorf("adapter %s: %w", a.Name, err) + } + if err := os.Rename(tmp, link); err != nil { + _ = os.Remove(tmp) + return fmt.Errorf("adapter %s: %w", a.Name, err) + } + } + return nil +} + +// loraModules is vLLM's --lora-modules for adapters loaded at start: one +// flag followed by every name=path, since vLLM keeps only the last of a +// repeated flag. +func loraModules(ads []Adapter) []string { + out := []string{"--lora-modules"} + for _, a := range ads { + out = append(out, a.Name+"="+a.Path) + } + return out +} diff --git a/internal/agent/engines/lora_test.go b/internal/agent/engines/lora_test.go new file mode 100644 index 0000000..e737607 --- /dev/null +++ b/internal/agent/engines/lora_test.go @@ -0,0 +1,129 @@ +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 +} diff --git a/internal/agent/engines/vllm.go b/internal/agent/engines/vllm.go index 1fff60b..70dcd4d 100644 --- a/internal/agent/engines/vllm.go +++ b/internal/agent/engines/vllm.go @@ -147,6 +147,20 @@ func (v *VLLM) serverCommand(s ShardSpec) (Command, error) { } } + // Adapters loaded at start are registered under their names, and a + // request names one as its model. Adapters loaded on demand start + // unregistered: the router loads each when a request first wants it + // here, by the link the agent made (RuntimeLoRADir). --max-loras is + // vLLM's default of one at start — raise it in the profile's arguments + // to mix more adapters in a batch — and the cap on demand, which + // WithRuntimeLoRA put in the arguments. + if len(p.Adapters) > 0 { + args = append(args, "--enable-lora") + if !p.RuntimeLoRA() { + args = append(args, loraModules(p.Adapters)...) + } + } + // Profile arguments last, so an operator can override anything above. args = append(args, p.Args...) diff --git a/internal/agent/supervisor/supervisor.go b/internal/agent/supervisor/supervisor.go index e0c365e..259cd26 100644 --- a/internal/agent/supervisor/supervisor.go +++ b/internal/agent/supervisor/supervisor.go @@ -870,6 +870,16 @@ func (s *Supervisor) KillOrphan(o Orphan) error { return nil } +// EngineDir is the working directory engines run in: the engines' user's +// home when they run as one, the work directory otherwise. A relative path +// an engine is given resolves against it. +func (s *Supervisor) EngineDir() string { + if s.as != nil { + return s.as.home + } + return s.cfg.WorkDir +} + // dropTo makes a command run as the engines' user, when there is one: its // credential, its home as the working directory, and its caches there. func (s *Supervisor) dropTo(c *exec.Cmd) { diff --git a/internal/cli/model.go b/internal/cli/model.go index 579e6af..e80be9b 100644 --- a/internal/cli/model.go +++ b/internal/cli/model.go @@ -277,10 +277,12 @@ func newAdaptersCmd(client func() *ctlClient) *cobra.Command { Use: "adapters [--add NAME=URL[@SCALE]] [--remove NAME]", Short: "List, add or remove a model's LoRA adapters", Long: "An adapter is asked for by its NAME as though it were a model, and is served\n" + - "by this model's replicas with it applied. URL is the adapter file, usually in\n" + - "the mirror; it must be a GGUF LoRA adapter for this model's architecture,\n" + - "which is checked before it is attached. Replicas load a changed set when\n" + - "restarted: llmgw ctl restart --model MODEL.", + "by this model's replicas with it applied. URL is the adapter, usually in the\n" + + "mirror. For llama.cpp it must be a GGUF LoRA adapter for this model's\n" + + "architecture, which is checked before it is attached; for vLLM it is a PEFT\n" + + "adapter directory (a mirror version ending in /) or a Hugging Face repository\n" + + "id, which vLLM checks when it loads it. SCALE applies on llama.cpp only.\n" + + "Replicas load a changed set when restarted: llmgw ctl restart --model MODEL.", Args: cobra.ExactArgs(1), RunE: func(cmd *cobra.Command, args []string) error { name := args[0] diff --git a/internal/control/adminapi/api.go b/internal/control/adminapi/api.go index 2a85e79..c37f72e 100644 --- a/internal/control/adminapi/api.go +++ b/internal/control/adminapi/api.go @@ -1846,8 +1846,12 @@ func (a *API) setAdapters(w http.ResponseWriter, r *http.Request) { for _, ad := range m.Adapters { have[ad.Source] = true } + // The file is read only where llama.cpp will load it: a vLLM adapter is + // a PEFT directory, which has no header to read, and vLLM itself refuses + // one that does not fit the model when it is loaded. + ggufs := slices.ContainsFunc(m.Profiles, func(p catalog.Profile) bool { return p.Engine == "llamacpp" }) for _, ad := range req.Adapters { - if have[ad.Source] { + if have[ad.Source] || !ggufs { continue // already checked when it was attached } h, err := gguf.KindFromURL(r.Context(), safedial.Client(safedial.ControlPlane, 30*time.Second), ad.Source) diff --git a/internal/control/catalog/catalog.go b/internal/control/catalog/catalog.go index ded29dc..8ea1f55 100644 --- a/internal/control/catalog/catalog.go +++ b/internal/control/catalog/catalog.go @@ -288,9 +288,18 @@ type Model struct { External *External `json:"external,omitempty" koanf:"external"` // Adapters are LoRA adapters served on this model's replicas, each asked // for by its own name as though it were a model. Every replica loads - // every adapter, applying one only to a request that names it. - Adapters []Adapter `json:"adapters,omitempty" koanf:"adapters"` - Fallbacks []string `json:"fallbacks,omitempty" koanf:"fallbacks"` + // every adapter, applying one only to a request that names it — unless + // MaxLoadedAdapters is set. + Adapters []Adapter `json:"adapters,omitempty" koanf:"adapters"` + // MaxLoadedAdapters, above zero, has each replica hold at most this many + // adapters at once, loaded when a request first asks for one and unloaded + // least recently used first to make room. The router sends a request for + // an adapter to a replica that already holds it while that one has room + // to take it. vLLM profiles only: llama-server loads adapters at start + // and cannot load one afterwards. Zero loads every adapter on every + // replica at start. + MaxLoadedAdapters int `json:"max_loaded_adapters,omitempty" koanf:"max_loaded_adapters"` + Fallbacks []string `json:"fallbacks,omitempty" koanf:"fallbacks"` // Canary sends a share of this model's requests to another catalog // model — the same model at new weights — to compare the two before // moving to it (`llmgw ctl model canary`). @@ -871,7 +880,7 @@ func (m Model) EngineProfile(name string) (engines.Profile, bool) { ModelPath: m.Source, ServedModelName: m.Name, Adapters: m.engineAdapters(), - }, true + }.WithRuntimeLoRA(m.MaxLoadedAdapters), true } return engines.Profile{}, false } @@ -1040,8 +1049,10 @@ func (m Model) IsExternal() bool { return m.External != nil } type Adapter struct { // Name is what clients ask for, like a model name. Name string `json:"name" koanf:"name"` - // Source is the adapter file (GGUF for llama.cpp), fetched by each node - // as the model's weights are. + // Source is the adapter: a GGUF file for llama.cpp; for vLLM a PEFT + // adapter directory (a mirror version ending in "/", a path on the node, + // or a Hugging Face repository id). Fetched by each node as the model's + // weights are. Source string `json:"source" koanf:"source"` // Scale is how strongly it is applied; zero takes 1. Scale float64 `json:"scale,omitempty" koanf:"scale"` @@ -1056,6 +1067,20 @@ func (a Adapter) EffectiveScale() float64 { } func validateAdapters(m Model) error { + if m.MaxLoadedAdapters < 0 { + return fmt.Errorf("model %q: max_loaded_adapters cannot be negative", m.Name) + } + if m.MaxLoadedAdapters > 0 { + // Loading on demand is vLLM's runtime LoRA API. llama-server can + // only change the scale of adapters it loaded at start, so a cap + // there would be a cap on nothing. + for _, p := range m.Profiles { + if p.Engine != "vllm" { + return fmt.Errorf("model %q profile %q: max_loaded_adapters needs vllm profiles; "+ + "llama.cpp loads every adapter at start", m.Name, p.Name) + } + } + } if len(m.Adapters) == 0 { return nil } @@ -1064,9 +1089,17 @@ func validateAdapters(m Model) error { } for _, p := range m.Profiles { // llama.cpp loads adapters at start and applies them per request; - // the other engines are not driven that way yet. - if p.Engine != "llamacpp" && p.Engine != "sim" { - return fmt.Errorf("model %q profile %q: adapters are served by llamacpp profiles only", m.Name, p.Name) + // vLLM loads them at start or on demand. The other engines are not + // driven either way yet. + if p.Engine != "llamacpp" && p.Engine != "vllm" && p.Engine != "sim" { + return fmt.Errorf("model %q profile %q: adapters are served by llamacpp and vllm profiles only", m.Name, p.Name) + } + // The variable is what switches vLLM to loading on demand, and the + // router loads adapters only for a model with a cap: set by hand, it + // would start replicas holding no adapter that nothing ever loads. + if _, set := p.Env[engines.RuntimeLoRAEnv]; set && m.MaxLoadedAdapters == 0 { + return fmt.Errorf("model %q profile %q: set max_loaded_adapters rather than %s", + m.Name, p.Name, engines.RuntimeLoRAEnv) } } seen := map[string]bool{} @@ -1080,6 +1113,9 @@ func validateAdapters(m Model) error { return fmt.Errorf("model %q: adapter %q is declared twice", m.Name, a.Name) case a.Scale < 0: return fmt.Errorf("model %q: adapter %q has a negative scale", m.Name, a.Name) + case m.MaxLoadedAdapters > 0 && (strings.ContainsAny(a.Name, `/\`) || a.Name == "." || a.Name == ".."): + // Loaded on demand, the name is a file's name on the node. + return fmt.Errorf("model %q: adapter %q cannot contain a slash when adapters load on demand", m.Name, a.Name) } seen[a.Name] = true } diff --git a/internal/control/catalog/lora_test.go b/internal/control/catalog/lora_test.go new file mode 100644 index 0000000..b93dc28 --- /dev/null +++ b/internal/control/catalog/lora_test.go @@ -0,0 +1,65 @@ +package catalog + +import ( + "slices" + "strings" + "testing" + + "gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines" +) + +func loraModel(engine string, limit int) Model { + m := model("coder", Profile{Name: "p", Engine: engine, GPUClass: "h100-80g", TP: 1}) + m.Adapters = []Adapter{{Name: "coder-sql", Source: "http://mirror/coder-sql/v1/"}} + m.MaxLoadedAdapters = limit + return m +} + +func TestAdaptersLoadOnDemandOnVLLMOnly(t *testing.T) { + if _, err := New([]Model{loraModel("vllm", 4)}); err != nil { + t.Errorf("vLLM with a cap: %v", err) + } + if _, err := New([]Model{loraModel("vllm", 0)}); err != nil { + t.Errorf("vLLM loading every adapter at start: %v", err) + } + _, err := New([]Model{loraModel("llamacpp", 4)}) + if err == nil || !strings.Contains(err.Error(), "max_loaded_adapters needs vllm") { + t.Errorf("llama.cpp with a cap: %v", err) + } + if _, err := New([]Model{loraModel("vllm", -1)}); err == nil { + t.Error("a negative cap was accepted") + } +} + +func TestRuntimeLoadingCannotBeTurnedOnByHand(t *testing.T) { + m := loraModel("vllm", 0) + m.Profiles[0].Env = map[string]string{engines.RuntimeLoRAEnv: "True"} + if _, err := New([]Model{m}); err == nil || !strings.Contains(err.Error(), "max_loaded_adapters") { + t.Errorf("want a pointer to max_loaded_adapters, got %v", err) + } +} + +func TestAnAdapterThatLoadsOnDemandIsNamedLikeAFile(t *testing.T) { + m := loraModel("vllm", 2) + m.Adapters[0].Name = "team/sql" + if _, err := New([]Model{m}); err == nil { + t.Error("an adapter with a slash in its name was accepted") + } + m.MaxLoadedAdapters = 0 + if _, err := New([]Model{m}); err != nil { + t.Errorf("loaded at start, the name names no file: %v", err) + } +} + +func TestTheEngineProfileOfAModelWithACapStartsVLLMReadyToLoad(t *testing.T) { + p, ok := loraModel("vllm", 3).EngineProfile("p") + if !ok { + t.Fatal("no profile") + } + if !p.RuntimeLoRA() || !slices.Equal(p.Args, []string{"--max-loras", "3"}) { + t.Errorf("env %v args %v", p.Env, p.Args) + } + if p, _ := loraModel("vllm", 0).EngineProfile("p"); p.RuntimeLoRA() || len(p.Args) != 0 { + t.Errorf("without a cap: env %v args %v", p.Env, p.Args) + } +} diff --git a/internal/docgen/reference_gen.go b/internal/docgen/reference_gen.go index abd4942..de38b4f 100644 --- a/internal/docgen/reference_gen.go +++ b/internal/docgen/reference_gen.go @@ -6,7 +6,7 @@ package docgen var Reference = []Block{ {Package: "catalog", Name: "Adapter", Title: "Adapter", Doc: "", Fields: []Field{ {Key: "name", Type: "string", Doc: "Name is what clients ask for, like a model name.", Nested: ""}, - {Key: "source", Type: "string", Doc: "Source is the adapter file (GGUF for llama.cpp), fetched by each node\nas the model's weights are.", Nested: ""}, + {Key: "source", Type: "string", Doc: "Source is the adapter: a GGUF file for llama.cpp; for vLLM a PEFT\nadapter directory (a mirror version ending in \"/\", a path on the node,\nor a Hugging Face repository id). Fetched by each node as the model's\nweights are.", Nested: ""}, {Key: "scale", Type: "float64", Doc: "Scale is how strongly it is applied; zero takes 1.", Nested: ""}, }}, {Package: "config", Name: "Admission", Title: "Admission", Doc: "", Fields: []Field{ @@ -137,7 +137,8 @@ var Reference = []Block{ {Key: "class", Type: "string", Doc: "Class is how much this model's capacity is worth against another's,\nwhen there is not enough for both: critical, standard or opportunistic.\nEmpty is standard, or opportunistic for a model loaded on demand. See\nEffectiveClass.", Nested: ""}, {Key: "owner", Type: "string", Doc: "Owner is the part of the tenancy tree whose capacity quota this model\ncounts against, written org:ID, team:ID or project:ID. Empty is shared:\nheld by the cluster, counted against nobody.", Nested: ""}, {Key: "external", Type: "*External", Doc: "External is a model served by someone else — an OpenAI-compatible API\noutside the cluster — rather than by replicas placed here. Such a model\nhas no profiles and nothing is placed for it; the router sends its\nrequests to the provider. Most useful named in another model's\nfallbacks: used only when the cluster has no capacity left. Its\nprompts leave the cluster, so which keys may name it is worth an\nallowlist.", Nested: "catalog.External"}, - {Key: "adapters", Type: "[]Adapter", Doc: "Adapters are LoRA adapters served on this model's replicas, each asked\nfor by its own name as though it were a model. Every replica loads\nevery adapter, applying one only to a request that names it.", Nested: ""}, + {Key: "adapters", Type: "[]Adapter", Doc: "Adapters are LoRA adapters served on this model's replicas, each asked\nfor by its own name as though it were a model. Every replica loads\nevery adapter, applying one only to a request that names it — unless\nMaxLoadedAdapters is set.", Nested: ""}, + {Key: "max_loaded_adapters", Type: "int", Doc: "MaxLoadedAdapters, above zero, has each replica hold at most this many\nadapters at once, loaded when a request first asks for one and unloaded\nleast recently used first to make room. The router sends a request for\nan adapter to a replica that already holds it while that one has room\nto take it. vLLM profiles only: llama-server loads adapters at start\nand cannot load one afterwards. Zero loads every adapter on every\nreplica at start.", Nested: ""}, {Key: "fallbacks", Type: "[]string", Doc: "", Nested: ""}, {Key: "canary", Type: "*Canary", Doc: "Canary sends a share of this model's requests to another catalog\nmodel — the same model at new weights — to compare the two before\nmoving to it (`llmgw ctl model canary`).", Nested: "catalog.Canary"}, {Key: "eval", Type: "*EvalGate", Doc: "Eval is a quality suite new weights must pass before they replace\nthis model's: `llmgw ctl model update` and `model promote` run it\nagainst them first, and stop below its score.", Nested: "catalog.EvalGate"}, diff --git a/internal/router/adapter_test.go b/internal/router/adapter_test.go index fea6f3c..25f243f 100644 --- a/internal/router/adapter_test.go +++ b/internal/router/adapter_test.go @@ -84,9 +84,9 @@ func TestAdaptersShareTheNamespaceOfModels(t *testing.T) { t.Error("an adapter named like another model was accepted") } bad := adapterModel() - bad.Profiles = []catalog.Profile{{Name: "gpu", Engine: "vllm", GPUClass: "h100-80g", TP: 1}} + bad.Profiles = []catalog.Profile{{Name: "gpu", Engine: "sglang", GPUClass: "h100-80g", TP: 1}} if _, err := catalog.New([]catalog.Model{bad}); err == nil { - t.Error("adapters on a vLLM profile were accepted") + t.Error("adapters on an SGLang profile were accepted") } c, err := catalog.New([]catalog.Model{adapterModel()}) if err != nil { diff --git a/internal/router/lora.go b/internal/router/lora.go new file mode 100644 index 0000000..5fae708 --- /dev/null +++ b/internal/router/lora.go @@ -0,0 +1,511 @@ +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")) +} diff --git a/internal/router/lora_test.go b/internal/router/lora_test.go new file mode 100644 index 0000000..a0412b1 --- /dev/null +++ b/internal/router/lora_test.go @@ -0,0 +1,452 @@ +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) + } +} diff --git a/internal/router/metrics.go b/internal/router/metrics.go index cd31eb7..e37beb5 100644 --- a/internal/router/metrics.go +++ b/internal/router/metrics.go @@ -31,6 +31,12 @@ type metrics struct { prefixDecisions *prometheus.CounterVec + // adapterAffinity is where requests for adapters that load on demand + // found them; adapterLoads and adapterEvictions what that cost. + adapterAffinity *prometheus.CounterVec + adapterLoads *prometheus.CounterVec + adapterEvictions *prometheus.CounterVec + // cacheRequests counts what the answer cache did, cacheBytes and // cacheEntries what it holds. cacheRequests *prometheus.CounterVec @@ -49,6 +55,19 @@ func newMetrics() *metrics { Help: "Prefix routing decisions by branch: cache_hit, cache_miss_min_load, " + "load_balance_override, no_tree, no_healthy_worker.", }, []string{"decision"}), + adapterAffinity: prometheus.NewCounterVec(prometheus.CounterOpts{ + Name: "llmgw_adapter_affinity_total", + Help: "Requests for LoRA adapters that load on demand, by where the chosen replica had it: " + + "hit (loaded), joined (being loaded for another request), loaded (loaded for this one).", + }, []string{"model", "result"}), + adapterLoads: prometheus.NewCounterVec(prometheus.CounterOpts{ + Name: "llmgw_adapter_loads_total", + Help: "LoRA adapters the router had a replica load, by outcome: ok or failed.", + }, []string{"model", "outcome"}), + adapterEvictions: prometheus.NewCounterVec(prometheus.CounterOpts{ + Name: "llmgw_adapter_evictions_total", + Help: "LoRA adapters the router had a replica unload to make room, by outcome: ok or failed.", + }, []string{"model", "outcome"}), requests: prometheus.NewCounterVec(prometheus.CounterOpts{ Name: "llmgw_requests_total", Help: "Requests handled, by model, pool and response status.", @@ -123,6 +142,7 @@ func (m *metrics) collectors() []prometheus.Collector { m.requests, m.duration, m.ttft, m.tokens, m.scorerValue, m.filtered, m.considered, m.admissionWait, m.admissionRejected, m.prefixDecisions, + m.adapterAffinity, m.adapterLoads, m.adapterEvictions, m.cacheRequests, m.cacheBytes, m.cacheEntries, }, slices.Concat(m.tenant.collectors(), m.guard.collectors(), m.token.collectors())...) } @@ -135,6 +155,16 @@ func (m *metrics) observePrefix(d prefix.Decision) { m.prefixDecisions.WithLabelValues(string(d)).Inc() } +// observeAdapter counts a load or an eviction the router made. +func (m *metrics) observeAdapter(model, what, outcome string) { + switch what { + case "load": + m.adapterLoads.WithLabelValues(model, outcome).Inc() + case "evict": + m.adapterEvictions.WithLabelValues(model, outcome).Inc() + } +} + func (m *metrics) observeAdmission(band, model string, wait time.Duration) { if band == "" { band = "none" diff --git a/internal/router/proxy/proxy.go b/internal/router/proxy/proxy.go index 1415937..8c4d164 100644 --- a/internal/router/proxy/proxy.go +++ b/internal/router/proxy/proxy.go @@ -168,6 +168,11 @@ type Request struct { // ContentType is the body's, when it is not JSON: an audio upload's // multipart form. ContentType string + // RetryStatus is one more status, besides the overloaded and unavailable + // ones, that is returned to the caller unwritten rather than relayed: a + // 404 for an adapter the router believed loaded and another router had + // unloaded, which a load and a second try fix. Zero adds none. + RetryStatus int // OnEvent, when set, is told the time since the previous streamed event // for each event after the first: about one token's time each. OnEvent func(gap time.Duration) @@ -277,7 +282,7 @@ func (p *Proxy) Do(ctx context.Context, w http.ResponseWriter, r Request) (Outco } defer resp.Body.Close() - if retryableStatus(resp.StatusCode) { + if retryableStatus(resp.StatusCode) || resp.StatusCode == r.RetryStatus { // Drain a little so the connection can be reused, then let the caller // pick somewhere else. Nothing has been written to the client. _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 4<<10)) diff --git a/internal/router/router.go b/internal/router/router.go index d129cf6..ac529ea 100644 --- a/internal/router/router.go +++ b/internal/router/router.go @@ -25,6 +25,7 @@ import ( "gitea.mynodes.xyz/ilker/llmgw/internal/account" "gitea.mynodes.xyz/ilker/llmgw/internal/accounting" + "gitea.mynodes.xyz/ilker/llmgw/internal/agent/engines" "gitea.mynodes.xyz/ilker/llmgw/internal/auth" "gitea.mynodes.xyz/ilker/llmgw/internal/config" "gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog" @@ -134,6 +135,8 @@ type Router struct { // prefix is nil unless approximate prefix routing is turned on. prefix *scheduler.PrefixScorer + // lora knows which replicas hold which adapters that load on demand. + lora *loraTracker // clusterView is nil until the control plane hands it over, and stays nil // in a router-only deployment. @@ -255,6 +258,7 @@ func New(cfg config.Server, log *slog.Logger) (*Router, error) { reg: reg, profiles: profiles, prefix: prefixScorer, + lora: newLoraTracker(log, metrics.observeAdapter), m: metrics, px: proxy.New(proxy.Config{ DialTimeout: cfg.Router.Timeouts.Dial, @@ -440,6 +444,9 @@ func toInfos(eps []config.Endpoint) []endpoint.Info { // It returns immediately. func (r *Router) Start(ctx context.Context) { r.poller.Sync(ctx) + if r.lora != nil { + go r.lora.run(ctx) + } if r.adm != nil { r.adm.Start() } @@ -654,7 +661,8 @@ func (r *Router) storeModels(models []catalog.Model) { // is applied to the request on the way there. for i, a := range entry.Adapters { m[a.Name] = entry - ads[a.Name] = adapterRef{base: entry.Name, index: i, scale: a.EffectiveScale()} + ads[a.Name] = adapterRef{base: entry.Name, index: i, scale: a.EffectiveScale(), + limit: entry.MaxLoadedAdapters, path: engines.RuntimeLoRAPath(a.Name, a.Source)} } } r.models.Store(&m) @@ -678,8 +686,16 @@ type adapterRef struct { base string index int scale float64 + // limit is the model's max_loaded_adapters: above zero, the adapter + // loads on demand, from path. + limit int + path string } +// onDemand reports whether the adapter is loaded when asked for rather than +// at start. +func (a adapterRef) onDemand() bool { return a.limit > 0 } + // adapter finds one by the name a client asked for. func (r *Router) adapter(name string) (adapterRef, bool) { if p := r.adapters.Load(); p != nil { @@ -730,6 +746,14 @@ func (r *Router) SetEndpoints(ctx context.Context, infos []endpoint.Info) { } } + if r.lora != nil { + ids := make(map[string]struct{}, len(infos)) + for _, i := range infos { + ids[i.ID] = struct{}{} + } + r.lora.keep(ids) + } + r.reg.Replace(infos) r.poller.Sync(ctx) } @@ -917,6 +941,15 @@ func (r *Router) dispatchOne(ctx context.Context, w http.ResponseWriter, req *pa MaxTokens: req.MaxTokens, }, } + // An adapter asked for by name runs on its model's replicas — not on a + // fallback to another model, which does not have it. + ad, isAdapter := r.adapter(req.Model) + isAdapter = isAdapter && ad.base == model + onDemand := isAdapter && ad.onDemand() && r.lora != nil + if onDemand { + sctx.Req.Adapter = req.Model + r.lora.syncUnknown(ctx, candidates, model) + } attempts := max(r.cfg.Retries+1, 1) var last dispatchResult @@ -925,6 +958,9 @@ func (r *Router) dispatchOne(ctx context.Context, w http.ResponseWriter, req *pa if len(candidates) == 0 { break } + if onDemand { + sctx.Adapters = r.lora.view(candidates, req.Model, ad.limit) + } _, sspan := tracer().Start(ctx, "schedule", trace.WithAttributes( attribute.String("llmgw.model", model), attribute.Int("llmgw.candidates", len(candidates)))) dec, err := profile.Schedule(sctx, candidates) @@ -945,6 +981,33 @@ func (r *Router) dispatchOne(ctx context.Context, w http.ResponseWriter, req *pa continue } + // The replica loads the adapter first if it does not hold it. + // Counted against its inflight while it does, so the requests + // arriving meanwhile see it busy. + releaseAdapter := func() {} + if onDemand { + how, rel, lerr := r.lora.acquire(ctx, dec.Endpoint.Info, model, req.Model, ad.path, ad.limit) + if lerr != nil { + release(true) + if ctx.Err() != nil { + return dispatchResult{endpointID: dec.Endpoint.ID, err: proxy.ErrClientGone} + } + res := dispatchResult{endpointID: dec.Endpoint.ID, pool: dec.Endpoint.Pool, + err: fmt.Errorf("adapter %s on %s: %w", req.Model, dec.Endpoint.ID, lerr)} + if errors.Is(lerr, errAdapterLoad) { + // Refused for what it is, it would be refused everywhere. + return res + } + r.log.Warn("a replica could not load an adapter, trying another", + "endpoint", dec.Endpoint.ID, "adapter", req.Model, "err", lerr) + last = res + candidates = without(candidates, dec.Endpoint.ID) + continue + } + releaseAdapter = rel + r.m.adapterAffinity.WithLabelValues(r.metricModel(model), how).Inc() + } + uctx, uspan := tracer().Start(ctx, "upstream", trace.WithSpanKind(trace.SpanKindClient), trace.WithAttributes(attribute.String("llmgw.endpoint", dec.Endpoint.ID), attribute.String("url.full", dec.Endpoint.URL+req.Path))) @@ -952,11 +1015,11 @@ func (r *Router) dispatchOne(ctx context.Context, w http.ResponseWriter, req *pa if dec.Endpoint.UpstreamModel != "" && req.ContentType == "" { body = withModel(body, dec.Endpoint.UpstreamModel) } - // An adapter asked for by name, served by its model's replicas: the - // engine is asked for the model, with the adapter applied. Not on a - // fallback to another model, which does not have it. - if a, ok := r.adapter(req.Model); ok && a.base == model { - body = withAdapter(body, model, a.index, a.scale) + // llama-server is asked for the model with the adapter applied to + // this request; vLLM serves each adapter as a model of its own name, + // which is what the client asked for already. + if isAdapter && !onDemand && dec.Endpoint.Engine != "vllm" { + body = withAdapter(body, model, ad.index, ad.scale) } // What computes the answer, said before the proxy writes the headers. // A model can have replicas on cards and on CPUs at once, and the two @@ -978,6 +1041,12 @@ func (r *Router) dispatchOne(ctx context.Context, w http.ResponseWriter, req *pa CaptureReply: req.captureBytes(), ContentType: req.ContentType, } + if onDemand { + // vLLM's answer for a model it does not have: the adapter was + // unloaded behind this router's back, by another router or a + // restart. Not relayed; loaded again and tried again. + preq.RetryStatus = http.StatusNotFound + } if req.Stream { gaps := r.m.token.gap.WithLabelValues(r.metricModel(model)) preq.OnEvent = func(d time.Duration) { gaps.Observe(d.Seconds()) } @@ -985,10 +1054,16 @@ func (r *Router) dispatchOne(ctx context.Context, w http.ResponseWriter, req *pa // A filter per attempt, for the model this attempt is served by. filter := r.outputFilter(uctx, req, model, &preq) out, perr := r.px.Do(uctx, w, preq) + releaseAdapter() uspan.SetAttributes(attribute.Int("http.response.status_code", out.Status)) endSpan(uspan, perr) - // A retryable failure is the endpoint's fault; a client disconnect is not. - release(perr == nil || errors.Is(perr, proxy.ErrClientGone)) + adapterGone := onDemand && !out.Wrote && out.Status == http.StatusNotFound + if adapterGone { + r.lora.forget(dec.Endpoint.ID, req.Model) + } + // A retryable failure is the endpoint's fault; a client disconnect is + // not, and nor is an adapter it was not known to have lost. + release(perr == nil || errors.Is(perr, proxy.ErrClientGone) || adapterGone) res := dispatchResult{endpointID: dec.Endpoint.ID, pool: dec.Endpoint.Pool, outcome: out, err: perr} if filter != nil { @@ -1011,6 +1086,13 @@ func (r *Router) dispatchOne(ctx context.Context, w http.ResponseWriter, req *pa return res } last = res + if adapterGone { + // The replica is fine; the next try may load the adapter on it + // again. + r.log.Info("a replica no longer held an adapter; loading it again", + "endpoint", dec.Endpoint.ID, "adapter", req.Model) + continue + } r.log.Warn("upstream attempt failed, trying another endpoint", "endpoint", dec.Endpoint.ID, "model", model, "err", perr) candidates = without(candidates, dec.Endpoint.ID) diff --git a/internal/router/scheduler/adapters.go b/internal/router/scheduler/adapters.go new file mode 100644 index 0000000..2f14155 --- /dev/null +++ b/internal/router/scheduler/adapters.go @@ -0,0 +1,58 @@ +package scheduler + +import ( + "gitea.mynodes.xyz/ilker/llmgw/internal/router/admission" + "gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint" +) + +// AdapterView is where a LoRA adapter that loads on demand is, at the moment +// of one decision. +type AdapterView struct { + // Holders are the replicas that have the adapter loaded, or are loading + // it now — a request sent to one of those waits for that load rather + // than starting a second. + Holders map[string]bool + // Free are the replicas with room for one more adapter without + // unloading another. + Free map[string]bool +} + +// AdapterAffinity keeps the replicas that already hold the requested +// adapter, while one of them has room for the request; failing that, the +// replicas that can load it without unloading another; failing that, +// everything. Requests for a model's own weights, and for adapters loaded at +// start, are left alone. +// +// A filter rather than a scorer, unlike prefix affinity. A miss here is not +// a slower prefill but a load — seconds, and perhaps another adapter thrown +// out — and the pickers sample: P2C would send a third of the requests for an +// adapter held by one of three replicas somewhere else, each a load. But it +// gives way to load, by the same measure as pool priority: a holder that is +// saturated does not keep its adapter's traffic, and the request goes where +// the other scorers send it, which loads it there. That is how an adapter +// in demand spreads to a second replica, and the only way it does. +func AdapterAffinity(targetWaiting, kvCeiling float64) Filter { + return adapterAffinityFilter{t: admission.Thresholds{TargetWaitingPerReplica: targetWaiting, KVCeiling: kvCeiling}} +} + +type adapterAffinityFilter struct{ t admission.Thresholds } + +func (adapterAffinityFilter) Name() string { return "adapter-affinity" } + +func (f adapterAffinityFilter) Filter(c *Ctx, eps []endpoint.Snapshot) []endpoint.Snapshot { + if c.Adapters == nil || c.Req.Adapter == "" { + return eps + } + for _, set := range []map[string]bool{c.Adapters.Holders, c.Adapters.Free} { + var keep []endpoint.Snapshot + for _, e := range eps { + if set[e.ID] && admission.Saturation([]endpoint.Snapshot{e}, 0, f.t) < 1 { + keep = append(keep, e) + } + } + if len(keep) > 0 { + return keep + } + } + return eps +} diff --git a/internal/router/scheduler/adapters_test.go b/internal/router/scheduler/adapters_test.go new file mode 100644 index 0000000..907d624 --- /dev/null +++ b/internal/router/scheduler/adapters_test.go @@ -0,0 +1,68 @@ +package scheduler + +import ( + "slices" + "testing" + + "gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint" +) + +func ids(eps []endpoint.Snapshot) []string { + var out []string + for _, e := range eps { + out = append(out, e.ID) + } + return out +} + +func adapterCtx(holders, free []string) *Ctx { + c := newCtx(Request{Model: "m", Adapter: "m-sql"}) + c.Adapters = &AdapterView{Holders: map[string]bool{}, Free: map[string]bool{}} + for _, h := range holders { + c.Adapters.Holders[h] = true + } + for _, f := range free { + c.Adapters.Free[f] = true + } + return c +} + +func TestAdapterAffinityKeepsTheHoldersFirstThenTheReplicasWithRoom(t *testing.T) { + f := AdapterAffinity(4, 0.9) + eps := []endpoint.Snapshot{ep("a"), ep("b"), ep("c")} + for _, c := range []struct { + name string + holders, free []string + want []string + }{ + {"held", []string{"b"}, []string{"a", "c"}, []string{"b"}}, + {"held nowhere, room on one", nil, []string{"c"}, []string{"c"}}, + {"no room anywhere", nil, nil, []string{"a", "b", "c"}}, + } { + if got := ids(f.Filter(adapterCtx(c.holders, c.free), eps)); !slices.Equal(got, c.want) { + t.Errorf("%s: kept %v, want %v", c.name, got, c.want) + } + } +} + +func TestASaturatedHolderDoesNotKeepItsAdaptersTraffic(t *testing.T) { + f := AdapterAffinity(4, 0.9) + eps := []endpoint.Snapshot{ep("a", waiting(8)), ep("b"), ep("c", kv(0.95))} + got := ids(f.Filter(adapterCtx([]string{"a"}, []string{"b", "c"}), eps)) + if !slices.Equal(got, []string{"b"}) { + t.Errorf("kept %v; the holder is saturated, and so is the other with room", got) + } +} + +func TestAdapterAffinityLeavesOtherRequestsAlone(t *testing.T) { + f := AdapterAffinity(4, 0.9) + eps := []endpoint.Snapshot{ep("a"), ep("b")} + c := adapterCtx([]string{"a"}, nil) + c.Req.Adapter = "" + if got := f.Filter(c, eps); len(got) != 2 { + t.Errorf("a request for no adapter was filtered: %v", ids(got)) + } + if got := f.Filter(newCtx(Request{Model: "m"}), eps); len(got) != 2 { + t.Errorf("a request with no view was filtered: %v", ids(got)) + } +} diff --git a/internal/router/scheduler/profiles.go b/internal/router/scheduler/profiles.go index 0a24a18..6a6c342 100644 --- a/internal/router/scheduler/profiles.go +++ b/internal/router/scheduler/profiles.go @@ -50,7 +50,7 @@ func Chat(c ProfileConfig) *Profile { c.normalize() return &Profile{ Name: "chat", - Filters: []Filter{BreakerOpen(), PoolPriority(c.TargetWaiting, c.KVCeiling)}, + Filters: []Filter{BreakerOpen(), PoolPriority(c.TargetWaiting, c.KVCeiling), AdapterAffinity(c.TargetWaiting, c.KVCeiling)}, Scorers: c.withPrefix([]WeightedScorer{ {Session(c.SessionLoadFactor), 3.0}, {Queue(c.TargetWaiting), 1.0}, @@ -81,7 +81,7 @@ func Batch(c ProfileConfig) *Profile { c.normalize() return &Profile{ Name: "batch", - Filters: []Filter{BreakerOpen(), PoolPriority(c.TargetWaiting, c.KVCeiling)}, + Filters: []Filter{BreakerOpen(), PoolPriority(c.TargetWaiting, c.KVCeiling), AdapterAffinity(c.TargetWaiting, c.KVCeiling)}, Scorers: []WeightedScorer{ {Queue(c.TargetWaiting), 3.0}, {KV(c.KVCeiling), 1.0}, @@ -99,7 +99,7 @@ func LongContext(c ProfileConfig) *Profile { c.normalize() return &Profile{ Name: "long-context", - Filters: []Filter{BreakerOpen(), PoolPriority(c.TargetWaiting, c.KVCeiling)}, + Filters: []Filter{BreakerOpen(), PoolPriority(c.TargetWaiting, c.KVCeiling), AdapterAffinity(c.TargetWaiting, c.KVCeiling)}, // Prefix matching earns more here than anywhere else: a long prompt is // mostly prefill, and prefill is what a cache hit skips. Scorers: c.withPrefix([]WeightedScorer{ diff --git a/internal/router/scheduler/scheduler.go b/internal/router/scheduler/scheduler.go index 4ab93fe..64541ce 100644 --- a/internal/router/scheduler/scheduler.go +++ b/internal/router/scheduler/scheduler.go @@ -30,12 +30,17 @@ type Request struct { PromptTokens int // MaxTokens is the requested generation length. MaxTokens int + // Adapter is the LoRA adapter asked for, when it is one that loads on + // demand; empty otherwise. + Adapter string } // Ctx is the context of a single scheduling decision. type Ctx struct { Req Request Now time.Time + // Adapters is where Req.Adapter is loaded; nil when it names none. + Adapters *AdapterView } // Filter drops candidates. diff --git a/notes/c24-lora.md b/notes/c24-lora.md new file mode 100644 index 0000000..fd0be42 --- /dev/null +++ b/notes/c24-lora.md @@ -0,0 +1,52 @@ +## CHANGELOG + +### LoRA adapters on vLLM, loaded on demand + +- **vLLM models can carry adapters.** A PEFT adapter directory (a mirror + version ending in `/`, a path on the nodes, or a Hugging Face repository + id) is registered at start with `--enable-lora --lora-modules`, and a + request names it as its model, as on llama.cpp. `scale` applies on + llama.cpp only. SGLang and TensorRT-LLM profiles still refuse adapters. +- **`max_loaded_adapters: N` on a model** has each replica hold at most N of + its adapters. Replicas start holding none (vLLM with + `VLLM_ALLOW_RUNTIME_LORA_UPDATING=True` and `--max-loras N`, unless the + profile's args set their own). The router loads an adapter on a replica + through `/v1/load_lora_adapter` when a request first wants it there, and + unloads the least recently used one no request is running on to make room. + vLLM profiles only; llama.cpp keeps loading every adapter at start, since + llama-server can rescale an adapter it started with but not load another. +- **Adapter affinity:** a request for an adapter goes to a replica that has + it, or is loading it, while that replica is not saturated; otherwise to one + with room for another adapter; otherwise where load sends it — which is how + a busy adapter spreads to a second replica. Requests arriving during a load + wait for it instead of loading the adapter again. +- **Self-correcting:** each replica's `/v1/models` is read when the router + first meets it and every 30 seconds, so adapters loaded before a restart or + by another router are known; a 404 for an adapter the router thought loaded + loads it again and retries. An adapter vLLM refuses (not found, rank above + `--max-lora-rank`) fails the request with vLLM's reason, without trying it + on every replica. +- **Metrics:** `llmgw_adapter_affinity_total{model,result=hit|joined|loaded}`, + `llmgw_adapter_loads_total{model,outcome}`, + `llmgw_adapter_evictions_total{model,outcome}`. +- The admin API and `ctl model adapters` read the file as GGUF only for a + model with llama.cpp profiles. + +## BACKLOG + +**24 LoRA affinity (branch `features-c2`):** built `max_loaded_adapters` +(catalog), vLLM adapter support at start and on demand (`engines/lora.go`, +`vllm.go`), the agent linking fetched adapters as `llmgw-lora/` in the +engines' working directory, the router's adapter tracker (`router/lora.go`: +loads/unloads with the engine key, LRU eviction never of an adapter in use, +one load per adapter per replica, `/v1/models` sync, 404 reload) and the +`adapter-affinity` scheduler filter in the chat, batch and long-context +profiles, with metrics and docs; tested against a fake of vLLM's HTTP +interface. **Left:** never run against a real vLLM — in particular that vLLM +resolves the relative `llmgw-lora/` path against its working directory +on every node (multi-node too) and that unloading an adapter mid-request is +harmless; the agent still fetches every adapter at replica start, so a new +adapter needs a restart (fetching on first load would need an agent command +and a proto change); load timeout and resync interval are constants, not +config; SGLang's runtime LoRA API is not driven; no per-replica "adapters +held" view in the console or admin API; the debug profile has no affinity.