llmgw/internal/control/modelserve/context.go

package modelserve

import (
	"context"
	"errors"
	"fmt"
	"math"
	"slices"
	"strconv"
	"strings"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/gguf"
	"gitea.mynodes.xyz/ilker/llmgw/internal/modelplan"
)

// A model's context changed after it was served: what one request may use,
// in the three places it lives — the catalog, the engine's arguments, and
// the memory its KV cache takes — together, checked against what the cards
// hold, then its replicas started again on it one by one.

// ContextRequest is the context to change to: tokens one request may use,
// and for llama.cpp how many at once (its pool is the two multiplied).
type ContextRequest struct {
	Tokens int `json:"tokens"`
	Slots  int `json:"slots,omitempty"`
}

// ContextProfile is what one profile needs before and after.
type ContextProfile struct {
	Profile string   `json:"profile"`
	Engine  string   `json:"engine"`
	Before  string   `json:"before,omitempty"`
	After   string   `json:"after,omitempty"`
	Args    []string `json:"args"`
}

// ContextChange is what was changed, and the replicas started again on it.
type ContextChange struct {
	Model     string           `json:"model"`
	Before    int              `json:"before"`
	After     int              `json:"after"`
	MaxOutput int              `json:"max_output"`
	Trained   int              `json:"trained,omitempty"`
	Profiles  []ContextProfile `json:"profiles"`
	Restarted int              `json:"restarted"`
	// Note is what could not be checked, said once.
	Note string `json:"note,omitempty"`
}

// shape is what sizes a context: a token's KV cache, the context the model
// was trained for, and its weights; and its KV heads, which decide how many
// cards tensor parallel can spread it over (0 unknown).
type shape struct {
	kvPerToken int64
	trained    int
	weights    int64
	kvHeads    int
}

// SetContext changes a served model's context.
func (s *Server) SetContext(ctx context.Context, name string, req ContextRequest, by string) (ContextChange, error) {
	if s.deps.Current == nil || s.deps.Models == nil {
		return ContextChange{}, errors.New("this control plane cannot change the catalog")
	}
	var m catalog.Model
	found := false
	for _, c := range s.deps.Current() {
		if c.Name == name {
			m, found = c, true
		}
	}
	switch {
	case !found:
		return ContextChange{}, fmt.Errorf("no model %q in the catalog", name)
	case s.deps.Models.Declared(name):
		return ContextChange{}, fmt.Errorf("model %q is declared in the cluster's settings: change its context there (llmgw ctl config edit)", name)
	case req.Tokens < 256:
		return ContextChange{}, errors.New("a context of at least 256 tokens")
	case req.Slots < 0 || req.Slots > 64:
		return ContextChange{}, errors.New("slots are 1 to 64")
	}
	sh, note := s.shapeOf(ctx, m)
	if sh.trained > 0 && req.Tokens > sh.trained {
		return ContextChange{}, fmt.Errorf("%s was trained for %d tokens; a longer context answers worse, and llama.cpp or vLLM refuses it",
			name, sh.trained)
	}
	out := ContextChange{Model: name, Before: m.ContextLength, After: req.Tokens, Trained: sh.trained, Note: note}
	profiles := slices.Clone(m.Profiles)
	for i, p := range profiles {
		next, change, err := resizeProfile(p, req, sh)
		if err != nil {
			return ContextChange{}, err
		}
		profiles[i] = next
		out.Profiles = append(out.Profiles, change)
	}
	out.MaxOutput = m.MaxOutput
	if m.MaxOutput == 0 || m.MaxOutput == maxOutput(m.ContextLength) {
		out.MaxOutput = maxOutput(req.Tokens)
	}
	out.MaxOutput = min(out.MaxOutput, req.Tokens)
	if _, err := s.deps.Models.SetContext(ctx, name, req.Tokens, out.MaxOutput, profiles, by); err != nil {
		return ContextChange{}, err
	}
	// Started again on it, a new replica before the old one stops.
	var ids []string
	if s.deps.Fleet != nil {
		for _, r := range s.deps.Fleet.Replicas() {
			if r.Model == name {
				ids = append(ids, r.ID)
			}
		}
	}
	if len(ids) > 0 && s.deps.Restart != nil {
		if err := s.deps.Restart(ids...); err != nil {
			return out, fmt.Errorf("the context is changed, and its replicas could not be started again on it: %w", err)
		}
		out.Restarted = len(ids)
	}
	return out, nil
}

