llmgw/internal/cli/train.go

package cli

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

	"github.com/spf13/cobra"

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

func newTrainCmd(client func() *ctlClient, asJSON *bool) *cobra.Command {
	c := &cobra.Command{
		Use:   "train",
		Short: "Train LoRA adapters on the nodes, on question/answer pairs",
		Long: "A node trains an adapter for a catalog model on pairs (llmgw ctl kb dataset\n" +
			"writes them from a knowledge base; read them through first). It trains with\n" +
			"its installed vLLM's torch and a small overlay, as the engines' user, on a card\n" +
			"with --gpu-memory free — beside an engine, if it leaves that much. What it\n" +
			"makes is kept by the control plane: llmgw ctl train download, or publish it\n" +
			"as the model's adapter.\n\n" +
			"An adapter learns how answers sound — words, format, tone — far better than\n" +
			"what they say. Facts belong to the knowledge base.",
	}

	var model, base, rev, node, data, owner string
	var public bool
	var gpuMem int
	var epochs, rank int
	var lr float64
	var noGGUF bool
	start := &cobra.Command{
		Use:   "start NAME --model MODEL --base-repo OWNER/NAME --node NODE --data FILE",
		Short: "Queue an adapter's training on a node",
		Example: "  llmgw ctl train start policies-tone --model qwen2.5-0.5b-instruct \\\n" +
			"    --base-repo Qwen/Qwen2.5-0.5B-Instruct --node cuda01 --gpu-memory 2300 --data policies.jsonl",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			b, err := os.ReadFile(data)
			if err != nil {
				return err
			}
			req := adminapi.TrainRequest{Name: args[0], Model: model, BaseRepo: base, BaseRevision: rev, Node: node,
				GPUMemoryMiB: gpuMem, Data: string(b), Params: map[string]any{}, Owner: owner, Public: public}
			if cmd.Flags().Changed("epochs") {
				req.Params["epochs"] = epochs
			}
			if cmd.Flags().Changed("rank") {
				req.Params["rank"] = rank
			}
			if cmd.Flags().Changed("lr") {
				req.Params["lr"] = lr
			}
			if noGGUF {
				f := false
				req.GGUF = &f
			}
			body, _ := json.Marshal(req)
			var v adminapi.TrainJobView
			if err := client().patient(2*time.Minute).do(cmd.Context(), "POST", "/api/v1/train/jobs", json.RawMessage(body), &v); err != nil {
				return err
			}
			who := "every caller of " + v.Model
			if v.Owner != "" {
				who = v.Owner + " only"
			}
			fmt.Fprintf(cmd.OutOrStdout(), "queued %s: %s for %s on %s, %d pairs; it will be seen by %s\n",
				v.ID, v.Name, v.Model, v.Node, v.Pairs, who)
			fmt.Fprintf(cmd.OutOrStdout(), "follow it with: llmgw ctl train show %s --follow\n", v.ID)
			return nil
		},
	}
	f := start.Flags()
	f.StringVar(&model, "model", "", "the catalog model the adapter is for")
	f.StringVar(&base, "base-repo", "", "the weights to train on, as safetensors on Hugging Face (owner/name)")
	f.StringVar(&rev, "base-revision", "", "a revision of them (default: the newest)")
	f.StringVar(&node, "node", "", "the node that trains")
	f.IntVar(&gpuMem, "gpu-memory", 4096, "MiB of a card the job needs free")
	f.StringVar(&data, "data", "", "the pairs, NDJSON of {\"q\", \"a\"}")
	f.IntVar(&epochs, "epochs", 6, "passes over the pairs")
	f.IntVar(&rank, "rank", 16, "the adapter's rank")
	f.Float64Var(&lr, "lr", 2e-4, "the learning rate")
	f.BoolVar(&noGGUF, "no-gguf", false, "skip the llama.cpp GGUF, even for a model llama.cpp serves")
	f.StringVar(&owner, "owner", "", "who the adapter belongs to — org:, team:, project: (path or ID), user: or key: — and only they see it "+
		"(default: the knowledge base's owner, when the pairs come from one)")
	f.BoolVar(&public, "public", false, "every caller of the model sees it; refused for pairs from a knowledge base")
	for _, req := range []string{"model", "base-repo", "node", "data"} {
		_ = start.MarkFlagRequired(req)
	}
	c.AddCommand(start)

	c.AddCommand(&cobra.Command{
		Use:   "list",
		Short: "Training jobs, newest first",
		Args:  cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			var jobs []adminapi.TrainJobView
			if err := client().get(cmd.Context(), "/api/v1/train/jobs", &jobs); err != nil {
				return err
			}
			if *asJSON {
				return json.NewEncoder(cmd.OutOrStdout()).Encode(jobs)
			}
			if len(jobs) == 0 {
				fmt.Fprintln(cmd.OutOrStdout(), "No training jobs (llmgw ctl train start).")
				return nil
			}
			w := tab()
			fmt.Fprintln(w, "JOB\tADAPTER\tMODEL\tOWNER\tNODE\tSTATE\tPROGRESS\tLOSS\tPEAK\tCREATED")
			for _, j := range jobs {
				fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\n", j.ID, j.Name, j.Model, or(j.Owner, "everyone"), j.Node, j.State,
					trainProgress(j), trainLoss(j), trainPeak(j), j.CreatedAt.Local().Format("01-02 15:04"))
			}
			return w.Flush()
		},
	})

	var follow bool
	show := &cobra.Command{
		Use:   "show JOB",
		Short: "How a training job goes, and what it made",
		Args:  cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			out := cmd.OutOrStdout()
			last := ""
			for {
				var j adminapi.TrainJobView
				if err := client().get(cmd.Context(), "/api/v1/train/jobs/"+args[0], &j); err != nil {
					return err
				}
				if *asJSON {
					return json.NewEncoder(out).Encode(j)
				}
				line := fmt.Sprintf("%s  %s  %s  loss %s  peak %s  %s", j.State, trainProgress(j), j.Node, trainLoss(j), trainPeak(j), j.Message)
				if !follow {
					fmt.Fprintf(out, "%s: adapter %s for %s, owned by %s, on %s (base %s), %d pairs\n", j.ID, j.Name, j.Model,
						or(j.Owner, "everyone"), j.Node, j.BaseRepo, j.Pairs)
					fmt.Fprintln(out, line)
					if len(j.Artifacts) > 0 {
						fmt.Fprintf(out, "made: %s (llmgw ctl train download %s --kind KIND)\n", strings.Join(j.Artifacts, ", "), j.ID)
					}
					return nil
				}
				if line != last {
					fmt.Fprintf(out, "%s  %s\n", time.Now().Format("15:04:05"), line)
					last = line
				}
				if store.TrainOver(j.State) {
					if j.State != store.TrainDone {
						return fmt.Errorf("%s %s", j.ID, j.State)
					}
					return nil
				}
				select {
				case <-cmd.Context().Done():
					return nil
				case <-time.After(3 * time.Second):
				}
			}
		},
	}
	show.Flags().BoolVarP(&follow, "follow", "f", false, "follow it until it ends")
	c.AddCommand(show)

	c.AddCommand(&cobra.Command{
		Use:   "cancel JOB",
		Short: "Stop a training job",
		Args:  cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			if err := client().do(cmd.Context(), "POST", "/api/v1/train/jobs/"+args[0]+"/cancel", nil, nil); err != nil {
				return err
			}
			fmt.Fprintf(cmd.OutOrStdout(), "%s cancelled\n", args[0])
			return nil
		},
	})
	c.AddCommand(&cobra.Command{
		Use:   "rm JOB",
		Short: "Remove an ended training job and what it made",
		Args:  cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			if err := client().do(cmd.Context(), "DELETE", "/api/v1/train/jobs/"+args[0], nil, nil); err != nil {
				return err
			}
			fmt.Fprintf(cmd.OutOrStdout(), "%s removed\n", args[0])
			return nil
		},
	})

	var restart bool
	publish := &cobra.Command{
		Use:   "publish JOB",
		Short: "Make what a finished job made its model's adapter, kept on the mirror",
		Long: "The GGUF for a model llama.cpp serves, the PEFT files for one vLLM serves, are kept\n" +
			"on the mirror and added to the model's adapters; a request then asks for the\n" +
			"adapter by its name, as a model. llama.cpp loads adapters when a replica starts:\n" +
			"--restart replaces the model's replicas, one at a time, to load it.",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			var v adminapi.TrainPublishView
			if err := client().patient(15*time.Minute).do(cmd.Context(), "POST", "/api/v1/train/jobs/"+args[0]+"/publish", nil, &v); err != nil {
				return err
			}
			out := cmd.OutOrStdout()
			fmt.Fprintf(out, "%s is %s's adapter, from %s\n", v.Adapter, v.Model, v.Source)
			if v.Owner != "" {
				fmt.Fprintf(out, "only %s sees it; to anyone else it is a model that does not exist\n", v.Owner)
			}
			if restart && strings.Contains(v.Note, "restart") {
				body, _ := json.Marshal(adminapi.RestartRequest{Models: []string{v.Model}})
				var res adminapi.RestartResult
				if err := client().post(cmd.Context(), "/api/v1/restart", body, &res); err != nil {
					return err
				}
				fmt.Fprintf(out, "%d replica(s) of %s restarting to load it; then ask for model %s\n", len(res.Queued), v.Model, v.Adapter)
				return nil
			}
			fmt.Fprintln(out, v.Note)
			return nil
		},
	}
	publish.Flags().BoolVar(&restart, "restart", false, "restart the model's replicas to load it (llama.cpp)")
	c.AddCommand(publish)

	var kind, outFile string
	download := &cobra.Command{
		Use:   "download JOB",
		Short: "Save what a training job made: the PEFT adapter (peft) or the GGUF (gguf)",
		Args:  cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			if outFile == "" {
				outFile = args[0] + ".tar.gz"
				if kind == "gguf" {
					outFile = args[0] + ".gguf"
				}
			}
			f, err := os.Create(outFile)
			if err != nil {
				return err
			}
			defer f.Close()
			cw := &countWriter{w: f}
			if err := client().patient(10*time.Minute).do(cmd.Context(), "GET",
				"/api/v1/train/jobs/"+args[0]+"/artifacts/"+kind, nil, io.Writer(cw)); err != nil {
				_ = os.Remove(outFile)
				return err
			}
			fmt.Fprintf(cmd.OutOrStdout(), "%s: %s\n", outFile, humanBytes(cw.n))
			return nil
		},
	}
	download.Flags().StringVar(&kind, "kind", "peft", "peft (a gzipped tar of the PEFT adapter) or gguf")
	download.Flags().StringVarP(&outFile, "out", "o", "", "the file to write")
	c.AddCommand(download)
	return c
}

type countWriter struct {
	w io.Writer
	n int64
}

func (c *countWriter) Write(p []byte) (int, error) {
	n, err := c.w.Write(p)
	c.n += int64(n)
	return n, err
}

func trainProgress(j adminapi.TrainJobView) string {
	if j.Steps == 0 {
		return "-"
	}
	return fmt.Sprintf("%d/%d", j.Step, j.Steps)
}

func trainLoss(j adminapi.TrainJobView) string {
	if j.Loss == 0 {
		return "-"
	}
	return fmt.Sprintf("%.3f", j.Loss)
}

func trainPeak(j adminapi.TrainJobView) string {
	if j.PeakMiB == 0 {
		return "-"
	}
	return fmt.Sprintf("%d MiB", j.PeakMiB)
}