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).
// MaxOutput is the longest answer, which may be changed with it or alone:
// Tokens 0 with a MaxOutput changes only that, in the catalog, which the
// router reads it from — nothing is started again.
type ContextRequest struct {
	Tokens    int `json:"tokens,omitempty"`
	Slots     int `json:"slots,omitempty"`
	MaxOutput int `json:"max_output,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"`
	// OutputOnly is a change of the longest answer alone, MaxOutputBefore
	// what it was (0 no limit); nothing was started again.
	OutputOnly      bool             `json:"output_only,omitempty"`
	MaxOutputBefore int              `json:"max_output_before,omitempty"`
	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.MaxOutput < 0:
		return ContextChange{}, errors.New("the longest answer is a number of tokens")
	case req.Tokens == 0 && req.MaxOutput > 0:
		return s.setMaxOutput(ctx, m, req.MaxOutput, by)
	case req.Tokens < 256:
		return ContextChange{}, errors.New("a context of at least 256 tokens")
	case req.MaxOutput > req.Tokens:
		return ContextChange{}, fmt.Errorf("an answer of %d tokens does not fit a context of %d", req.MaxOutput, req.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 req.MaxOutput > 0 {
		out.MaxOutput = req.MaxOutput
	}
	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
}

// setMaxOutput changes only the longest answer: the catalog's, which the
// router holds a request's max_tokens to. No engine reads it, so no replica
// is started again.
func (s *Server) setMaxOutput(ctx context.Context, m catalog.Model, maxOut int, by string) (ContextChange, error) {
	if m.ContextLength > 0 && maxOut > m.ContextLength {
		return ContextChange{}, fmt.Errorf("an answer of %d tokens does not fit %s's context of %d: change the context with it",
			maxOut, m.Name, m.ContextLength)
	}
	if _, err := s.deps.Models.SetContext(ctx, m.Name, m.ContextLength, maxOut, m.Profiles, by); err != nil {
		return ContextChange{}, err
	}
	return ContextChange{Model: m.Name, Before: m.ContextLength, After: m.ContextLength, MaxOutput: maxOut,
		OutputOnly: true, MaxOutputBefore: m.MaxOutput}, 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))
	case "strata":
		// Strata's setup takes the context, and writes the engine's start
		// script with it at every start; it sizes its KV cache and expert
		// cache to the card itself, so there is no card check to make here.
		p.Args = setArg(p.Args, "--context", strconv.Itoa(req.Tokens))
		change.Args = p.Args
		return p, change, nil
	case "freetoken":
		// Its experts live in RAM and stream to the card, so the weights
		// are not what has to fit there, and the card check below would
		// refuse what it serves.
		p.Args = setArg(p.Args, "--max-seq-len-override", strconv.Itoa(req.Tokens))
		change.Args = p.Args
		return p, change, nil
	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
}