llmgw/internal/control/modelserve/modelserve.go

// Package modelserve takes a model from a registry repository to a replica
// answering requests, as one job.
//
// Adding a model was four commands and a hand-written YAML file: plan it,
// fetch it into the mirror, write its catalog entry with a memory figure
// somebody guessed, scale it — and then find out whether it loads. This runs
// those steps in the control plane, where the mirror, the registry, the
// catalog and the fleet already are, and reports each one as it goes, so
// whatever asks — ctl, a page — only has to watch.
package modelserve

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"regexp"
	"slices"
	"strconv"
	"strings"
	"sync"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/canary"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
	"gitea.mynodes.xyz/ilker/llmgw/internal/gguf"
	"gitea.mynodes.xyz/ilker/llmgw/internal/guard"
	"gitea.mynodes.xyz/ilker/llmgw/internal/hf"
	"gitea.mynodes.xyz/ilker/llmgw/internal/mirror"
	"gitea.mynodes.xyz/ilker/llmgw/internal/modelplan"
	"gitea.mynodes.xyz/ilker/llmgw/internal/notify"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// Registry reads a repository.
type Registry interface {
	Repo(ctx context.Context, id string) (hf.Repo, error)
}

// Mirror fetches files and reports on them.
type Mirror interface {
	Fetch(ctx context.Context, req mirror.FetchRequest) (mirror.FetchResponse, error)
	List(ctx context.Context) (mirror.ModelsResponse, error)
}

// Catalog adds a model.
type Catalog interface {
	Add(ctx context.Context, m catalog.Model, by string) error
}

// Fleet is what the job watches once the model is in the catalog.
type Fleet interface {
	Replicas() []reconciler.Replica
	ModelHealth() []reconciler.ModelHealth
	EngineKey(replicaID string) string
}

// Deps are what a job needs.
type Deps struct {
	Registry Registry
	Mirror   Mirror
	Catalog  Catalog
	Fleet    Fleet
	// LlamaVersion is the llama.cpp build the nodes install, which decides
	// which architectures load.
	LlamaVersion string
	HTTP         *http.Client
	Log          *slog.Logger
	// Poll is how often the job looks at the mirror and the fleet.
	Poll time.Duration
	// ReadyTimeout bounds the wait for a replica: the node fetching the
	// weights from the mirror and the engine loading them.
	ReadyTimeout time.Duration

	// For updates: the catalog in force, the service that changes it, and
	// what replaces replicas. Without them updates are not offered.
	Current func() []catalog.Model
	Models  Models
	Restart func(ids ...string) error

	// Pool is the fleet's ready nodes as a plan sees them: what decides
	// which file suits the hardware, and how it is spread.
	Pool func() []modelplan.PoolNode
	// UnplacedGrace is how long a model may go unplaced before the job
	// stops waiting for it: a minute unless set.
	UnplacedGrace time.Duration
	// Notify is told when a model's repository publishes new weights or a
	// canary is rolled back; nil tells no one.
	Notify *notify.Notifier
	// CanaryStats is how each model's canary and stable side served, by
	// model; ResetCanary forgets one. From the router; without them a
	// canary is never rolled back by itself.
	CanaryStats func() map[string]canary.Stat
	ResetCanary func(model string)
	// Eval runs the quality suites an update or a promotion is held to;
	// without it, one held to a suite is refused.
	Eval Evaluator
}

// errBusy is a model the cluster has the hardware for, all of it in use now.
var errBusy = errors.New("the hardware it fits is in use")

// Request is what to serve.
type Request struct {
	// Repo is the registry repository, owner/name.
	Repo string `json:"repo"`
	// Quant picks the file by its quantisation, Q4_K_M unless given.
	Quant string `json:"quant,omitempty"`
	// File picks it by name instead.
	File string `json:"file,omitempty"`
	// Projector is a vision model's mmproj file in the repository, fetched
	// beside the weights and loaded with them; empty serves text alone.
	Projector string `json:"projector,omitempty"`
	// Format is "gguf", run by llama.cpp, or "safetensors": the original
	// weights with their config and tokenizer, run by vLLM on GPUs of
	// GPUClass, across TP of them.
	Format string `json:"format,omitempty"`
	// Where it runs: on the CPU, or on GPUsPerNode cards of GPUClass on
	// each of Nodes nodes. Left empty, the plan decides: the best way the
	// cluster can run the chosen file. TP is GPUsPerNode's older name.
	CPU         bool   `json:"cpu,omitempty"`
	GPUClass    string `json:"gpu_class,omitempty"`
	GPUsPerNode int    `json:"gpus_per_node,omitempty"`
	Nodes       int    `json:"nodes,omitempty"`
	TP          int    `json:"tp,omitempty"`
	// Fallbacks adds the next best ways, lower in priority, so a model whose
	// first choice is busy runs on another rather than waiting.
	Fallbacks bool `json:"fallbacks,omitempty"`
	// Name is the catalog name; the repository suggests one otherwise.
	Name        string   `json:"name,omitempty"`
	Aliases     []string `json:"aliases,omitempty"`
	Description string   `json:"description,omitempty"`
	// Replicas to run; 0 makes it on demand, loaded when first asked for.
	Replicas int `json:"replicas"`
	// GuardFormat makes it a judge model for the guardrails (role: guard):
	// llamaguard3, qwen3guard, shieldgemma or json. Empty serves clients.
	GuardFormat string `json:"guard_format,omitempty"`
	// Class is what it may take from others and they from it when there is
	// not room for both: critical, standard or opportunistic. Empty is the
	// catalog's default for its role.
	Class string `json:"class,omitempty"`
	// CtxPool is llama.cpp's --ctx-size; ContextLength what one request may
	// use. They default to 16384 and 4096, what four parallel slots divide;
	// llama.cpp is given as many slots as the pool holds requests.
	CtxPool       int `json:"ctx_pool,omitempty"`
	ContextLength int `json:"context_length,omitempty"`
	// Tools says the model makes tool calls the engine returns as calls: it
	// is written into the catalog, where clients and editors read it. Check
	// a running replica with llmgw ctl tool-check before saying so.
	Tools bool `json:"tools,omitempty"`
	// Embeddings serves it as an embedding model, which its GGUF header
	// says by itself for most: --embedding, many slots, a batch that holds
	// a whole passage, /v1/embeddings.
	Embeddings bool `json:"embeddings,omitempty"`
	// sized is the caller having said --ctx-pool or --context-length, which
	// an embedding model's defaults then leave alone.
	sized bool
	// CPUCores is what a replica computes on; zero takes half the smallest
	// node it fits on.
	CPUCores int `json:"cpu_cores,omitempty"`
	// By is who asked, for the catalog's record.
	By string `json:"-"`
}

