llmgw/internal/cli/modelplan.go

package cli

import (
	"fmt"
	"strings"

	"github.com/spf13/cobra"

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

// newModelPlanCmd answers whether a model fits this cluster, and how.
//
// The arithmetic it replaces is the last thing standing between having GPUs and
// serving a model on them. Deciding tp, gpu_class and a context that will
// actually fit means knowing the model's layer count, its KV head count, what
// its weights weigh at a chosen quantisation, and how much of a card the
// runtime keeps — and getting it wrong is not an error message, it is an
// out-of-memory twenty minutes into a load.
//
// It reads the model from its own repository and the hardware from the running
// fleet, so neither half is a number somebody retyped.
func newModelPlanCmd(client func() *ctlClient) *cobra.Command {
	var (
		context int
		params  int64
		precs   []string
	)
	c := &cobra.Command{
		Use:   "plan <huggingface-repo>",
		Short: "Work out whether a model fits this cluster, and how it should be split",
		Long: "Reads the model's architecture from its repository and the hardware from\n" +
			"the running fleet, then prints every way it could run — fewest GPUs\n" +
			"first — and a catalog entry for the best one.\n\n" +
			"Every candidate is listed rather than one answer, because the choice is\n" +
			"not the planner's to make: with cards to spare you want the\n" +
			"full-precision fit, and with one card free you want to know a quantised\n" +
			"one exists.",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			repo := args[0]
			m, err := modelplan.FetchModel(cmd.Context(), repo)
			if err != nil {
				return err
			}
			if params > 0 {
				m.Params = params
			}
			if m.Params == 0 {
				return fmt.Errorf("%s does not publish a parameter count; pass --params", repo)
			}

			var chosen []modelplan.Precision
			for _, name := range precs {
				p, err := modelplan.PrecisionByName(name)
				if err != nil {
					return err
				}
				chosen = append(chosen, p)
			}

			var nodes []adminapi.NodeView
			if err := client().get(cmd.Context(), "/api/v1/nodes", &nodes); err != nil {
				return err
			}
			hw := hardwareFrom(nodes)
			if len(hw) == 0 {
				return fmt.Errorf("no node is ready, so there is nothing to plan against")
			}

			out := cmd.OutOrStdout()
			fmt.Fprintf(out, "%s: %.1fB parameters, %d layers, %d KV heads, trained to %d tokens\n",
				m.Name, float64(m.Params)/1e9, m.Layers, m.KVHeads, m.MaxContext)

			want := context
			if want <= 0 {
				want = m.MaxContext
			}
			fmt.Fprintf(out, "Planned for a %d-token pool. %s a token of context.\n\n",
				want, humanBytes(m.KVBytesPerToken(modelplan.BF16)))

			fits := modelplan.Fits(m, hw, want, chosen)
			if len(fits) == 0 {
				fmt.Fprintln(out, "It does not fit this cluster at any offered precision.")
				fmt.Fprintln(out, "\nWhat would change that:")
				fmt.Fprintln(out, "  - a smaller context pool, with --context")
				fmt.Fprintln(out, "  - a heavier quantisation, with --precision q4_k_m")
				fmt.Fprintln(out, "  - more cards free; llmgw ctl status shows what is in use")
				return errDoesNotFit
			}

			w := tab()
			fmt.Fprintln(w, "  GPUS\tPRECISION\tPER CARD\tFREE\tWHERE")
			for _, f := range fits {
				where := f.Class
				if f.CPUOnly {
					where = "cpu"
				} else if f.Nodes > 1 {
					where = fmt.Sprintf("%s, %d nodes on %s", f.Class, f.Nodes, f.Fabric)
				}
				fmt.Fprintf(w, "  %d\t%s\t%.1f GB\t%.1f GB\t%s\n",
					f.TP, f.Precision.Name, f.PerGPUGB, f.Headroom, where)
			}
			_ = w.Flush()

			fmt.Fprintf(out, "\nA catalog entry for the first of those:\n\n%s\n",
				catalogEntry(m, fits[0], repo, want))
			fmt.Fprintln(out, "Save it and add it with:  llmgw ctl model add -f <file>")
			return nil
		},
	}
	c.Flags().IntVar(&context, "context", 0, "token pool to plan for (default: the model's trained context)")
	c.Flags().Int64Var(&params, "params", 0, "parameter count, when the repository does not publish one")
	c.Flags().StringSliceVar(&precs, "precision", nil, "precisions to consider (default: bf16, fp8, q8_0, q4_k_m)")
	return c
}

// errDoesNotFit gives a script something to test without parsing prose.
var errDoesNotFit = fmt.Errorf("the model does not fit this cluster")

