llmgw/internal/control/modelserve/plan.go

package modelserve

import (
	"context"
	"errors"
	"fmt"
	"strings"

	"gitea.mynodes.xyz/ilker/llmgw/internal/gguf"
	"gitea.mynodes.xyz/ilker/llmgw/internal/hf"
	"gitea.mynodes.xyz/ilker/llmgw/internal/modelplan"
	"gitea.mynodes.xyz/ilker/llmgw/internal/strata"
)

// Planning: which of a repository's files suits the cluster's hardware, and
// how — one card, several of a node's, several nodes' over a fabric, or a
// CPU node — worked out from the files themselves before any is fetched.
// The model's shape is read from one GGUF file's header, a few megabytes
// from the registry, or from its config.json; the hardware is the fleet as
// it is now.

// PlanRequest is what to plan.
type PlanRequest struct {
	Repo   string `json:"repo"`
	Format string `json:"format,omitempty"` // gguf (default) or safetensors
	// Context is the token pool: llama.cpp's --ctx-size, or vLLM's
	// --max-model-len. 0 takes 16384 and 8192.
	Context int `json:"context,omitempty"`
}

// PlannedArtifact is one file a repository offers, and the best way to run it.
type PlannedArtifact struct {
	modelplan.Artifact
	Best *modelplan.Option `json:"best,omitempty"`
}

// Plan is every way a repository's files could run here, best first.
type Plan struct {
	Repo     string `json:"repo"`
	Revision string `json:"revision"`
	Format   string `json:"format"`
	Arch     string `json:"architecture,omitempty"`
	License  string `json:"license,omitempty"`
	Params   int64  `json:"params,omitempty"`
	Context  int    `json:"context"`
	// KVBytesPerToken is one token of context, the whole model: what an
	// option needs for another context is worked out from it. 0 is unknown.
	KVBytesPerToken int64 `json:"kv_bytes_per_token,omitempty"`
	// Experts is a mixture-of-experts model's routed experts; 0 for dense.
	Experts int  `json:"experts,omitempty"`
	Gated   bool `json:"gated,omitempty"`
	// Servable is false when the nodes' engine cannot run the model at all,
	// and Reason says why.
	Servable bool   `json:"servable"`
	Reason   string `json:"reason,omitempty"`
	// Note is anything the plan could not know, said once.
	Note       string             `json:"note,omitempty"`
	Projectors []string           `json:"projectors,omitempty"`
	Artifacts  []PlannedArtifact  `json:"artifacts"`
	Options    []modelplan.Option `json:"options"`
}

// Recommended is the best option of all, if any runs.
func (p Plan) Recommended() (modelplan.Option, bool) {
	if len(p.Options) == 0 {
		return modelplan.Option{}, false
	}
	return p.Options[0], true
}

// ggufSets are a repository's GGUF weights: models, not projectors or
// importance matrices.
func ggufSets(repo hf.Repo) []hf.FileSet {
	var weights []hf.File
	for _, f := range repo.Files {
		lower := strings.ToLower(f.Path)
		if strings.HasSuffix(lower, ".gguf") && !strings.Contains(lower, "mmproj") && !strings.Contains(lower, "imatrix") {
			weights = append(weights, f)
		}
	}
	return hf.Group(weights)
}

// Plan works out where a repository's files could run.
func (s *Server) Plan(ctx context.Context, req PlanRequest) (Plan, error) {
	if s.deps.Pool == nil {
		return Plan{}, errors.New("this control plane cannot see its fleet here")
	}
	req.Repo = strings.Trim(strings.TrimSpace(req.Repo), "/")
	if !repoPattern.MatchString(req.Repo) {
		return Plan{}, fmt.Errorf("%q is not a repository; it looks like owner/name", req.Repo)
	}
	repo, err := s.deps.Registry.Repo(ctx, req.Repo)
	if err != nil {
		return Plan{}, err
	}
	return s.plan(ctx, repo, req)
}