// The states a job and its steps can be in.
const (
	Running = "running"
	Done    = "done"
	Failed  = "failed"
	Skipped = "skipped"
)

// Step is one stage of a job.
type Step struct {
	Name   string    `json:"name"`
	State  string    `json:"state"`
	Detail string    `json:"detail,omitempty"`
	At     time.Time `json:"at"`
}

// Job is one model being served, and where it stands.
type Job struct {
	ID string `json:"id"`
	// Kind is "update" for a model moved to new weights, empty for one
	// served.
	Kind    string    `json:"kind,omitempty"`
	Request Request   `json:"request"`
	State   string    `json:"state"`
	Error   string    `json:"error,omitempty"`
	Model   string    `json:"model,omitempty"`
	Steps   []Step    `json:"steps"`
	Started time.Time `json:"started"`
	// Reply is what the model said to the smoke test, and TokensPerSecond
	// how fast it said it.
	Reply           string  `json:"reply,omitempty"`
	TokensPerSecond float64 `json:"tokens_per_second,omitempty"`

	// trace is the job's spans, when tracing is on.
	trace jobTrace
	// Eval is the quality suite the new weights were asked before they
	// replaced the old, and how they scored, without each case.
	Eval *store.EvalRun `json:"eval,omitempty"`
}

// Server runs jobs and remembers them.
type Server struct {
	deps Deps
	mu   sync.Mutex
	jobs map[string]*Job
	seq  int
	// updates is what the last checks of the registry found, by model.
	updates map[string]UpdateInfo
	// origins are those worked out for models whose definition does not
	// keep one, nil for a model that could not be traced.
	origins map[string]*catalog.Origin
}

// New builds a server.
func New(d Deps) *Server {
	if d.Poll <= 0 {
		d.Poll = 2 * time.Second
	}
	if d.ReadyTimeout <= 0 {
		d.ReadyTimeout = 30 * time.Minute
	}
	if d.UnplacedGrace <= 0 {
		d.UnplacedGrace = time.Minute
	}
	if d.HTTP == nil {
		d.HTTP = &http.Client{Timeout: 5 * time.Minute}
	}
	if d.Log == nil {
		d.Log = slog.Default()
	}
	return &Server{deps: d, jobs: map[string]*Job{}, updates: map[string]UpdateInfo{}, origins: map[string]*catalog.Origin{}}
}

