llmgw/internal/cli/model.go

package cli

import (
	"encoding/json"
	"fmt"
	"os"
	"os/user"
	"slices"
	"strconv"
	"strings"
	"time"

	"github.com/knadh/koanf/parsers/yaml"
	"github.com/knadh/koanf/providers/file"
	"github.com/knadh/koanf/v2"
	"github.com/spf13/cobra"

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

// newModelCmd adds and removes models without a restart.
//
// Adding a model used to mean editing a file on the control plane and
// restarting it: a settling window, adoption churn, and on one occasion two
// minutes of outage. For a fixed production set that is a reasonable trade; for
// a platform whose promise is that people can try things, it is the difference
// between trying one and not bothering.
func newModelCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	c := &cobra.Command{
		Use:   "model",
		Short: "Add or remove models while the cluster is running",
	}
	c.AddCommand(newModelAddCmd(client), newModelPresetsCmd(), newModelRemoveCmd(client), newModelPlanCmd(client),
		newResourcesCmd(client), newClassCmd(client), newOwnerCmd(client),
		newTargetCmd(client, asJSON), newAdaptersCmd(client),
		newModelServeCmd(client), newModelUpdatesCmd(client, asJSON), newModelUpdateCmd(client),
		newModelRollbackCmd(client), newModelOriginCmd(client),
		newModelCanaryCmd(client), newModelPromoteCmd(client), newModelCanariesCmd(client, asJSON),
		newModelGateCmd(client), newModelFallbacksCmd(client))
	return c
}

func newModelAddCmd(client func() *ctlClient) *cobra.Command {
	var path string
	var skipCheck bool
	c := &cobra.Command{
		Use:   "add PRESET | -f <file.yaml>",
		Short: "Add a model that comes with llmgw, or one from a definition file",
		Long: "A preset is a model that comes with llmgw, its workflow, nodes and licence\n" +
			"already found and tried: " + strings.Join(catalog.PresetNames(), ", ") + " (llmgw ctl model presets).\n" +
			"Adding a ComfyUI one also has the nodes install ComfyUI, when the cluster\n" +
			"has not been told otherwise.\n\n" +
			"A definition is the same YAML a model takes in server.yaml, so an\n" +
			"entry can be moved between the two by copying it.\n\n" +
			"The catalog is rebuilt and validated before anything is written, so a\n" +
			"definition that collides with an existing name, or names a fallback that\n" +
			"does not exist, is refused rather than stored and then ignored on every\n" +
			"start. Once accepted it takes effect on the reconciler's next pass — no\n" +
			"restart — and survives one.",
		Example: "  llmgw ctl model add triposr\n" +
			"  llmgw ctl model add -f qwen.yaml",
		Args: cobra.MaximumNArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			var m catalog.Model
			switch {
			case len(args) == 1 && path != "":
				return fmt.Errorf("a preset or a file, not both")
			case len(args) == 1:
				p, ok := catalog.PresetNamed(args[0])
				if !ok {
					return fmt.Errorf("no preset %q: there are %s; any other model is added from a file (-f)",
						args[0], strings.Join(catalog.PresetNames(), ", "))
				}
				m = p.Model
			case path == "":
				return fmt.Errorf("a preset (%s) or --file, the model definition as YAML",
					strings.Join(catalog.PresetNames(), ", "))
			default:
				var err error
				if m, err = readModelFile(path); err != nil {
					return err
				}
			}

			// Before the catalog is touched, because a source that does not
			// resolve is otherwise discovered by a replica failing to start
			// minutes later on a node — with an error naming the node rather
			// than the typo.
			// Every source the model's profiles load, not the model's alone.
			if !skipCheck {
				for _, src := range m.WeightSources() {
					size, err := sourcecheck.Check(cmd.Context(), src)
					if err != nil {
						return fmt.Errorf("%w\n  Pass --no-source-check to add it anyway, "+
							"which is right when the mirror is not up yet or the nodes "+
							"reach it by an address this machine does not.", err)
					}
					if _, isURL := sourcecheck.HTTPSource(src); isURL {
						fmt.Fprintln(cmd.OutOrStdout(), sourceNote(src, size))
					}
				}
			}

			body, err := json.Marshal(m)
			if err != nil {
				return err
			}

			var res map[string]string
			if err := client().post(cmd.Context(), "/api/v1/models", body, &res); err != nil {
				return err
			}
			fmt.Fprintf(cmd.OutOrStdout(),
				"%s added. It will be placed on the next reconcile pass; watch it with:\n"+
					"  llmgw ctl models\n", m.Name)
			if m.ServesComfyUI() && !m.IsExternal() {
				return ensureComfyUI(cmd, client, m.Name)
			}
			return nil
		},
	}
	c.Flags().StringVarP(&path, "file", "f", "", "YAML file holding the model definition")
	c.Flags().BoolVar(&skipCheck, "no-source-check", false,
		"add it without asking whether the weights are actually at that URL")
	return c
}

