llmgw/internal/cli/serverinit_choices.go

package cli

import (
	"fmt"
	"math"
	"net"
	"os"
	"os/exec"
	"path"
	"slices"
	"strconv"
	"strings"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/agent/provision"
	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
	"gitea.mynodes.xyz/ilker/llmgw/internal/pgha"
	"gitea.mynodes.xyz/ilker/llmgw/internal/preflight"
)

// What a first install lets a person choose beyond which parts it has: which
// starter model, how the judge decides, what the portal speaks, the limits a
// key starts with, and whether pictures and 3D come by themselves. Each zero
// value is what init did before it could be chosen, so `llmgw server init`
// with no flags installs what it always did.
//
// All of it is written into server.yaml when init writes it, and the cluster's
// keys among it become the cluster's first settings when the control plane
// first starts. A server.yaml already there is never overwritten, so a second
// run changes none of it: the settings page and `llmgw ctl config` do.

// starterSpec is a model init can fetch and serve so there is something to ask.
type starterSpec struct {
	Name, Alias, Source string
	// SizeGB is the weights' size, for the memory a node needs; 0 is not known.
	SizeGB float64
	About  string
}

// starterChoices are the starters offered, smallest first: the first is the
// default. Each is a single GGUF file, Q4_K_M.
var starterChoices = []starterSpec{
	{starterName, starterAlias, starterSource, 0.5, "answers at once on any CPU; fast, not clever"},
	{"qwen2.5-1.5b-instruct", "qwen-1.5b", "hf:Qwen/Qwen2.5-1.5B-Instruct-GGUF/qwen2.5-1.5b-instruct-q4_k_m.gguf", 1.1,
		"still quick on a CPU, and noticeably better"},
	{"qwen2.5-3b-instruct", "qwen-3b", "hf:Qwen/Qwen2.5-3B-Instruct-GGUF/qwen2.5-3b-instruct-q4_k_m.gguf", 2.1,
		"slow on a CPU, quick on any GPU"},
	{"qwen3-4b", "", "hf:Qwen/Qwen3-4B-GGUF/Qwen3-4B-Q4_K_M.gguf", 2.5,
		"thinks before it answers; a GPU is worth having"},
	{"qwen3-8b", "", "hf:Qwen/Qwen3-8B-GGUF/Qwen3-8B-Q4_K_M.gguf", 5.0,
		"a real assistant; a GPU with 8 GB or more"},
	{"qwen3-14b", "", "hf:Qwen/Qwen3-14B-GGUF/Qwen3-14B-Q4_K_M.gguf", 9.0,
		"better still; a GPU with 16 GB or more"},
}

// starterFrom reads --starter: a name from starterChoices, or the hf: source
// of any GGUF file, named after the file.
func starterFrom(s string) (starterSpec, error) {
	s = strings.TrimSpace(s)
	for _, c := range starterChoices {
		if s == c.Name || (c.Alias != "" && s == c.Alias) {
			return c, nil
		}
	}
	if rest, ok := strings.CutPrefix(s, "hf:"); ok && strings.Count(rest, "/") >= 2 && strings.HasSuffix(strings.ToLower(rest), ".gguf") {
		name := strings.ToLower(strings.TrimSuffix(path.Base(rest), path.Ext(rest)))
		return starterSpec{Name: name, Source: s, About: "chosen at install"}, nil
	}
	var names []string
	for _, c := range starterChoices {
		names = append(names, c.Name)
	}
	return starterSpec{}, fmt.Errorf("starter %q: one of %s, or hf:<owner>/<repo>/<file>.gguf", s, strings.Join(names, ", "))
}

func starterNames() []string {
	var out []string
	for _, c := range starterChoices {
		out = append(out, c.Name)
	}
	return out
}

// readChoices takes the flags that are not read straight into init: the
// starter by its name, the judge's categories by commas. Then it checks them
// all.
func (in *serverInit) readChoices(starter, categories string) error {
	if strings.TrimSpace(starter) != "" {
		s, err := starterFrom(starter)
		if err != nil {
			return err
		}
		in.starter = s
	}
	in.judge.Categories = nil
	for _, c := range strings.Split(categories, ",") {
		if c = strings.TrimSpace(c); c != "" {
			in.judge.Categories = append(in.judge.Categories, c)
		}
	}
	return in.checkChoices()
}