func (s *Server) plan(ctx context.Context, repo hf.Repo, req PlanRequest) (Plan, error) {
	// A pack Strata fetches itself: one way, Strata's, whatever its files.
	if fam, size, ok := strata.ForFetchedRepo(repo.ID); ok {
		art := modelplan.Artifact{Name: size.Name, Format: "strata", Bytes: int64(size.DownloadGB * 1e9), Quant: size.Name}
		p := Plan{Repo: repo.ID, Revision: repo.Revision, Format: "strata", Arch: repo.Architecture, License: repo.License,
			Context: strataContext, Servable: true, Note: "Strata fetches this pack onto the node itself on its first start"}
		p.Options = modelplan.StrataOptions(fam.Name, []modelplan.Artifact{art}, s.deps.Pool())
		if len(p.Options) == 0 && fam.MinCC > 0 {
			p.Note = fmt.Sprintf("Strata's %s engine needs an NVIDIA card of compute capability %.1f or newer (RTX 30 on); "+
				"a node with Strata and such a card, and %d GB of RAM, runs it", fam.Title, fam.MinCC, size.MemoryGB())
		}
		pa := PlannedArtifact{Artifact: art}
		if len(p.Options) > 0 {
			pa.Best = &p.Options[0]
		}
		p.Artifacts = []PlannedArtifact{pa}
		return p, nil
	}
	p := Plan{Repo: repo.ID, Revision: repo.Revision, Format: req.Format, Arch: repo.Architecture,
		License: repo.License, Params: repo.Params, Context: req.Context, Gated: repo.Gated.Is(), Servable: true,
		Projectors: Projectors(repo)}
	var arts []modelplan.Artifact
	var shape modelplan.Shape
	switch req.Format {
	case "", "gguf":
		p.Format = "gguf"
		if p.Context <= 0 {
			p.Context = 16384
		}
		sets := ggufSets(repo)
		if len(sets) == 0 {
			return p, fmt.Errorf("%s has no GGUF files", repo.ID)
		}
		small := sets[0]
		for _, st := range sets {
			arts = append(arts, modelplan.Artifact{Name: st.Name, Format: "gguf", Bytes: st.Bytes,
				Shards: st.Shards, Quant: modelplan.QuantOf(st.Name)})
			if st.Bytes < small.Bytes {
				small = st
			}
		}
		// Every quantisation of a model has its shape: the smallest file's
		// header says it, read from the registry without fetching the file.
		base := ""
		if b, ok := s.deps.Registry.(interface{ Base() string }); ok {
			base = b.Base()
		}
		v := gguf.Check(ctx, repo.DownloadURL(base, small.Files[0].Path), s.deps.LlamaVersion)
		if v.Checked && !v.Servable {
			p.Servable, p.Reason = false, v.Reason
		}
		shape.KVBytesPerToken = v.Header.KVBytesPerToken()
		shape.KVHeads = int(v.Header.HeadCountKV) //nolint:gosec // a header's head count
		if p.Arch == "" {
			p.Arch = v.Header.Architecture
		}
	case "safetensors":
		if p.Context <= 0 {
			p.Context = 8192
		}
		set, err := pickSafetensors(repo, Request{})
		if err != nil {
			return p, err
		}
		arts = append(arts, modelplan.Artifact{Name: "original", Format: "safetensors", Bytes: set.Bytes,
			Shards: len(set.Files), Quant: "bf16"})
		m, err := modelplan.FetchModel(ctx, repo.ID)
		if err == nil {
			shape.KVBytesPerToken = m.KVBytesPerToken(modelplan.BF16)
			shape.KVHeads = m.KVHeads
			shape.Experts = m.Experts
			if p.Params == 0 {
				p.Params = m.Params
			}
		} else {
			p.Note = "its config.json could not be read: " + err.Error()
		}
	default:
		return p, fmt.Errorf("the format is gguf or safetensors, not %q", req.Format)
	}
	if shape.KVBytesPerToken == 0 && p.Note == "" {
		p.Note = "the model's shape was not readable, so the context's memory is not counted"
	}
	shape.Params = p.Params
	p.KVBytesPerToken, p.Experts = shape.KVBytesPerToken, shape.Experts
	p.Options = modelplan.Options(arts, shape, s.deps.Pool(), p.Context)
	if !p.Servable {
		p.Options = nil
	}
	// A family Strata runs: its files on one card with the experts in RAM.
	// Strata reads them itself, whether or not llama.cpp could.
	if fam, ok := strata.ForRepo(repo.ID); ok && p.Format == "gguf" {
		if so := modelplan.StrataOptions(fam.Name, arts, s.deps.Pool()); len(so) > 0 {
			p.Options = append(p.Options, so...)
			modelplan.Rank(p.Options)
			if !p.Servable {
				p.Note = strings.TrimSpace(p.Note + " llama.cpp cannot run it (" + p.Reason + "); Strata can.")
				p.Servable, p.Reason = true, ""
			}
		}
	}
	for _, a := range arts {
		pa := PlannedArtifact{Artifact: a}
		for i := range p.Options {
			if p.Options[i].Artifact == a.Name {
				pa.Best = &p.Options[i]
				break
			}
		}
		p.Artifacts = append(p.Artifacts, pa)
	}
	return p, nil
}