llmgw/internal/agent/engines/lora.go

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
}