llmgw/internal/cli/model.go

package cli

import (
	"encoding/json"
	"fmt"
	"os/user"

	"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/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"
}