llmgw/internal/control/modelserve/gpus.go

package modelserve

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

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

// A served model moved to more cards, or fewer: how many each replica
// spreads it over, on one node. The profile changes — tensor parallel for
// vLLM, SGLang and TensorRT-LLM, llama.cpp's layer split for llama.cpp —
// checked against the model's KV heads, what each card would hold, and the
// nodes there are; then its replicas start again on it.

// GPUsRequest is how many cards a replica is to use.
type GPUsRequest struct {
	GPUs int `json:"gpus"`
	// Profile is the one to change, when the model has more than one
	// profile on cards.
	Profile string `json:"profile,omitempty"`
	// InPlace stops each replica before its replacement starts: for when no
	// node has the cards free beside it. It does not answer while the new
	// one loads.
	InPlace bool `json:"in_place,omitempty"`
}

// GPUsChange is what was changed, and the replicas started again on it.
type GPUsChange struct {
	Model   string   `json:"model"`
	Profile string   `json:"profile"`
	Engine  string   `json:"engine"`
	Class   string   `json:"gpu_class"`
	Before  int      `json:"before"`
	After   int      `json:"after"`
	Args    []string `json:"args"`
	// PerCard is what each card holds after, when it could be worked out.
	PerCard   string `json:"per_card,omitempty"`
	Restarted int    `json:"restarted"`
	InPlace   bool   `json:"in_place,omitempty"`
	// Note is what could not be checked, said once.
	Note string `json:"note,omitempty"`
}

// tpFlags are each engine's own spelling of a tensor-parallel width, which
// the agent sets from the profile: a profile's own would contradict it.
var tpFlags = map[string][]string{
	"vllm":   {"--tensor-parallel-size", "-tp"},
	"sglang": {"--tp-size", "--tensor-parallel-size", "--tp"},
	"trtllm": {"--tp_size"},
	// llama.cpp's proportions were for the cards it had.
	"llamacpp": {"--tensor-split", "-ts"},
}

// SetGPUs changes how many cards a served model's replicas use.
func (s *Server) SetGPUs(ctx context.Context, name string, req GPUsRequest, by string) (GPUsChange, error) {
	if s.deps.Current == nil || s.deps.Models == nil {
		return GPUsChange{}, 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 GPUsChange{}, fmt.Errorf("no model %q in the catalog", name)
	case s.deps.Models.Declared(name):
		return GPUsChange{}, fmt.Errorf("model %q is declared in the cluster's settings: change its profile there (llmgw ctl config edit)", name)
	case req.GPUs < 1 || req.GPUs > 64:
		return GPUsChange{}, errors.New("a replica uses 1 to 64 cards")
	}
	i, err := gpuProfile(m, req.Profile)
	if err != nil {
		return GPUsChange{}, err
	}
	p := m.Profiles[i]
	before := max(p.GPUsPerNode, p.TP, 1) * max(p.Nodes, p.PP, 1)
	out := GPUsChange{Model: name, Profile: p.Name, Engine: p.Engine, Class: p.GPUClass, Before: before, After: req.GPUs}
	if before == req.GPUs {
		return out, fmt.Errorf("%s's %s profile already uses %d card(s)", name, p.Name, before)
	}

	sh, note := s.shapeOf(ctx, m)
	next, err := respread(p, req.GPUs, sh)
	if err != nil {
		return out, err
	}
	out.Args = next.Args
	switch kv := int64(contextOf(m, p)) * sh.kvPerToken; {
	case sh.weights == 0 || sh.kvPerToken == 0 || contextOf(m, p) == 0:
		if note == "" {
			note = "its context is not known, so what each card holds is not checked"
		}
	default:
		need := modelplan.CardNeed(p.Engine, sh.weights, kv, req.GPUs)
		out.PerCard = fmt.Sprintf("%.1f GB a card of %d", need, req.GPUs)
		if budget := modelplan.CardBudget(p.GPUClass); budget > 0 {
			out.PerCard += fmt.Sprintf(", of %.0f usable", budget)
			if need > budget {
				return out, fmt.Errorf("on %d %s cards each would hold %.1f GB, more than the %.0f GB they hold: "+
					"more cards, or a shorter context (llmgw ctl model context)", req.GPUs, p.GPUClass, need, budget)
			}
		}
	}
	out.Note = note

	ids := s.replicasOf(name)
	inPlace, err := s.roomFor(p.GPUClass, req.GPUs, name, len(ids) > 0, req.InPlace)
	if err != nil {
		return out, err
	}
	profiles := slices.Clone(m.Profiles)
	profiles[i] = next
	if _, err := s.deps.Models.SetProfiles(ctx, name, profiles, by); err != nil {
		return out, err
	}
	if len(ids) == 0 {
		return out, nil
	}
	restart := s.deps.Restart
	if inPlace {
		restart, out.InPlace = s.deps.RestartInPlace, true
	}
	if restart == nil {
		return out, errors.New("the profile is changed, and this control plane cannot start its replicas again: llmgw ctl restart")
	}
	if err := restart(ids...); err != nil {
		return out, fmt.Errorf("the profile is changed, and its replicas could not be started again on it: %w", err)
	}
	out.Restarted = len(ids)
	return out, nil
}

