llmgw/internal/cli/model_update.go

package cli

import (
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"github.com/spf13/cobra"

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

func newModelUpdatesCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	var check bool
	c := &cobra.Command{
		Use:   "updates [MODEL]",
		Short: "Show which models' repositories have published new weights",
		Long: "A model served from the registry remembers the repository and commit its\n" +
			"files came from, and the control plane asks the registry every\n" +
			"mirror.update_check (six hours by default) whether those files have\n" +
			"changed there. A new commit that leaves them alone is not an update.\n\n" +
			"--check asks now rather than showing the last answer.",
		Args: cobra.MaximumNArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			var list []modelserve.UpdateInfo
			if check || len(args) == 1 {
				req := map[string]string{}
				if len(args) == 1 {
					req["model"] = args[0]
				}
				body, _ := json.Marshal(req)
				if err := client().post(cmd.Context(), "/api/v1/updates/check", body, &list); err != nil {
					return err
				}
			} else if err := client().get(cmd.Context(), "/api/v1/updates", &list); err != nil {
				return err
			}
			if *asJSON {
				return emitJSON(list)
			}
			if len(list) == 0 {
				fmt.Fprintln(cmd.OutOrStdout(), "no model has an origin to compare yet; one is checked a few minutes after the control plane starts")
				return nil
			}
			w := tabTo(cmd.OutOrStdout())
			fmt.Fprintln(w, "MODEL\tREPOSITORY\tRUNS\tLATEST\tUPDATE\tCHECKED")
			for _, u := range list {
				fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\n", u.Model, u.Repo, catalog.ShortRevision(u.Current),
					catalog.ShortRevision(u.Latest), updateCell(u), ago(u.CheckedAt, time.Now()))
			}
			if err := w.Flush(); err != nil {
				return err
			}
			for _, u := range list {
				if u.Available {
					fmt.Fprintf(cmd.OutOrStdout(), "\nTo move one to its new weights: llmgw ctl model update MODEL\n")
					break
				}
			}
			return nil
		},
	}
	c.Flags().BoolVar(&check, "check", false, "ask the registry now")
	return c
}

// updateCell says what an update would be, or why there is none.
func updateCell(u modelserve.UpdateInfo) string {
	switch {
	case u.Err != "":
		return "could not ask: " + truncate(u.Err, 60)
	case u.Available:
		return fmt.Sprintf("yes: %d file(s), %s", len(u.Changed), bytesShort(u.Bytes))
	case len(u.Missing) > 0:
		return "the repository no longer has " + strings.Join(u.Missing, ", ")
	case u.Latest != "" && u.Latest != u.Current:
		return "no: new commit, same weights"
	}
	return "-"
}

func newModelUpdateCmd(client func() *ctlClient) *cobra.Command {
	return &cobra.Command{
		Use:   "update MODEL",
		Short: "Move a model to the weights its repository publishes now",
		Long: "The new files are fetched into the mirror beside the ones in use and\n" +
			"checked against the registry's checksums and the nodes' llama.cpp; then\n" +
			"the catalog names them and the replicas are replaced one at a time, each\n" +
			"new one up before an old one stops, and one is asked to say hello.\n\n" +
			"The old files stay on the mirror: `llmgw ctl model rollback MODEL` goes\n" +
			"back to them. A model from server.yaml is changed there instead.",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			var job modelserve.Job
			if err := client().post(cmd.Context(), "/api/v1/models/"+args[0]+"/update", []byte("{}"), &job); err != nil {
				return err
			}
			return followJob(cmd, client, job, "runs the new weights")
		},
	}
}

func newModelRollbackCmd(client func() *ctlClient) *cobra.Command {
	return &cobra.Command{
		Use:   "rollback MODEL",
		Short: "Move a model back to the weights its last update replaced",
		Args:  cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			var out map[string]string
			if err := client().post(cmd.Context(), "/api/v1/models/"+args[0]+"/rollback", []byte("{}"), &out); err != nil {
				return err
			}
			fmt.Fprintln(cmd.OutOrStdout(), out["result"])
			return nil
		},
	}
}

func newModelOriginCmd(client func() *ctlClient) *cobra.Command {
	return &cobra.Command{
		Use:   "origin MODEL OWNER/REPO",
		Short: "Say which registry repository a model's weights came from",
		Long: "For a model fetched before llmgw kept this, and that it could not find on\n" +
			"its own. The repository is checked: it must have the commit the mirror's\n" +
			"directory is named after, holding the model's files at the sizes the\n" +
			"mirror has. From then on its updates are watched.",
		Args: cobra.ExactArgs(2),
		RunE: func(cmd *cobra.Command, args []string) error {
			body, _ := json.Marshal(map[string]string{"repo": args[1]})
			var o struct{ Repo, Revision string }
			if err := client().post(cmd.Context(), "/api/v1/models/"+args[0]+"/origin", body, &o); err != nil {
				return err
			}
			fmt.Fprintf(cmd.OutOrStdout(), "%s came from %s at %s; its updates are watched\n", args[0], o.Repo,
				catalog.ShortRevision(o.Revision))
			return nil
		},
	}
}