// readModelFile reads a model's definition. By koanf, the same parser and
// the same tags that read server.yaml — so an entry can be moved between the
// two by copying it, and a field that works in one cannot silently mean
// nothing in the other.
func readModelFile(path string) (catalog.Model, error) {
	var m catalog.Model
	k := koanf.New(".")
	if err := k.Load(file.Provider(path), yaml.Parser()); err != nil {
		return m, fmt.Errorf("%s: %w", path, err)
	}
	if err := k.UnmarshalWithConf("", &m, koanf.UnmarshalConf{Tag: "koanf"}); err != nil {
		return m, fmt.Errorf("%s: %w", path, err)
	}
	if m.Name == "" {
		return m, fmt.Errorf("%s: the definition has no name", path)
	}
	return m, nil
}

func newModelRemoveCmd(client func() *ctlClient) *cobra.Command {
	return &cobra.Command{
		Use:   "remove <name>",
		Short: "Remove a model that was added at runtime",
		Long: "Only models added at runtime. One declared in server.yaml would come\n" +
			"back on the next restart, so removing it here would leave the cluster\n" +
			"disagreeing with its own configuration until then.\n\n" +
			"Its replicas are stopped on the next pass, after the endpoint has been\n" +
			"withdrawn, so traffic stops before the engine does.",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			var res map[string]string
			if err := client().delete(cmd.Context(), "/api/v1/models/"+args[0], &res); err != nil {
				return err
			}
			fmt.Fprintf(cmd.OutOrStdout(), "%s removed.\n", args[0])
			return nil
		},
	}
}

// actor names who made a change, for the audit trail on a stored model.
func actor() string {
	// Under sudo, the person who ran it rather than root.
	if u := os.Getenv("SUDO_USER"); u != "" && os.Geteuid() == 0 {
		return u
	}
	if u, err := user.Current(); err == nil && u.Username != "" {
		return u.Username
	}
	return "unknown"
}

