llmgw/internal/cli/model.go

package cli

import (
	"encoding/json"
	"fmt"
	"os/user"
	"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"
)

// 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) *cobra.Command {
	c := &cobra.Command{
		Use:   "model",
		Short: "Add or remove models while the cluster is running",
	}
	c.AddCommand(newModelAddCmd(client), newModelRemoveCmd(client), newModelPlanCmd(client))
	return c
}

func newModelAddCmd(client func() *ctlClient) *cobra.Command {
	var path string
	var skipCheck bool
	c := &cobra.Command{
		Use:   "add -f <file.yaml>",
		Short: "Add a model from a definition file",
		Long: "The 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.",
		Args: cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			if path == "" {
				return fmt.Errorf("--file is required: the model definition, as YAML")
			}
			// Read 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.
			k := koanf.New(".")
			if err := k.Load(file.Provider(path), yaml.Parser()); err != nil {
				return fmt.Errorf("%s: %w", path, err)
			}
			var m catalog.Model
			if err := k.UnmarshalWithConf("", &m, koanf.UnmarshalConf{Tag: "koanf"}); err != nil {
				return fmt.Errorf("%s: %w", path, err)
			}
			if m.Name == "" {
				return fmt.Errorf("%s: the definition has no name", path)
			}

			// 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.
			if !skipCheck {
				size, err := checkSource(cmd.Context(), m.Source)
				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 := httpSource(m.Source); isURL {
					fmt.Fprintln(cmd.OutOrStdout(), sourceNote(m.Source, 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)
			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
}

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 {
	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
}