// checkChoices says what among the choices cannot be installed.
func (in *serverInit) checkChoices() error {
	switch in.language {
	case "", "en", "tr":
	default:
		return fmt.Errorf("language %q: en or tr", in.language)
	}
	switch in.judge.Fail {
	case "", "closed", "open":
	default:
		return fmt.Errorf("judge fail %q: closed or open", in.judge.Fail)
	}
	for _, c := range in.judge.Categories {
		if !slices.Contains(judgeCategories, c) {
			return fmt.Errorf("judge category %q: one of %s", c, strings.Join(judgeCategories, ", "))
		}
	}
	switch {
	case in.judge.Replicas < 0 || in.judge.Replicas > 8:
		return fmt.Errorf("judge replicas %d: from 1 to 8", in.judge.Replicas)
	case in.judge.Budget < 0 || in.judge.Budget > time.Minute:
		return fmt.Errorf("judge budget %s: at most a minute", in.judge.Budget)
	case in.starterCtx < 0 || in.starterCtx > 262144:
		return fmt.Errorf("starter context %d: at most 262144 tokens", in.starterCtx)
	case in.starterCtx > 0 && in.starterCtx < 512:
		return fmt.Errorf("starter context %d: at least 512 tokens", in.starterCtx)
	case in.defaultRPM < 0 || in.defaultTPM < 0:
		return fmt.Errorf("default limits cannot be negative")
	case in.maxKeyTTL < 0 || (in.maxKeyTTL > 0 && in.maxKeyTTL < time.Hour):
		return fmt.Errorf("max key ttl %s: an hour at least", in.maxKeyTTL)
	case in.evictAfterIdle < 0 || in.minLifetime < 0:
		return fmt.Errorf("on-demand times cannot be negative")
	case in.slurm && in.noNode:
		return fmt.Errorf("slurm runs slurmctld on this machine's node: it needs the node (no --no-node)")
	case in.vip != "" && !in.standbys:
		return fmt.Errorf("a floating address moves between control planes: it needs the standbys")
	}
	if in.slurm {
		// The controller is this machine's node, named by its hostname: a
		// name Slurm does not take would keep the server from starting.
		if err := (config.SlurmSettings{Enabled: true, Controller: localNodeID()}).Check(); err != nil {
			return fmt.Errorf("slurm: this machine's hostname: %w", err)
		}
	}
	if in.vip != "" {
		ip, _, err := net.ParseCIDR(in.vip)
		if err != nil || ip.To4() == nil {
			return fmt.Errorf("floating address %q: an IPv4 address with its prefix, as 10.0.0.160/24", in.vip)
		}
	}
	return nil
}

// starterModelSpec is the starter this install serves.
func (in *serverInit) starterModelSpec() starterSpec {
	if in.starter.Source == "" {
		return starterChoices[0]
	}
	return in.starter
}

// minMemoryGB is what a node needs free for the starter: its weights and
// room for its context, never less than the 2 GB it always asked.
func (s starterSpec) minMemoryGB() int {
	return max(2, int(math.Ceil(s.SizeGB+1.5)))
}

// servingMemoryMB is what this machine needs, beyond what the control plane
// does, for the models it will serve itself: the starter's and the judges'
// reservations, and a gigabyte more because the scheduler counts a machine's
// memory in whole gigabytes, rounded down — 7.75 GB is 7. Zero when it serves
// none of them; what names those it does.
//
// Asked before anything is installed: the 8 GB machine the guide was first
// recorded on finished every step and then had no room for the starter, the
// one thing a person asks first.
func (in *serverInit) servingMemoryMB(judgesElsewhere bool) (mb int, what string) {
	if in.noNode {
		return 0, ""
	}
	gb := 0
	var parts []string
	if !in.noModel && (!in.noMirror || in.mirrorURL != "") {
		gb += in.starterModelSpec().minMemoryGB()
		parts = append(parts, "the starter model")
	}
	// Only judges that will be installed: none without a mirror to fetch
	// them from, or from a bundle that does not carry one.
	if !judgesElsewhere && in.judgeComes() {
		n := in.judgeFor().Replicas
		gb += n * judgeMemoryGB
		parts = append(parts, fmt.Sprintf("%d judges", n))
	}
	if gb == 0 {
		return 0, ""
	}
	return (gb + 1) * 1024, strings.Join(parts, " and ")
}