// hardwareFrom turns the fleet into pools a replica could actually occupy.
//
// Grouped by class and fabric together, because a replica may never span two
// fabrics: eight free cards split across two fabrics are not eight cards for a
// model that needs eight, and treating them as one pool would plan something
// that hangs on its first all-reduce.
func hardwareFrom(nodes []adminapi.NodeView) []modelplan.Hardware {
	type key struct{ class, fabric string }
	pools := map[key]*modelplan.Hardware{}

	for _, n := range nodes {
		if n.State != "ready" {
			continue
		}
		fabric := ""
		if len(n.Fabrics) > 0 {
			fabric = n.Fabrics[0]
		}
		class := ""
		if len(n.Classes) > 0 {
			class = n.Classes[0]
		}
		k := key{class, fabric}
		if pools[k] == nil {
			pools[k] = &modelplan.Hardware{Class: class, Fabric: fabric}
		}
		p := pools[k]
		p.Nodes++
		// The widest a single node could contribute, not the total: a tensor
		// group takes the same number of cards from every node it spans.
		if n.GPUFree > p.GPUsPerNode {
			p.GPUsPerNode = n.GPUFree
		}
		if n.MemoryGB > p.HostMemGB {
			p.HostMemGB = n.MemoryGB
		}
		// Card memory from the class label, which carries it by construction:
		// "h100-80g" is eighty gigabytes and says so.
		if class != "" && p.GPUMemGB == 0 {
			p.GPUMemGB = memFromClass(class)
		}
	}

	out := make([]modelplan.Hardware, 0, len(pools))
	for _, p := range pools {
		if p.Class != "" && (p.GPUMemGB == 0 || p.GPUsPerNode == 0) {
			continue // a GPU pool with no free cards is not a pool
		}
		out = append(out, *p)
	}
	return out
}

// memFromClass reads the card size out of a class label such as "h100-80g".
func memFromClass(class string) int {
	i := strings.LastIndex(class, "-")
	if i < 0 || !strings.HasSuffix(class, "g") {
		return 0
	}
	var gb int
	if _, err := fmt.Sscanf(class[i+1:], "%dg", &gb); err != nil {
		return 0
	}
	return gb
}

// catalogEntry renders a fit as the YAML that would serve it.
func catalogEntry(m modelplan.Model, f modelplan.Fit, source string, context int) string {
	var b strings.Builder
	name := strings.ToLower(m.Name)
	fmt.Fprintf(&b, "name: %s\n", name)
	if f.Precision.GGUF {
		// A GGUF quantisation lives in a file, not a repository, and llama.cpp
		// is given a path. Printing the repository id here would produce an
		// entry that looks right and cannot load — which is the failure this
		// whole command exists to prevent.
		fmt.Fprintf(&b, "# %s is a GGUF quantisation: download the %s file and give its path.\n",
			f.Precision.Name, f.Precision.Name)
		fmt.Fprintf(&b, "source: /var/lib/llmgw-models/%s-%s.gguf\n", name, f.Precision.Name)
	} else {
		fmt.Fprintf(&b, "source: %s\n", source)
	}
	fmt.Fprintf(&b, "routing_profile: chat\n")
	fmt.Fprintf(&b, "context_length: %d\n", context)
	fmt.Fprintf(&b, "scaling: { min: 1, max: 1 }\n")
	fmt.Fprintf(&b, "profiles:\n")

	switch {
	case f.CPUOnly:
		fmt.Fprintf(&b, "  - name: cpu\n    engine: %s\n", f.Precision.Engine)
		fmt.Fprintf(&b, "    min_memory_gb: %d\n", int(f.PerGPUGB)+2)
	default:
		fmt.Fprintf(&b, "  - name: tp%d\n    engine: %s\n", f.TP, f.Precision.Engine)
		fmt.Fprintf(&b, "    gpu_class: %s\n    tp: %d\n", f.Class, f.TP)
		if f.Nodes > 1 {
			fmt.Fprintf(&b, "    nodes: %d\n    gpus_per_node: %d\n", f.Nodes, f.GPUsPerNode)
		}
		if q := f.Precision.Quantization; q != "" {
			fmt.Fprintf(&b, "    args: [\"--quantization\", \"%s\"]\n", q)
		}
	}
	fmt.Fprintf(&b, "    priority: 100\n    weight: 1\n")
	// Measured rather than guessed is the rule everywhere else in this system,
	// and a planner should say so rather than pretend its estimate is one.
	fmt.Fprintf(&b, "    startup_timeout: 10m\n")
	b.WriteString("\n# weight should come from `llmgw ctl bench` once it is running;\n")
	b.WriteString("# 1 is a placeholder, not a measurement.\n")
	return b.String()
}

// humanBytes renders a per-token cache size, which is always small.
func humanBytes(n int64) string {
	switch {
	case n >= 1<<20:
		return fmt.Sprintf("%.1f MiB", float64(n)/(1<<20))
	default:
		return fmt.Sprintf("%.0f KiB", float64(n)/1024)
	}
}