// shapeOf reads what sizes a model's context: a GGUF file's header from the
// mirror, else the original weights' config from their repository.
func (s *Server) shapeOf(ctx context.Context, m catalog.Model) (shape, string) {
	var sh shape
	if m.Origin != nil {
		for _, f := range m.Origin.Files {
			sh.weights += f.Bytes
		}
	}
	switch {
	case strings.HasSuffix(strings.ToLower(m.Source), ".gguf"):
		h, err := gguf.FromURL(ctx, nil, m.Source)
		if err != nil {
			return sh, "its header could not be read (" + err.Error() + "), so the memory it needs is not checked"
		}
		sh.kvPerToken, sh.trained = h.KVBytesPerToken(), int(h.ContextLength) //nolint:gosec // a header's context
		sh.kvHeads = int(h.HeadCountKV)                                       //nolint:gosec // a header's head count
	case m.Origin != nil && m.Origin.Repo != "":
		mm, err := modelplan.FetchModel(ctx, m.Origin.Repo)
		if err != nil {
			return sh, "its config could not be read (" + err.Error() + "), so the memory it needs is not checked"
		}
		sh.kvPerToken, sh.trained = mm.KVBytesPerToken(modelplan.BF16), mm.MaxContext
		sh.kvHeads = mm.KVHeads
	default:
		return sh, "where it came from is not known, so the memory it needs is not checked"
	}
	if sh.weights == 0 {
		return sh, "its files' sizes are not known, so the memory it needs is not checked"
	}
	return sh, ""
}

// resizeProfile is a profile with the context asked for: its engine's
// arguments, and its memory where it is on CPUs, refused where its cards
// would not hold it.
func resizeProfile(p catalog.Profile, req ContextRequest, sh shape) (catalog.Profile, ContextProfile, error) {
	change := ContextProfile{Profile: p.Name, Engine: p.Engine}
	var pool int
	switch p.Engine {
	case "llamacpp":
		slots := req.Slots
		if slots == 0 {
			slots = max(argInt(p.Args, "--parallel"), 1)
		}
		pool = req.Tokens * slots
		p.Args = setArg(setArg(p.Args, "--ctx-size", strconv.Itoa(pool)), "--parallel", strconv.Itoa(slots))
	case "vllm":
		pool = req.Tokens
		p.Args = setArg(p.Args, "--max-model-len", strconv.Itoa(pool))
	case "sglang":
		pool = req.Tokens
		p.Args = setArg(p.Args, "--context-length", strconv.Itoa(pool))
	default:
		return p, change, fmt.Errorf("a %s profile (%s) does not take a context this way", p.Engine, p.Name)
	}
	change.Args = p.Args
	if sh.kvPerToken == 0 || sh.weights == 0 {
		return p, change, nil
	}
	if p.CPUOnly() {
		need := sh.weights + int64(pool)*sh.kvPerToken + int64(1<<30) + sh.weights/20
		gb := int(math.Ceil(modelplan.GB(need)))
		change.Before, change.After = fmt.Sprintf("%d GB of RAM", p.MinMemoryGB), fmt.Sprintf("%d GB of RAM", gb)
		p.MinMemoryGB = gb
		return p, change, nil
	}
	cards := max(p.GPUsPerNode, p.TP, 1) * max(p.Nodes, p.PP, 1)
	need := modelplan.CardNeed(p.Engine, sh.weights, int64(pool)*sh.kvPerToken, cards)
	change.After = fmt.Sprintf("%.1f GB a card of %d", need, cards)
	if budget := modelplan.CardBudget(p.GPUClass); budget > 0 {
		change.After += fmt.Sprintf(", of %.0f usable", budget)
		if need > budget {
			return p, change, fmt.Errorf("with %d tokens the %s profile needs %.1f GB on each of its %d %s cards, more than the %.0f GB they hold: "+
				"a shorter context, or the model on more cards", req.Tokens, p.Name, need, cards, p.GPUClass, budget)
		}
	}
	return p, change, nil
}

// setArg sets a flag's value in arguments, or adds it.
func setArg(args []string, flag, value string) []string {
	out := slices.Clone(args)
	for i := 0; i < len(out); i++ {
		if out[i] == flag && i+1 < len(out) {
			out[i+1] = value
			return out
		}
		if v, ok := strings.CutPrefix(out[i], flag+"="); ok && v != "" {
			out[i] = flag + "=" + value
			return out
		}
	}
	return append(out, flag, value)
}

// argInt is a flag's whole-number value, 0 when absent.
func argInt(args []string, flag string) int {
	for i, a := range args {
		if a == flag && i+1 < len(args) {
			n, _ := strconv.Atoi(args[i+1])
			return n
		}
		if v, ok := strings.CutPrefix(a, flag+"="); ok {
			n, _ := strconv.Atoi(v)
			return n
		}
	}
	return 0
}