From 2a1e0e6c4c38949b29ea8a923888ca64a8cfaffd Mon Sep 17 00:00:00 2001 From: Ilker Manap Date: Mon, 5 Oct 2026 01:14:38 +0200 Subject: [PATCH] LoRA adapters on vLLM, loaded on demand with replica affinity: max_loaded_adapters A vLLM model can carry PEFT adapters, registered at start, or with max_loaded_adapters held at most N per replica: the router loads an adapter through vLLM's runtime LoRA API when a request first wants it, sends the adapter's requests to a replica that holds it while that one is not saturated, loads each cold adapter once, and unloads the least recently used idle one to make room. llama.cpp keeps loading every adapter at start. Co-Authored-By: Claude Opus 5.5 (1M context) --- docs/admin-api.openapi.json | 8 +- docs/install.md | 49 ++ docs/reference.md | 22 +- internal/agent/agent.go | 9 + internal/agent/engines/lora.go | 136 ++++++ internal/agent/engines/lora_test.go | 129 ++++++ internal/agent/engines/vllm.go | 14 + internal/agent/supervisor/supervisor.go | 10 + internal/cli/model.go | 10 +- internal/control/adminapi/api.go | 6 +- internal/control/catalog/catalog.go | 54 ++- internal/control/catalog/lora_test.go | 65 +++ internal/docgen/reference_gen.go | 5 +- internal/router/adapter_test.go | 4 +- internal/router/lora.go | 511 +++++++++++++++++++++ internal/router/lora_test.go | 452 ++++++++++++++++++ internal/router/metrics.go | 30 ++ internal/router/proxy/proxy.go | 7 +- internal/router/router.go | 98 +++- internal/router/scheduler/adapters.go | 58 +++ internal/router/scheduler/adapters_test.go | 68 +++ internal/router/scheduler/profiles.go | 6 +- internal/router/scheduler/scheduler.go | 5 + notes/c24-lora.md | 52 +++ 24 files changed, 1773 insertions(+), 35 deletions(-) create mode 100644 internal/agent/engines/lora.go create mode 100644 internal/agent/engines/lora_test.go create mode 100644 internal/control/catalog/lora_test.go create mode 100644 internal/router/lora.go create mode 100644 internal/router/lora_test.go create mode 100644 internal/router/scheduler/adapters.go create mode 100644 internal/router/scheduler/adapters_test.go create mode 100644 notes/c24-lora.md diff --git a/docs/admin-api.openapi.json b/docs/admin-api.openapi.json index 74c449b6..feca378a 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 114315a7..cbd4b05b 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 d8a24849..09e60806 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 b7923586..4ac71923 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 00000000..000e2bff --- /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 00000000..e7376077 --- /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 1fff60be..70dcd4d2 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 e0c365e8..259cd268 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 579e6af7..e80be9b4 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 2a85e79e..c37f72ec 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 ded29dc4..8ea1f55b 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 00000000..b93dc284 --- /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 abd4942b..de38b4f0 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 fea6f3cc..25f243ff 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 00000000..5fae708f --- /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 00000000..a0412b1b --- /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 cd31eb7b..e37beb5f 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 14159376..8c4d1648 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 d129cf63..ac529eaa 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 00000000..2f141558 --- /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 00000000..907d6240 --- /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 0a24a18b..6a6c3425 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 4ab93feb..64541cee 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 00000000..fd0be429 --- /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.