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/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
	"gitea.mynodes.xyz/ilker/llmgw/internal/gguf"
	"gitea.mynodes.xyz/ilker/llmgw/internal/hf"
	"gitea.mynodes.xyz/ilker/llmgw/internal/mirror"
)

// 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
}

// 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"`
	// 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"`
	// CtxPool is llama.cpp's --ctx-size; ContextLength what one request may
	// use. They default to 16384 and 4096, what four parallel slots divide.
	CtxPool       int `json:"ctx_pool,omitempty"`
	ContextLength int `json:"context_length,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"`
	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"`
}

// Server runs jobs and remembers them.
type Server struct {
	deps Deps
	mu   sync.Mutex
	jobs map[string]*Job
	seq  int
}

// 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.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{}}
}

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")
	}
	if req.Quant == "" && req.File == "" {
		req.Quant = "Q4_K_M"
	}
	if req.CtxPool <= 0 {
		req.CtxPool = 16384
	}
	if req.ContextLength <= 0 {
		req.ContextLength = 4096
	}
	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.jobs[j.ID] = j
	s.mu.Unlock()
	go s.run(j)
	return s.snapshot(j), nil
}

// 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()
	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()
	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
	}
	set, err := pick(repo, req)
	if err != nil {
		s.fail(j, "registry", err)
		return
	}
	name := req.Name
	if name == "" {
		name = hf.ModelName(repo.ID)
	}
	s.mu.Lock()
	j.Model = name
	s.mu.Unlock()
	s.step(j, "registry", Done, fmt.Sprintf("%s, %.1f GB in %d file(s), revision %s",
		set.Name, gb(set.Bytes), len(set.Files), repo.ShortRevision()))

	// 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' llama.cpp loads it, and how much memory it needs.
	s.step(j, "check", Running, "reading its header from the mirror")
	var header gguf.Header
	if !strings.HasSuffix(source, "/") {
		v := gguf.Check(ctx, source, 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
	}
	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))

	// 4. The catalog.
	m := catalog.Model{
		Name: name, Aliases: req.Aliases, Description: req.Description, Source: source,
		RoutingProfile: "chat", ContextLength: req.ContextLength, MaxOutput: 1024,
		Profiles: []catalog.Profile{{
			Name: "cpu", Engine: "llamacpp", Weight: 1, Priority: 100,
			Args: []string{"--ctx-size", strconv.Itoa(req.CtxPool)}, MinMemoryGB: memGB,
		}},
		Scaling: catalog.Scaling{Min: req.Replicas, Max: max(req.Replicas, 1)},
	}
	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)
	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)
	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
	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
	for _, f := range set.Files {
		file := mirror.LastPathElement(f.Path)
		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) (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) > time.Minute {
				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.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) (string, float64, error) {
	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
}