var repoPattern = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._-]*/[A-Za-z0-9][A-Za-z0-9._-]*$`)

// Start begins a job and returns at once; the job runs on its own.
func (s *Server) Start(req Request) (Job, error) {
	req.Repo = strings.Trim(strings.TrimSpace(req.Repo), "/")
	if !repoPattern.MatchString(req.Repo) {
		return Job{}, fmt.Errorf("%q is not a repository; it looks like owner/name", req.Repo)
	}
	if req.Replicas < 0 {
		return Job{}, errors.New("replicas cannot be negative")
	}
	switch req.Format {
	case "", "gguf":
		req.Format = "gguf"
	case "safetensors":
		if req.Projector != "" {
			return Job{}, errors.New("a projector is for a GGUF model; vLLM reads a vision model's own files")
		}
		if req.CPU {
			return Job{}, errors.New("original weights run on vLLM, on GPUs; a GGUF file runs on the CPU")
		}
	default:
		return Job{}, fmt.Errorf("the format is gguf or safetensors, not %q", req.Format)
	}
	if req.GuardFormat != "" && !guard.KnownFormat(req.GuardFormat) {
		return Job{}, fmt.Errorf("the guard format is one of %s, not %q", strings.Join(guard.Formats, ", "), req.GuardFormat)
	}
	switch req.Class {
	case "", catalog.ClassCritical, catalog.ClassStandard, catalog.ClassOpportunistic:
	default:
		return Job{}, fmt.Errorf("the class is critical, standard or opportunistic, not %q", req.Class)
	}
	if req.GuardFormat != "" && req.Replicas == 0 {
		return Job{}, errors.New("a judge model must always run: give it --replicas 1 or more, not --on-demand")
	}
	req.sized = req.CtxPool > 0 || req.ContextLength > 0
	if req.CtxPool <= 0 {
		req.CtxPool = 16384
	}
	if req.ContextLength <= 0 {
		req.ContextLength = 4096
	}
	if req.ContextLength > req.CtxPool {
		return Job{}, fmt.Errorf("a request of %d tokens cannot fit a pool of %d: give --ctx-pool at least --context-length",
			req.ContextLength, req.CtxPool)
	}
	s.mu.Lock()
	s.seq++
	j := &Job{ID: fmt.Sprintf("serve-%d-%d", time.Now().Unix(), s.seq), Request: req,
		State: Running, Started: time.Now()}
	s.keepJob(j)
	s.mu.Unlock()
	go s.run(j)
	return s.snapshot(j), nil
}

// keptJobs is how many finished jobs are remembered: enough for the console's
// list, and not one more for every model ever served.
const keptJobs = 100

// keepJob records a job, and forgets the oldest finished ones beyond
// keptJobs. Called with s.mu held.
func (s *Server) keepJob(j *Job) {
	s.jobs[j.ID] = j
	var done []*Job
	for _, o := range s.jobs {
		if o.State != Running {
			done = append(done, o)
		}
	}
	if len(done) <= keptJobs {
		return
	}
	slices.SortFunc(done, func(a, b *Job) int { return a.Started.Compare(b.Started) })
	for _, o := range done[:len(done)-keptJobs] {
		delete(s.jobs, o.ID)
	}
}

// Jobs is the jobs since the control plane started — every one running, and
// the last keptJobs finished — the newest first.
func (s *Server) Jobs() []Job {
	s.mu.Lock()
	all := make([]*Job, 0, len(s.jobs))
	for _, j := range s.jobs {
		all = append(all, j)
	}
	s.mu.Unlock()
	out := make([]Job, 0, len(all))
	for _, j := range all {
		out = append(out, s.snapshot(j))
	}
	slices.SortFunc(out, func(a, b Job) int { return b.Started.Compare(a.Started) })
	return out
}

// Get returns a job as it stands.
func (s *Server) Get(id string) (Job, bool) {
	s.mu.Lock()
	j, ok := s.jobs[id]
	s.mu.Unlock()
	if !ok {
		return Job{}, false
	}
	return s.snapshot(j), true
}

func (s *Server) snapshot(j *Job) Job {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := *j
	out.Steps = slices.Clone(j.Steps)
	return out
}

// step starts a new step, or updates the detail of the current one.
func (s *Server) step(j *Job, name, state, detail string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	j.traceStep(name, state, detail)
	if n := len(j.Steps); n > 0 && j.Steps[n-1].Name == name {
		j.Steps[n-1].State, j.Steps[n-1].Detail, j.Steps[n-1].At = state, detail, time.Now()
		return
	}
	j.Steps = append(j.Steps, Step{Name: name, State: state, Detail: detail, At: time.Now()})
}

func (s *Server) fail(j *Job, name string, err error) {
	s.step(j, name, Failed, err.Error())
	s.mu.Lock()
	j.State, j.Error = Failed, err.Error()
	j.traceEnd(err)
	s.mu.Unlock()
	s.deps.Log.Warn("serving a model failed", "job", j.ID, "repo", j.Request.Repo, "step", name, "err", err)
}

func (s *Server) run(j *Job) {
	ctx := context.Background()
	req := j.Request

	// 1. The registry: which files.
	s.step(j, "registry", Running, "reading "+req.Repo)
	repo, err := s.deps.Registry.Repo(ctx, req.Repo)
	if err != nil {
		s.fail(j, "registry", err)
		return
	}
	if repo.Gated.Is() {
		s.fail(j, "registry", fmt.Errorf("%s is gated: it needs an account's approval the mirror does not have", req.Repo))
		return
	}
	if req.TP > 0 && req.GPUsPerNode == 0 {
		req.GPUsPerNode = req.TP
	}
	// The plan: which file suits the hardware, and how it spreads over it.
	planCtx := req.CtxPool
	if req.Format == "safetensors" {
		planCtx = req.ContextLength
	}
	var plan Plan
	planned := false
	if s.deps.Pool != nil {
		if plan, err = s.plan(ctx, repo, PlanRequest{Repo: repo.ID, Format: req.Format, Context: planCtx}); err == nil {
			planned = true
		} else {
			s.deps.Log.Warn("planning a model", "repo", repo.ID, "err", err)
		}
	}
	if req.Format == "gguf" && req.Quant == "" && req.File == "" {
		if rec, ok := plan.Recommended(); planned && ok {
			req.File = rec.Artifact
		} else {
			req.Quant = "Q4_K_M"
		}
	}
	pickSet := pick
	if req.Format == "safetensors" {
		pickSet = pickSafetensors
	}
	set, err := pickSet(repo, req)
	if err != nil {
		s.fail(j, "registry", err)
		return
	}
	layout, err := chooseLayout(req, plan, planned, set)
	if err != nil {
		s.fail(j, "registry", err)
		return
	}
	if req.Projector != "" {
		i := slices.IndexFunc(repo.Files, func(f hf.File) bool {
			return (f.Path == req.Projector || mirror.LastPathElement(f.Path) == req.Projector) &&
				strings.Contains(strings.ToLower(f.Path), "mmproj")
		})
		if i < 0 {
			s.fail(j, "registry", fmt.Errorf("%s has no vision projector %q; it has: %s", repo.ID, req.Projector,
				strings.Join(Projectors(repo), ", ")))
			return
		}
		set.Files = append(set.Files, repo.Files[i])
		set.Bytes += repo.Files[i].Bytes
	}
	name := req.Name
	if name == "" {
		name = hf.ModelName(repo.ID)
	}
	s.mu.Lock()
	j.Model = name
	s.mu.Unlock()
	where := make([]string, 0, len(layout))
	for _, o := range layout {
		where = append(where, o.Label())
	}
	s.step(j, "registry", Done, fmt.Sprintf("%s, %.1f GB in %d file(s), revision %s; runs on %s",
		set.Name, gb(set.Bytes), len(set.Files), repo.ShortRevision(), strings.Join(where, ", else ")))

	// 2. The mirror: fetch, and check what arrived against what was published.
	source, err := s.mirror(ctx, j, repo, set, name)
	if err != nil {
		s.fail(j, "mirror", err)
		return
	}

	// 3. Whether the nodes' engine loads it, and how much memory it needs.
	if req.Format == "safetensors" {
		// vLLM reads the model's shape from its config.json, which came
		// with the weights.
		s.step(j, "check", Done, fmt.Sprintf("%.1f GB of weights for vLLM", gb(set.Bytes)))
		s.addAndStart(ctx, j, repo, set, name, source, profilesFor(layout, req, 0, 0), hardwareExists(layout))
		return
	}
	s.step(j, "check", Running, "reading its header from the mirror")
	var header gguf.Header
	if u := headerURL(source, set); u != "" {
		v := gguf.Check(ctx, u, s.deps.LlamaVersion)
		if !v.Servable {
			s.fail(j, "check", fmt.Errorf("llama.cpp %s on the nodes cannot run this: %s", s.deps.LlamaVersion, v.Reason))
			return
		}
		header = v.Header
	}
	if header.Embedding {
		req.Embeddings = true
	}
	if req.Embeddings && !req.sized {
		// Many passages at once rather than four: an import of a hundred
		// thousand chunks went at the pace of four slots with the card a
		// third busy.
		req.ContextLength, req.CtxPool = embedSlotTokens, embedSlotTokens*embedSlots
	}
	// On one card, the weights and the context's cache have to fit it: a
	// 70B with a 65536-token pool on a 24 GB card loaded its weights and then
	// failed to make its context, every time, as did an 8B with 131072.
	if len(layout) > 0 && layout[0].Class != "" && !layout[0].CPU {
		fitted, why, err := fitPoolToCard(req, header, set.Bytes, layout[0].Class, max(layout[0].GPUsPerNode, 1))
		if err != nil {
			s.fail(j, "check", err)
			return
		}
		if why != "" {
			s.step(j, "check", Running, why)
			req = fitted
		}
	}
	j.Request = req
	memGB := int((header.MemoryBytes(set.Bytes, req.CtxPool) + 1<<30 - 1) >> 30)
	basis := fmt.Sprintf("weights %.1f GB + KV cache %.1f GB for %d tokens + engine",
		gb(set.Bytes), gb(int64(req.CtxPool)*header.KVBytesPerToken()), req.CtxPool)
	if memGB == 0 {
		// A set of shards, or a header that did not say: weights plus half
		// again plus two, the rough figure from before shapes were read.
		memGB = int(gb(set.Bytes)*1.5) + 2
		basis = "weights plus half again plus two; the header gave no shape"
	}
	s.step(j, "check", Done, fmt.Sprintf("%s, needs about %d GB (%s)", orUnknown(header.Architecture), memGB, basis))

	cores := req.CPUCores
	if cores <= 0 {
		if nc, ok := s.deps.Fleet.(interface{ NodeCPUs(int) []int }); ok {
			cores = catalog.SuggestCores(nc.NodeCPUs(memGB))
		}
	}
	s.addAndStart(ctx, j, repo, set, name, source, profilesFor(layout, req, memGB, cores), hardwareExists(layout))
}

// chooseLayout is where a model runs: what was asked, else the plan's best
// for the chosen file; with Fallbacks, the next best of other shapes after.
func chooseLayout(req Request, plan Plan, planned bool, set hf.FileSet) ([]modelplan.Option, error) {
	artifact := set.Name
	engine := "llamacpp"
	if req.Format == "safetensors" {
		artifact, engine = "original", "vllm"
	}
	var opts []modelplan.Option
	for _, o := range plan.Options {
		if o.Artifact == artifact {
			opts = append(opts, o)
		}
	}
	var out []modelplan.Option
	switch {
	case req.CPU || req.GPUClass != "":
		o := modelplan.Option{Artifact: artifact, Format: req.Format, Engine: engine, CPU: req.CPU}
		if !req.CPU {
			o.Class, o.GPUsPerNode, o.Nodes = req.GPUClass, max(req.GPUsPerNode, 1), max(req.Nodes, 1)
		}
		if o.Nodes > 1 && engine == "llamacpp" {
			return nil, errors.New("llama.cpp does not span nodes; serve the original weights with vLLM for that")
		}
		// The plan's own reckoning of that way, if it has one: what it holds
		// and on how many hosts.
		for _, p := range opts {
			if p.CPU == o.CPU && p.Class == o.Class && p.GPUsPerNode == o.GPUsPerNode && max(p.Nodes, 1) == o.Nodes {
				o = p
			}
		}
		out = append(out, o)
	case len(opts) > 0:
		out = append(out, opts[0])
	case planned:
		return nil, fmt.Errorf("%s fits nowhere in this cluster with a %d-token context: no card holds it, "+
			"nor does any node's memory; `llmgw ctl model plan %s` shows what each file needs", set.Name, plan.Context, plan.Repo)
	case engine == "vllm":
		return nil, errors.New("original weights run on vLLM, on GPUs: name the GPU class")
	default:
		out = append(out, modelplan.Option{Artifact: artifact, Format: req.Format, Engine: engine, CPU: true})
	}
	if req.Fallbacks {
		shape := func(o modelplan.Option) string {
			return fmt.Sprintf("%v/%s/%d/%d", o.CPU, o.Class, o.GPUsPerNode, o.Nodes)
		}
		seen := map[string]bool{shape(out[0]): true}
		for _, o := range opts {
			if len(out) == 3 {
				break
			}
			if !seen[shape(o)] {
				seen[shape(o)] = true
				out = append(out, o)
			}
		}
	}
	return out, nil
}

// profilesFor turns the ways chosen into catalog profiles, the first
// preferred.
func profilesFor(layout []modelplan.Option, req Request, memGB, cores int) []catalog.Profile {
	var out []catalog.Profile
	for i, o := range layout {
		p := catalog.Profile{Engine: o.Engine, Weight: 1, Priority: 100 - 10*i}
		switch {
		case o.CPU:
			p.Name = "cpu"
			p.Args = llamaArgs(req)
			p.MinMemoryGB, p.CPUCores = memGB, cores
		case o.Engine == "llamacpp":
			p.Name = fmt.Sprintf("gpu-%dx%s", o.GPUsPerNode, o.Class)
			p.GPUClass, p.GPUsPerNode = o.Class, o.GPUsPerNode
			p.Args = llamaArgs(req)
		default:
			p.Name = fmt.Sprintf("gpu-%dx%s", o.GPUs(), o.Class)
			p.GPUClass, p.TP = o.Class, o.GPUsPerNode
			if o.Nodes > 1 {
				p.PP, p.Nodes, p.GPUsPerNode = o.Nodes, o.Nodes, o.GPUsPerNode
			}
			p.Args = []string{"--max-model-len", strconv.Itoa(req.ContextLength)}
		}
		out = append(out, p)
	}
	return out
}

// llamaArgs size llama.cpp's pool and its slots together. The pool is shared
// by the slots, so slots of the context the catalog promises are as many as
// the pool holds: the engine's default of four turned a 32768 pool into
// slots of 8192, and the gateway, told 32768, sent requests the engine then
// refused.
func llamaArgs(req Request) []string {
	args := []string{"--ctx-size", strconv.Itoa(req.CtxPool),
		"--parallel", strconv.Itoa(max(1, req.CtxPool/req.ContextLength))}
	if req.Embeddings {
		// A passage is embedded in one batch: with the default 512, one of
		// 516 tokens was refused.
		n := strconv.Itoa(req.ContextLength)
		args = append([]string{"--embedding"}, append(args, "--batch-size", n, "--ubatch-size", n)...)
	}
	return args
}

// An embedding model's slots, when its caller did not size them: passages
// of up to a thousand tokens, thirty-two at a time.
const (
	embedSlotTokens = 1024
	embedSlots      = 32
)

// maxOutput is the longest answer the catalog allows: a quarter of the
// context, from 1024 to 8192. 1024 everywhere cut an editor's agent off in
// the middle of the file it was writing.
func maxOutput(contextLength int) int {
	return min(8192, max(1024, contextLength/4))
}

// addAndStart puts the model in the catalog with its profile, waits for a
// replica and asks it to answer.
// hardwareExists is whether the cluster has what the ways chosen need, busy
// or not: the plan counted hosts for it, or it was asked for by hand.
func hardwareExists(layout []modelplan.Option) bool {
	for _, o := range layout {
		if o.Hosts > 0 || (o.Class != "" && o.NeedGB == 0) {
			return true
		}
	}
	return false
}

func (s *Server) addAndStart(ctx context.Context, j *Job, repo hf.Repo, set hf.FileSet, name, source string,
	profs []catalog.Profile, hardware bool) {
	req := j.Request
	// 4. The catalog.
	m := catalog.Model{
		Name: name, Aliases: req.Aliases, Description: req.Description, Source: source,
		RoutingProfile: "chat", ContextLength: req.ContextLength, MaxOutput: maxOutput(req.ContextLength),
		Profiles:     profs,
		Scaling:      catalog.Scaling{Min: req.Replicas, Max: max(req.Replicas, 1)},
		Capabilities: catalog.Capabilities{Vision: req.Projector != "", Tools: req.Tools, Embeddings: req.Embeddings},
		Origin:       &catalog.Origin{Repo: repo.ID, Revision: repo.Revision, Files: originFiles(set)},
	}
	if req.GuardFormat != "" {
		m.Role, m.Guard, m.RoutingProfile = catalog.RoleGuard, &catalog.GuardModel{Format: req.GuardFormat}, ""
	}
	m.Class = req.Class
	s.step(j, "catalog", Running, "adding "+name)
	if err := s.deps.Catalog.Add(ctx, m, req.By); err != nil {
		s.fail(j, "catalog", err)
		return
	}
	if req.Replicas == 0 {
		s.step(j, "catalog", Done, name+" added on demand: it loads when first asked for")
		s.finish(j)
		return
	}
	s.step(j, "catalog", Done, fmt.Sprintf("%s added, %d replica(s)", name, req.Replicas))

	// 5. A replica, ready.
	rep, err := s.waitReady(ctx, j, name, hardware)
	if errors.Is(err, errBusy) {
		// In the catalog, and the hardware for it exists: it starts when
		// the cards are free, with nothing more for anyone to do.
		s.step(j, "placement", Skipped, err.Error())
		s.finish(j)
		return
	}
	if err != nil {
		s.fail(j, "placement", err)
		return
	}

	// 6. Does it answer.
	s.step(j, "smoke test", Running, "asking "+rep.ID+" to say hello")
	reply, tps, err := s.smoke(ctx, rep, name, req.Embeddings)
	if err != nil {
		s.fail(j, "smoke test", err)
		return
	}
	s.mu.Lock()
	j.Reply, j.TokensPerSecond = reply, tps
	s.mu.Unlock()
	s.step(j, "smoke test", Done, fmt.Sprintf("answered at %.1f tokens/s: %q", tps, reply))
	s.finish(j)
}

func (s *Server) finish(j *Job) {
	s.mu.Lock()
	j.State = Done
	j.traceEnd(nil)
	s.mu.Unlock()
	s.deps.Log.Info("model served", "job", j.ID, "repo", j.Request.Repo, "model", j.Model)
}

// pick chooses the file set to serve.
func pick(repo hf.Repo, req Request) (hf.FileSet, error) {
	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") {
			continue
		}
		weights = append(weights, f)
	}
	sets := hf.Group(weights)
	if len(sets) == 0 {
		return hf.FileSet{}, fmt.Errorf("%s has no GGUF files; the nodes run llama.cpp, which needs one", repo.ID)
	}
	var names []string
	for _, st := range sets {
		names = append(names, st.Name)
		if req.File != "" && (st.Name == req.File || slices.ContainsFunc(st.Files, func(f hf.File) bool { return f.Path == req.File })) {
			return st, nil
		}
	}
	if req.File == "" {
		want := strings.ToLower(req.Quant)
		var match []hf.FileSet
		for _, st := range sets {
			if strings.Contains(strings.ToLower(st.Name), want) {
				match = append(match, st)
			}
		}
		if len(match) == 1 {
			return match[0], nil
		}
		if len(match) > 1 {
			// Both a split and a joined copy of one quantisation: the joined
			// one, which is one download to verify rather than several.
			slices.SortFunc(match, func(a, b hf.FileSet) int { return a.Shards - b.Shards })
			return match[0], nil
		}
	}
	return hf.FileSet{}, fmt.Errorf("%s has no file matching %q; it has: %s",
		repo.ID, or(req.File, req.Quant), strings.Join(names, ", "))
}

// mirror fetches the set into the mirror and returns the source a catalog
// entry names: the file, or the directory for a set of shards.
func (s *Server) mirror(ctx context.Context, j *Job, repo hf.Repo, set hf.FileSet, name string) (string, error) {
	version := repo.ShortRevision()
	s.step(j, "mirror", Running, "asking the mirror for "+set.Name)
	var base string
	// Files already there from this commit, checksums matching, are not
	// fetched again: a model moved back and then forward again.
	have := map[string]bool{}
	if got, err := s.deps.Mirror.List(ctx); err == nil {
		for _, st := range got.Models {
			if st.Model == name && st.Version == version && st.State == mirror.Present {
				for _, f := range set.Files {
					if mirror.LastPathElement(f.Path) == st.File && (f.SHA256 == "" || strings.EqualFold(f.SHA256, st.SHA256)) {
						have[st.File] = true
						base = mirror.Entry{Model: name, Version: version, File: st.File}.URL(got.Base)
					}
				}
			}
		}
	}
	for _, f := range set.Files {
		file := mirror.LastPathElement(f.Path)
		if have[file] {
			continue
		}
		resp, err := s.deps.Mirror.Fetch(ctx, mirror.FetchRequest{
			Source: repo.DownloadURL("", f.Path), Model: name, Version: version, File: file,
		})
		// Already there is not a failure: the same commit fetched before.
		if err != nil && !strings.Contains(err.Error(), "already") {
			return "", err
		}
		base = resp.URL
	}
	deadline := time.Now().Add(6 * time.Hour)
	for {
		got, err := s.deps.Mirror.List(ctx)
		if err != nil {
			return "", err
		}
		var done, total int64
		present, failed := 0, ""
		sums := map[string]string{}
		for _, st := range got.Models {
			if st.Model != name || st.Version != version {
				continue
			}
			done += st.BytesDone
			total += max(st.BytesTotal, st.BytesDone)
			switch st.State {
			case mirror.Present:
				present++
				sums[st.File] = st.SHA256
			case mirror.Failed:
				failed = st.Err
			}
		}
		if failed != "" {
			return "", fmt.Errorf("the mirror could not fetch it: %s", failed)
		}
		if present >= len(set.Files) {
			// What arrived against what the registry published: a mirror
			// copies one file to every node, and is the one place a corrupted
			// download is caught once instead of loaded everywhere.
			for _, f := range set.Files {
				file := mirror.LastPathElement(f.Path)
				if f.SHA256 != "" && sums[file] != "" && !strings.EqualFold(f.SHA256, sums[file]) {
					return "", fmt.Errorf("%s arrived with sha256 %s, and the registry published %s",
						file, sums[file][:12], f.SHA256[:12])
				}
			}
			source := base
			if len(set.Files) > 1 {
				source = strings.TrimSuffix(source, mirror.LastPathElement(source))
			}
			s.step(j, "mirror", Done, fmt.Sprintf("%.1f GB on the mirror, checksums match the registry", gb(total)))
			return source, nil
		}
		if total > 0 {
			s.step(j, "mirror", Running, fmt.Sprintf("fetching: %.1f of %.1f GB (%d%%)",
				gb(done), gb(total), done*100/total))
		}
		if time.Now().After(deadline) {
			return "", errors.New("the mirror had not finished after six hours")
		}
		time.Sleep(s.deps.Poll)
	}
}

// waitReady waits for a replica of the model to serve, reporting what it is
// doing and giving up early when it cannot be placed or keeps failing.
func (s *Server) waitReady(ctx context.Context, j *Job, name string, hardware bool) (reconciler.Replica, error) {
	s.step(j, "placement", Running, "waiting for the control plane to place it")
	started := time.Now()
	deadline := started.Add(s.deps.ReadyTimeout)
	for {
		for _, h := range s.deps.Fleet.ModelHealth() {
			if h.Model != name {
				continue
			}
			if h.Parked {
				return reconciler.Replica{}, fmt.Errorf("it failed to start %d times and is parked: %s", h.Failures, h.Reason)
			}
			// A minute's grace from when this step began, not from the job:
			// the download before it can take hours, and the first pass may
			// not have run yet.
			if h.Unplaceable && time.Since(started) > s.deps.UnplacedGrace {
				if hardware {
					return reconciler.Replica{}, fmt.Errorf("%w: it is in the catalog and starts when there is room, "+
						"with nothing more to do; what holds it now:\n%s", errBusy, h.Placement)
				}
				return reconciler.Replica{}, fmt.Errorf("no node can host it:\n%s", h.Placement)
			}
		}
		for _, r := range s.deps.Fleet.Replicas() {
			if r.Model != name {
				continue
			}
			if r.State == reconciler.Ready {
				s.linkReplica(j, r)
				s.step(j, "placement", Done, fmt.Sprintf("%s ready on %s at %s", r.ID,
					strings.Join(r.NodeIDs(), ","), r.Endpoint))
				return r, nil
			}
			s.step(j, "placement", Running, fmt.Sprintf("%s on %s is %s", r.ID,
				strings.Join(r.NodeIDs(), ","), r.State))
		}
		if time.Now().After(deadline) {
			return reconciler.Replica{}, fmt.Errorf("no replica was ready after %s", s.deps.ReadyTimeout)
		}
		select {
		case <-ctx.Done():
			return reconciler.Replica{}, ctx.Err()
		case <-time.After(s.deps.Poll):
		}
	}
}

// smoke sends one short request straight to the replica, with its engine key.
func (s *Server) smoke(ctx context.Context, r reconciler.Replica, name string, embeddings bool) (string, float64, error) {
	if embeddings {
		return s.smokeEmbeddings(ctx, r, name)
	}
	body, _ := json.Marshal(map[string]any{
		"model":      name,
		"max_tokens": 32,
		"messages":   []map[string]string{{"role": "user", "content": "Say hello in one short sentence."}},
	})
	req, err := http.NewRequestWithContext(ctx, http.MethodPost,
		strings.TrimRight(r.Endpoint, "/")+"/v1/chat/completions", bytes.NewReader(body))
	if err != nil {
		return "", 0, err
	}
	req.Header.Set("Content-Type", "application/json")
	if k := s.deps.Fleet.EngineKey(r.ID); k != "" {
		req.Header.Set("Authorization", "Bearer "+k)
	}
	start := time.Now()
	resp, err := s.deps.HTTP.Do(req)
	if err != nil {
		return "", 0, err
	}
	defer resp.Body.Close()
	b, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if resp.StatusCode != http.StatusOK {
		return "", 0, fmt.Errorf("the replica answered %s: %s", resp.Status, strings.TrimSpace(string(b)))
	}
	var out struct {
		Choices []struct {
			Message struct {
				Content string `json:"content"`
			} `json:"message"`
		} `json:"choices"`
		Usage struct {
			CompletionTokens int `json:"completion_tokens"`
		} `json:"usage"`
	}
	if err := json.Unmarshal(b, &out); err != nil || len(out.Choices) == 0 {
		return "", 0, fmt.Errorf("the replica's answer was not a completion: %s", strings.TrimSpace(string(b)))
	}
	secs := time.Since(start).Seconds()
	tps := 0.0
	if secs > 0 {
		tps = float64(out.Usage.CompletionTokens) / secs
	}
	return strings.TrimSpace(out.Choices[0].Message.Content), tps, nil
}

func gb(b int64) float64 { return float64(b) / (1 << 30) }

func or(a, b string) string {
	if a != "" {
		return a
	}
	return b
}

func orUnknown(arch string) string {
	if arch == "" {
		return "architecture not read"
	}
	return arch
}

// Projectors are a repository's vision projectors: the mmproj files a vision
// model loads beside its weights.
func Projectors(repo hf.Repo) []string {
	var out []string
	for _, f := range repo.Files {
		if strings.HasSuffix(strings.ToLower(f.Path), ".gguf") && strings.Contains(strings.ToLower(f.Path), "mmproj") {
			out = append(out, f.Path)
		}
	}
	return out
}

// headerURL is the file whose header says what the model is: the source
// itself, or in a directory the first part of the weights, not a projector.
func headerURL(source string, set hf.FileSet) string {
	if !strings.HasSuffix(source, "/") {
		return source
	}
	var names []string
	for _, f := range set.Files {
		name := mirror.LastPathElement(f.Path)
		lower := strings.ToLower(name)
		if strings.HasSuffix(lower, ".gguf") && !strings.Contains(lower, "mmproj") {
			names = append(names, name)
		}
	}
	if len(names) == 0 {
		return ""
	}
	slices.Sort(names)
	return source + names[0]
}

// SafetensorsSet is what serving a repository's original weights fetches.
func SafetensorsSet(repo hf.Repo) (hf.FileSet, error) { return pickSafetensors(repo, Request{}) }

// pickSafetensors chooses a repository's original weights and the files an
// engine reads beside them: config.json, the tokenizer, the shard index.
func pickSafetensors(repo hf.Repo, _ Request) (hf.FileSet, error) {
	var weights []hf.File
	for _, f := range repo.Files {
		if !strings.Contains(f.Path, "/") && strings.HasSuffix(f.Path, ".safetensors") {
			weights = append(weights, f)
		}
	}
	sets := hf.Group(weights)
	if len(sets) == 0 {
		return hf.FileSet{}, fmt.Errorf("%s has no safetensors weights at its top", repo.ID)
	}
	// Sharded "model-00001-of-00004" or a single "model.safetensors"; a
	// repository with a second, consolidated copy keeps it under another name.
	set := sets[0]
	for _, st := range sets {
		if strings.HasPrefix(st.Name, "model") {
			set = st
			break
		}
	}
	companions := hf.Companions(repo.Files)
	if !slices.ContainsFunc(companions, func(f hf.File) bool { return f.Path == "config.json" }) {
		return hf.FileSet{}, fmt.Errorf("%s has no config.json, which vLLM reads the model's shape from", repo.ID)
	}
	for _, f := range companions {
		set.Files = append(set.Files, f)
		set.Bytes += f.Bytes
	}
	return set, nil
}

// smokeEmbeddings asks an embedding model for one vector, and says its
// length: an embedding model does not chat.
func (s *Server) smokeEmbeddings(ctx context.Context, r reconciler.Replica, name string) (string, float64, error) {
	body, _ := json.Marshal(map[string]any{"model": name, "input": "hello"})
	req, err := http.NewRequestWithContext(ctx, http.MethodPost,
		strings.TrimRight(r.Endpoint, "/")+"/v1/embeddings", bytes.NewReader(body))
	if err != nil {
		return "", 0, err
	}
	req.Header.Set("Content-Type", "application/json")
	if k := s.deps.Fleet.EngineKey(r.ID); k != "" {
		req.Header.Set("Authorization", "Bearer "+k)
	}
	resp, err := s.deps.HTTP.Do(req)
	if err != nil {
		return "", 0, err
	}
	defer resp.Body.Close()
	b, _ := io.ReadAll(io.LimitReader(resp.Body, 4<<20))
	if resp.StatusCode != http.StatusOK {
		return "", 0, fmt.Errorf("the replica answered %s: %s", resp.Status, strings.TrimSpace(string(b)))
	}
	var out struct {
		Data []struct {
			Embedding []float64 `json:"embedding"`
		} `json:"data"`
	}
	if err := json.Unmarshal(b, &out); err != nil || len(out.Data) == 0 || len(out.Data[0].Embedding) == 0 {
		return "", 0, fmt.Errorf("the replica answered no vector: %s", strings.TrimSpace(string(b)))
	}
	return fmt.Sprintf("a vector of %d dimensions", len(out.Data[0].Embedding)), 0, nil
}

// embeds is whether the catalog serves name as an embedding model, for a
// smoke test of one already there (a canary, an update).
func (s *Server) embeds(name string) bool {
	if s.deps.Current == nil {
		return false
	}
	for _, m := range s.deps.Current() {
		if m.Name == name {
			return m.Capabilities.Embeddings
		}
	}
	return false
}

// fitPoolToCard shrinks a context pool the cards cannot hold to one they
// can, saying so, and refuses weights that do not fit at all. The slots stay
// as many as asked while each keeps 2048 tokens, then fewer.
func fitPoolToCard(req Request, h gguf.Header, weights int64, class string, gpus int) (Request, string, error) {
	maxTok := h.PoolThatFits(weights, gguf.ClassMemGB(class), gpus)
	switch {
	case maxTok == 0 || req.CtxPool <= maxTok:
		return req, "", nil
	case maxTok < 0:
		return req, "", fmt.Errorf("its weights (%.1f GB) do not fit %d card(s) of %s with room to run: "+
			"choose a smaller quantisation, or more cards", gb(weights), gpus, class)
	}
	slots := max(1, req.CtxPool/max(req.ContextLength, 1))
	for slots > 1 && maxTok/slots < 2048 {
		slots--
	}
	per := (maxTok / slots) / 1024 * 1024
	if per < 1024 {
		return req, "", fmt.Errorf("its weights (%.1f GB) leave room on %s for a context of %d tokens only",
			gb(weights), class, maxTok)
	}
	was := req.CtxPool
	req.ContextLength, req.CtxPool = per, per*slots
	return req, fmt.Sprintf("a %d-token pool does not fit %s beside %.1f GB of weights: %d slots of %d tokens instead",
		was, class, gb(weights), slots, per), nil
}