func newScaleCmd(client func() *ctlClient) *cobra.Command {
	var min, max int
	var idle time.Duration
	c := &cobra.Command{
		Use:   "scale <model> --min N [--max M]",
		Short: "Change how many replicas a model runs",
		Long: "The replica count is --min. --max is a ceiling, not a target: a model\n" +
			"with a minimum runs exactly that many, and nothing scales it up under\n" +
			"load. Leaving --max out makes it follow --min.\n\n" +
			"--min 0 with a --max above it is the on-demand model: nothing runs until\n" +
			"somebody asks, and it stops again after the idle timeout. That loads one\n" +
			"replica however high --max is. --min 0 --max 0 is a model kept in the\n" +
			"catalog and deliberately not running.\n\n" +
			"Only models added at runtime. One declared in server.yaml is read again\n" +
			"on every start, so a number set here would hold until the next restart\n" +
			"and then revert — change it in the file instead.\n\n" +
			"It takes effect on the reconciler's next pass, and survives a restart.",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			if !cmd.Flags().Changed("min") {
				return fmt.Errorf("--min is required: it is the replica count")
			}
			body, err := json.Marshal(adminapi.ScaleRequest{Min: min, Max: max, IdleTimeout: idle})
			if err != nil {
				return err
			}
			var res adminapi.ScaleResponse
			if err := client().post(cmd.Context(),
				"/api/v1/models/"+args[0]+"/scale", body, &res); err != nil {
				return err
			}
			// What is in force, not what was asked for: the catalog fills in a
			// max that was left out, and printing the request back would show a
			// number the cluster is not using.
			out := cmd.OutOrStdout()
			switch {
			case res.OnDemand:
				fmt.Fprintf(out, "%s is on-demand: loaded when asked for, up to %d, stopped when idle.\n",
					res.Model, res.Max)
			case res.Min == 0:
				fmt.Fprintf(out, "%s will not be placed; it stays in the catalog.\n", res.Model)
			default:
				fmt.Fprintf(out, "%s: %d replicas (max %d). It takes effect on the next pass; watch it with:\n"+
					"  llmgw ctl models\n", res.Model, res.Min, res.Max)
			}
			return nil
		},
	}
	c.Flags().IntVar(&min, "min", 0, "replicas to keep running; this is the count")
	c.Flags().IntVar(&max, "max", 0, "ceiling for an on-demand model; follows --min when left out")
	c.Flags().DurationVar(&idle, "idle-timeout", 0,
		"how long an on-demand model stays up unasked-for; zero takes the cluster default")
	return c
}

func newResourcesCmd(client func() *ctlClient) *cobra.Command {
	var memory, cores int
	var profile string
	c := &cobra.Command{
		Use:   "resources <model> --memory GB [--cores N]",
		Short: "Set what a model reserves of a node: memory, and CPU cores",
		Long: "--memory is the host RAM a replica needs: for a model with no GPU, its\n" +
			"weights, KV cache and engine together. Placement reserves it on the node,\n" +
			"so a node holds only as many replicas as it has memory for.\n\n" +
			"--cores is the CPU cores a replica with no GPU computes on. They are\n" +
			"reserved the same way, and become the engine's thread count, so two\n" +
			"replicas on one node stop competing for the same cores. Zero reserves\n" +
			"none and leaves the thread count to the node.\n\n" +
			"Running replicas keep what they were started with; `llmgw ctl restart\n" +
			"--model M` moves them to the new figures.",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			if !cmd.Flags().Changed("memory") {
				return fmt.Errorf("--memory is required")
			}
			body, err := json.Marshal(adminapi.ResourcesRequest{Profile: profile, MemoryGB: memory, CPUCores: cores})
			if err != nil {
				return err
			}
			var m catalog.Model
			if err := client().post(cmd.Context(), "/api/v1/models/"+args[0]+"/resources", body, &m); err != nil {
				return err
			}
			for _, p := range m.Profiles {
				fmt.Fprintf(cmd.OutOrStdout(), "%s/%s: %d GB memory, %d cores\n", m.Name, p.Name, p.MinMemoryGB, p.CPUCores)
			}
			fmt.Fprintf(cmd.OutOrStdout(), "Running replicas keep their old figures until restarted: llmgw ctl restart --model %s\n", m.Name)
			return nil
		},
	}
	c.Flags().IntVar(&memory, "memory", 0, "host RAM a replica needs, in GB")
	c.Flags().IntVar(&cores, "cores", 0, "CPU cores a replica with no GPU computes on; 0 reserves none")
	c.Flags().StringVar(&profile, "profile", "", "one profile; every profile when left out")
	return c
}