// checkMemory refuses, before anything is installed, a machine too small for
// the models it would serve as well as the control plane.
func (in *serverInit) checkMemory() error {
	if in.unitPath == "" {
		return nil // a development run, served by nothing
	}
	// Run again on a machine already serving — to change --host, say, as
	// the docs advise — it installs no models it does not already serve. A
	// first run that stopped half way is checked again: it still would.
	if in.alreadyServing() {
		return nil
	}
	meminfo, err := os.ReadFile("/proc/meminfo")
	if err != nil {
		return nil
	}
	total, ok := pgha.MemTotal(string(meminfo))
	if !ok {
		return nil
	}
	extra, what := in.servingMemoryMB(len(in.judgeMachines) > 0)
	if extra == 0 {
		return nil
	}
	have, need := int(total>>20), preflight.Needs[preflight.RoleControl].MemoryMB+extra
	if have >= need {
		return nil
	}
	return fmt.Errorf("memory: this machine has %d MB; the control plane and %s it would serve need %d. "+
		"Give it more, or leave them off it: --no-model, --no-guard-model, or --no-node to serve models only on nodes added later", have, what, need)
}

// defaultContext is the context a client may use of the starter.
const defaultContext = 4096

// starterContext is the context a client may use of the starter.
func (in *serverInit) starterContext() int {
	if in.starterCtx > 0 {
		return in.starterCtx
	}
	return defaultContext
}

// judgeSpec is how the guardrails' judge decides.
type judgeSpec struct {
	// Fail is what is done with a request no judge read in time: closed
	// refuses it, open serves it flagged.
	Fail       string
	Categories []string
	Replicas   int
	Budget     time.Duration
}

// judgeCategories are the ones a judge can be asked about, as Qwen3Guard
// names them; defaultJudgeCategories those it is asked about unless told.
var (
	judgeCategories = []string{"Violent", "Non-violent Illegal Acts", "Sexual Content or Sexual Acts",
		"Suicide & Self-Harm", "Jailbreak", "PII", "Unethical Acts", "Politically Sensitive Topics", "Copyright Violation"}
	defaultJudgeCategories = judgeCategories[:5]
)

// judgeFor is the judge as this install sets it up, the defaults filled in.
func (in *serverInit) judgeFor() judgeSpec {
	j := in.judge
	if j.Fail == "" {
		j.Fail = "closed"
	}
	if len(j.Categories) == 0 {
		j.Categories = defaultJudgeCategories
	}
	if j.Replicas <= 0 {
		j.Replicas = 2
		if in.singleMachine() {
			j.Replicas = 1
		}
	}
	if j.Budget <= 0 {
		j.Budget = 5 * time.Second
	}
	return j
}

// singleMachine is an install where this machine is the only node there
// will be, as far as init knows: a second judge replica there shares the
// same cores and is lost with the same machine, so it buys nothing but
// half the speed of each.
func (in *serverInit) singleMachine() bool {
	return !in.noNode && !in.standbys && len(in.patroni) == 0 && len(in.patroniNodes) == 0
}

// The configuration's limits, as init writes them unless told otherwise.
const (
	defaultRPM       = 60
	defaultTPM       = 200000
	defaultMaxKeyTTL = 2160 * time.Hour
)

