Files
llmgw/internal/control/modelserve/modelserve.go
T
Ilker ManapandClaude Opus 5.5 8ffde2b80c
CI / test (push) Successful in 5m18s
A big model's download no longer runs out its startup timeout
The startup timeout runs from when a node's download last moved, so a
model of hundreds of gigabytes on a slow mirror is not failed while it
arrives. ctl model serve --env sets engine variables (NCCL_DEBUG=INFO)
on every node, and a layout given by hand keeps expert parallelism for a
model of experts. docs/large-models.md: how a model is split, what each
way needs of a card, and a checklist for bringing one up on a fabric.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2026-10-20 09:45:31 +02:00

1167 lines
40 KiB
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"`
// Env is set in the engine's environment, on every node of a replica:
// NCCL_DEBUG=INFO to see how its cards talk, NCCL_IB_* to steer them.
// Only the variables engines are tuned with; the catalog refuses others.
Env map[string]string `json:"env,omitempty"`
// IfAbsent adds the model only if nobody has meanwhile: the leader's own
// serving of a judge a policy names gives way to the one init or an
// operator added while it fetched, rather than replacing it.
IfAbsent bool `json:"-"`
// MemoryGB is what a replica on CPUs reserves, in place of the figure
// worked out from its weights and context: a judge's, whose own sizing
// is known.
MemoryGB int `json:"memory_gb,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())
}
// What the first way needs of each card, for the context asked for.
need := ""
if first := layout[0]; first.NeedGB > 0 {
tokens := req.ContextLength
if req.Format != "safetensors" && req.CtxPool > 0 {
tokens = req.CtxPool
}
need = "; needs " + first.WithContext(plan.KVBytesPerToken, tokens).MemoryText()
}
s.step(j, "registry", Done, fmt.Sprintf("%s, %.1f GB in %d file(s), revision %s; runs on %s%s",
set.Name, gb(set.Bytes), len(set.Files), repo.ShortRevision(), strings.Join(where, ", else "), need))
// 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
}
}
// Under the lock: a job is read by Get while it runs, and the race
// detector caught this write beside it.
s.mu.Lock()
j.Request = req
s.mu.Unlock()
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"
}
if req.MemoryGB > 0 {
memGB, basis = req.MemoryGB, "as asked, in place of "+basis
}
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
}
}
// A way the plan did not offer is still a model of experts.
o.ExpertParallel = engine == "vllm" && plan.Experts > 1 && o.GPUs() > 1
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, Env: req.Env}
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)}
if o.ExpertParallel {
// The experts whole on each card rather than each sliced
// across them: less traffic between the cards for every token.
p.Args = append(p.Args, "--enable-expert-parallel")
}
}
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)...)
}
if req.GuardFormat != "" {
// No prompt cache: a judge's prompts, each a different text, are
// not asked again, and llama.cpp's default of up to 8 GB grew one
// judge to 3.2 GB over a day.
args = append(args, "--cache-ram", "0")
}
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)
add := s.deps.Catalog.Add
if adder, ok := s.deps.Catalog.(interface {
AddNew(ctx context.Context, m catalog.Model, by string) error
}); ok && req.IfAbsent {
add = adder.AddNew
}
if err := add(ctx, m, req.By); err != nil {
if req.IfAbsent && errors.Is(err, catalog.ErrExists) {
s.step(j, "catalog", Skipped, name+" was added meanwhile; left as it is")
s.finish(j)
return
}
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
}