// gpuProfile is the model's profile on cards to change: the one named, or
// its only one.
func gpuProfile(m catalog.Model, named string) (int, error) {
	var onCards []int
	for i, p := range m.Profiles {
		if named != "" && p.Name == named {
			if p.CPUOnly() {
				return 0, fmt.Errorf("%s's %s profile runs on CPUs; give it a gpu_class first", m.Name, p.Name)
			}
			return i, nil
		}
		if !p.CPUOnly() {
			onCards = append(onCards, i)
		}
	}
	switch {
	case named != "":
		return 0, fmt.Errorf("%s has no profile %q", m.Name, named)
	case len(onCards) == 0:
		return 0, fmt.Errorf("%s runs on CPUs only; a profile on cards needs a gpu_class", m.Name)
	case len(onCards) > 1:
		var names []string
		for _, i := range onCards {
			names = append(names, m.Profiles[i].Name)
		}
		return 0, fmt.Errorf("%s has several profiles on cards (%s): name one", m.Name, strings.Join(names, ", "))
	}
	return onCards[0], nil
}

// respread is a profile spread over n cards of one node.
func respread(p catalog.Profile, n int, sh shape) (catalog.Profile, error) {
	for _, f := range tpFlags[p.Engine] {
		p.Args = dropArg(p.Args, f)
	}
	switch p.Engine {
	case "vllm", "sglang", "trtllm":
		if !modelplan.SplitsHeads(sh.kvHeads, n) {
			return p, fmt.Errorf("the model has %d KV heads, which tensor parallel cannot spread over %d cards: "+
				"a number that divides them, or that they divide", sh.kvHeads, n)
		}
		p.TP, p.GPUsPerNode = n, n
	case "llamacpp":
		// llama.cpp has no tensor parallel: it splits the layers over a
		// node's cards, in proportion to their free memory.
		p.TP, p.GPUsPerNode = 1, n
	default:
		return p, fmt.Errorf("a %s replica runs on one card", p.Engine)
	}
	// One node: a replica spanning nodes is laid out by hand, with a fabric.
	p.Nodes, p.PP = 1, 0
	return p, nil
}

// contextOf is the tokens a profile's KV cache holds: llama.cpp's pool, or
// the engine's limit, or the model's context.
func contextOf(m catalog.Model, p catalog.Profile) int {
	switch p.Engine {
	case "llamacpp":
		if n := argInt(p.Args, "--ctx-size"); n > 0 {
			return n
		}
	case "vllm":
		if n := argInt(p.Args, "--max-model-len"); n > 0 {
			return n
		}
	case "sglang":
		if n := argInt(p.Args, "--context-length"); n > 0 {
			return n
		}
	case "strata":
		if n := argInt(p.Args, "--context"); n > 0 {
			return n
		}
	case "freetoken":
		if n := argInt(p.Args, "--max-seq-len-override"); n > 0 {
			return n
		}
	}
	return m.ContextLength
}

func (s *Server) replicasOf(model string) []string {
	var ids []string
	if s.deps.Fleet != nil {
		for _, r := range s.deps.Fleet.Replicas() {
			if r.Model == model {
				ids = append(ids, r.ID)
			}
		}
	}
	return ids
}

// roomFor says whether the new replicas fit: a node with n cards of the
// class at all, and free beside the running ones — or, asked to start them
// in place, once the model's own are stopped. True is in place.
func (s *Server) roomFor(class string, n int, model string, running, inPlace bool) (bool, error) {
	if s.deps.Pool == nil {
		return inPlace, nil
	}
	own := map[string]int{}
	if s.deps.Fleet != nil {
		for _, r := range s.deps.Fleet.Replicas() {
			if r.Model == model {
				for _, sh := range r.Shards {
					own[sh.NodeID] += len(sh.GPUIndices)
				}
			}
		}
	}
	most, freeMost, besideOK, inPlaceOK := 0, 0, false, false
	for _, node := range s.deps.Pool() {
		for _, k := range modelplan.Kinds(node) {
			if k.Class != class {
				continue
			}
			most, freeMost = max(most, k.GPUs), max(freeMost, k.FreeGPUs)
			besideOK = besideOK || k.FreeGPUs >= n
			inPlaceOK = inPlaceOK || k.FreeGPUs+own[node.ID] >= n
		}
	}
	switch {
	case most == 0:
		return false, fmt.Errorf("no ready node has %s cards", class)
	case most < n:
		return false, fmt.Errorf("no node has %d %s cards; the most one has is %d", n, class, most)
	case !running || besideOK:
		return false, nil
	case !inPlaceOK:
		return false, fmt.Errorf("no node has %d %s cards free, even with this model's own stopped (%d at most free now): "+
			"other models hold them", n, class, freeMost)
	case !inPlace:
		return false, fmt.Errorf("no node has %d %s cards free beside the running replica (%d at most); "+
			"--in-place stops each replica first, and the model does not answer while the new one loads", n, class, freeMost)
	}
	return true, nil
}

// dropArg is arguments without a flag and its value.
func dropArg(args []string, flag string) []string {
	out := make([]string, 0, len(args))
	for i := 0; i < len(args); i++ {
		switch {
		case args[i] == flag:
			i++ // its value
		case strings.HasPrefix(args[i], flag+"="):
		default:
			out = append(out, args[i])
		}
	}
	return out
}