// tuneConfig writes this install's choices into the server.yaml template:
// the lines the template has, changed, and the blocks it has not, added.
// Unchosen, the template comes back as it was.
func (in *serverInit) tuneConfig(body string) string {
	set := func(old, now string) {
		if old != now {
			body = strings.Replace(body, old, now, 1)
		}
	}
	if in.signupOpen {
		set("  signup_open: false", "  signup_open: true")
	}
	if in.defaultRPM > 0 {
		set("  default_rpm: 60", "  default_rpm: "+strconv.Itoa(in.defaultRPM))
	}
	if in.defaultTPM > 0 {
		set("  default_tpm: 200000", "  default_tpm: "+strconv.Itoa(in.defaultTPM))
	}
	if in.maxKeyTTL > 0 && in.maxKeyTTL != defaultMaxKeyTTL {
		set("  max_key_ttl: 2160h   # 90 days", fmt.Sprintf("  max_key_ttl: %s   # %s", durationYAML(in.maxKeyTTL), days(in.maxKeyTTL)))
	}
	if in.language != "" {
		set("  # language: tr", "  language: "+in.language)
	}
	var extra []string
	if in.media {
		extra = append(extra, "# Pictures and 3D, chosen at install: the nodes install ComfyUI, and the\n"+
			"# leader adds a picture model and a 3D model by itself once a card fits\n"+
			"# one ('llmgw ctl model presets auto off' stops it).\n"+
			"engines:\n  comfyui: "+provision.ComfyUIVersion()+"\npresets:\n  auto: true\n")
	}
	if in.slurm {
		extra = append(extra, slurmSection(localNodeID()))
	}
	if in.evictAfterIdle > 0 || in.minLifetime > 0 {
		b := "# How models that load when asked for take turns on a card, chosen at\n# install.\non_demand:\n"
		if in.evictAfterIdle > 0 {
			b += "  evict_after_idle: " + durationYAML(in.evictAfterIdle) + "\n"
		}
		if in.minLifetime > 0 {
			b += "  min_lifetime: " + durationYAML(in.minLifetime) + "\n"
		}
		extra = append(extra, b)
	}
	if len(extra) > 0 {
		body += "\n" + strings.Join(extra, "\n")
	}
	return body
}

// slurmSection is server.yaml's slurm block as an install chooses it:
// slurmctld on this machine's node, slurmd on every node.
func slurmSection(controller string) string {
	return "# Slurm beside llmgw, chosen at install: slurmctld on this machine, slurmd on\n" +
		"# every node, those that join later too. Each node is llmgw's, drained in\n" +
		"# Slurm, until 'llmgw ctl slurm lend NODE' gives it to Slurm; 'llmgw ctl slurm\n" +
		"# reclaim NODE' takes it back once its jobs end.\n" +
		"slurm:\n  enabled: true\n  controller: " + yamlString(controller) + "\n  all_nodes: true\n"
}

// localNodeID is the node id this machine's agent takes: its hostname.
func localNodeID() string {
	if h, err := os.Hostname(); err == nil && h != "" {
		return h
	}
	return "localhost"
}

// durationYAML is a duration as a person writes it: 2h, 90m, 1m30s.
func durationYAML(d time.Duration) string {
	switch {
	case d == 0:
		return "0s"
	case d%time.Hour == 0:
		return fmt.Sprintf("%dh", d/time.Hour)
	case d%time.Minute == 0:
		return fmt.Sprintf("%dm", d/time.Minute)
	}
	return d.String()
}

func days(d time.Duration) string {
	n := int(d.Hours() / 24)
	if n == 1 {
		return "1 day"
	}
	return fmt.Sprintf("%d days", n)
}