func newClassCmd(client func() *ctlClient) *cobra.Command {
	return &cobra.Command{
		Use:   "class <model> critical|standard|opportunistic|default",
		Short: "Set how much a model's capacity is worth against another's",
		Long: "When a model cannot be placed, replicas of a lower class are stopped to\n" +
			"make room for it:\n\n" +
			"  critical       never gives up the replicas of its minimum\n" +
			"  standard       gives way to critical\n" +
			"  opportunistic  gives way to both; the default for an on-demand model\n\n" +
			"Within a class, a model below its minimum may take what another grew into\n" +
			"under load, but not its minimum. \"default\" goes back to what the model's\n" +
			"scaling implies: opportunistic when on demand, standard otherwise.",
		Args: cobra.ExactArgs(2),
		RunE: func(cmd *cobra.Command, args []string) error {
			class := args[1]
			if class == "default" {
				class = ""
			}
			body, _ := json.Marshal(map[string]string{"class": class})
			var res map[string]string
			if err := client().post(cmd.Context(), "/api/v1/models/"+args[0]+"/class", body, &res); err != nil {
				return err
			}
			fmt.Fprintf(cmd.OutOrStdout(), "%s is %s\n", res["model"], res["class"])
			return nil
		},
	}
}

func newAdaptersCmd(client func() *ctlClient) *cobra.Command {
	var add, remove []string
	c := &cobra.Command{
		Use:   "adapters <model> [--add NAME=URL[@SCALE]] [--remove NAME]",
		Short: "List, add or remove a model's LoRA adapters",
		Long: "An adapter is asked for by its NAME as though it were a model, and is served\n" +
			"by this model's replicas with it applied. URL is the adapter, usually in the\n" +
			"mirror. For llama.cpp it must be a GGUF LoRA adapter for this model's\n" +
			"architecture, which is checked before it is attached; for vLLM it is a PEFT\n" +
			"adapter directory (a mirror version ending in /) or a Hugging Face repository\n" +
			"id, which vLLM checks when it loads it. SCALE applies on llama.cpp only.\n" +
			"Replicas load a changed set when restarted: llmgw ctl restart --model MODEL.",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			name := args[0]
			var models []adminapi.ModelView
			if err := client().get(cmd.Context(), "/api/v1/models", &models); err != nil {
				return err
			}
			var current []catalog.Adapter
			found := false
			for _, m := range models {
				if m.Name == name {
					current, found = m.Adapters, true
				}
			}
			if !found {
				return fmt.Errorf("no model %q", name)
			}
			out := cmd.OutOrStdout()
			if len(add) == 0 && len(remove) == 0 {
				if len(current) == 0 {
					fmt.Fprintf(out, "%s has no adapters\n", name)
				}
				for _, a := range current {
					fmt.Fprintf(out, "%s\tscale %g\t%s\n", a.Name, a.EffectiveScale(), a.Source)
				}
				return nil
			}
			next := slices.DeleteFunc(slices.Clone(current), func(a catalog.Adapter) bool { return slices.Contains(remove, a.Name) })
			for _, spec := range add {
				n, rest, ok := strings.Cut(spec, "=")
				if !ok || n == "" || rest == "" {
					return fmt.Errorf("--add %q is not NAME=URL[@SCALE]", spec)
				}
				ad := catalog.Adapter{Name: n, Source: rest}
				if i := strings.LastIndex(rest, "@"); i > 0 {
					if s, err := strconv.ParseFloat(rest[i+1:], 64); err == nil {
						ad.Source, ad.Scale = rest[:i], s
					}
				}
				next = append(next, ad)
			}
			body, _ := json.Marshal(adminapi.AdaptersRequest{Adapters: next})
			var res adminapi.AdaptersRequest
			if err := client().post(cmd.Context(), "/api/v1/models/"+name+"/adapters", body, &res); err != nil {
				return err
			}
			fmt.Fprintf(out, "%s serves %d adapter(s); restart its replicas to load them: llmgw ctl restart --model %s\n",
				name, len(res.Adapters), name)
			return nil
		},
	}
	c.Flags().StringArrayVar(&add, "add", nil, "NAME=URL[@SCALE]; repeatable")
	c.Flags().StringArrayVar(&remove, "remove", nil, "an adapter's name; repeatable")
	return c
}