// initFlags is this install as `llmgw server init`'s flags, for doing it
// again from a script; what is left at its default is left out.
func (in *serverInit) initFlags() []string {
	var out []string
	add := func(flag string, v ...string) {
		if len(v) == 0 {
			out = append(out, flag)
			return
		}
		out = append(out, flag+" "+shellWord(v[0]))
	}
	if in.admin != "" {
		add("--admin", in.admin)
	}
	for _, h := range in.hosts {
		add("--host", h)
	}
	if in.language != "" {
		add("--language", in.language)
	}
	if in.signupOpen {
		add("--signup-open")
	}
	if in.defaultRPM > 0 && in.defaultRPM != defaultRPM {
		add("--default-rpm", strconv.Itoa(in.defaultRPM))
	}
	if in.defaultTPM > 0 && in.defaultTPM != defaultTPM {
		add("--default-tpm", strconv.Itoa(in.defaultTPM))
	}
	if in.maxKeyTTL > 0 && in.maxKeyTTL != defaultMaxKeyTTL {
		add("--max-key-ttl", durationYAML(in.maxKeyTTL))
	}
	if in.postgres {
		add("--postgres")
	}
	for _, t := range in.patroni {
		add("--patroni", t)
	}
	for _, m := range in.judgeMachines {
		add("--judge-node", m.target)
	}
	if in.secondMirror > 0 {
		add("--second-mirror", fmt.Sprint(in.secondMirror))
	}
	switch {
	case in.mirrorURL != "":
		add("--mirror", in.mirrorURL)
	case in.noMirror:
		add("--no-mirror")
	case in.mirrorPath != "" && in.mirrorPath != mirrorDir:
		add("--mirror-dir", in.mirrorPath)
	}
	if in.bundle != "" {
		add("--bundle", in.bundle)
		add("--bundle-sha256", in.bundleSHA)
	}
	if in.noNode {
		add("--no-node")
	}
	if in.noModel {
		add("--no-model")
	} else {
		if s := in.starterModelSpec(); s.Name != starterName {
			v := s.Name
			if !slices.ContainsFunc(starterChoices, func(c starterSpec) bool { return c.Name == s.Name }) {
				v = s.Source
			}
			add("--starter", v)
		}
		if in.starterCtx > 0 && in.starterCtx != defaultContext {
			add("--starter-context", strconv.Itoa(in.starterCtx))
		}
	}
	if in.noGuardModel {
		add("--no-guard-model")
	} else {
		// Against the defaults for an install of the same shape: one judge
		// on one machine, two otherwise.
		shape := serverInit{noNode: in.noNode, standbys: in.standbys, patroni: in.patroni, patroniNodes: in.patroniNodes}
		j, d := in.judgeFor(), shape.judgeFor()
		if j.Fail != d.Fail {
			add("--judge-fail", j.Fail)
		}
		if !slices.Equal(j.Categories, d.Categories) {
			add("--judge-categories", strings.Join(j.Categories, ","))
		}
		if j.Replicas != d.Replicas {
			add("--judge-replicas", strconv.Itoa(j.Replicas))
		}
		if j.Budget != d.Budget {
			add("--judge-budget", durationYAML(j.Budget))
		}
	}
	if in.noFirewall {
		add("--no-firewall")
	}
	if in.standbys {
		add("--standbys")
	}
	if in.vip != "" {
		add("--vip", in.vip)
	}
	for _, w := range in.webhooks {
		for _, s := range strings.Split(w, ",") {
			if s = strings.TrimSpace(s); s != "" {
				add("--webhook", s)
			}
		}
	}
	if in.media {
		add("--media")
	}
	if in.slurm {
		add("--slurm")
	}
	if in.evictAfterIdle > 0 {
		add("--evict-after-idle", durationYAML(in.evictAfterIdle))
	}
	if in.minLifetime > 0 {
		add("--min-lifetime", durationYAML(in.minLifetime))
	}
	return out
}

// shellWord quotes a word for a shell when it needs it.
func shellWord(s string) string {
	if s != "" && !strings.ContainsAny(s, " \t\n'\"\\$`&;|<>(){}*?[]#~!") {
		return s
	}
	return "'" + strings.ReplaceAll(s, "'", `'\''`) + "'"
}

// alreadyServing reports whether this machine's gateway is set up and
// running: a server.yaml, and its unit active.
func (in *serverInit) alreadyServing() bool {
	if _, err := os.Stat(in.configPath()); err != nil {
		return false
	}
	if in.serviceActive != nil {
		return in.serviceActive()
	}
	return exec.Command("systemctl", "is-active", "--quiet", "llmgw-server").Run() == nil
}