Files
llmgw/internal/control/modelserve/modelserve.go
T
Ilker ManapandClaude Opus 5.5 5732a5ebda
CI / test (push) Successful in 9m26s
The serve forms start at the longest context that fits
The mirror page's form began every GGUF model at 4096 tokens, the
console's at 4096 or the plan's figure. Both now start at what the model
was trained for, cut to the most whose KV cache fits beside the weights
where the form places it; the mirror page's slots are as many of that, up
to four, as the room left holds. Plan carries the trained context.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-11-08 14:38:41 +01:00

1261 lines
44 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"
"gitea.mynodes.xyz/ilker/llmgw/internal/strata"
)
// 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
// RestartInPlace stops each replica before its replacement starts, for
// when there is no room beside it.
RestartInPlace 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"`
// Engine is the engine the file runs on: "strata" for a family Strata
// runs, "llamacpp" to keep it off Strata; empty, the plan's best.
Engine string `json:"engine,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 = DefaultCtxPool
}
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
}
// A pack Strata fetches itself: no mirror; its repository is the source.
if _, size, ok := strata.ForFetchedRepo(repo.ID); ok {
plan, err := s.plan(ctx, repo, PlanRequest{Repo: repo.ID})
if err != nil {
s.fail(j, "registry", err)
return
}
set := hf.FileSet{Name: size.Name}
layout, err := chooseLayout(Request{Format: "gguf", Engine: "strata", GPUClass: req.GPUClass}, plan, true, set)
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()
if !req.sized {
req.ContextLength = strataContext
}
s.step(j, "registry", Done, fmt.Sprintf("Strata's ready-made %s pack, revision %s; runs on %s",
size.Name, repo.ShortRevision(), layout[0].Label()))
s.step(j, "mirror", Done, fmt.Sprintf("not through the mirror: Strata fetches about %.0f GB onto the node itself", size.DownloadGB))
profiles := profilesFor(layout, req, 0, 0)
profiles[0].StartupTimeout = 4 * time.Hour
s.mu.Lock()
j.Request = req
s.mu.Unlock()
s.addAndStart(ctx, j, repo, set, name, repo.ID, profiles, hardwareExists(layout))
return
}
// 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 layout[0].Engine == "strata" {
// Strata prepares it on the node from the directory of its files,
// and keeps the experts in the RAM the plan counted.
if !strings.HasSuffix(source, "/") {
source = strings.TrimSuffix(source, mirror.LastPathElement(source))
}
if !req.sized {
req.ContextLength = strataContext
}
s.step(j, "check", Done, fmt.Sprintf("Strata's %s %s, prepared on the node; %.0f GB of RAM for its experts",
layout[0].StrataFamily, layout[0].StrataSize, layout[0].NeedGB))
s.mu.Lock()
j.Request = req
s.mu.Unlock()
s.addAndStart(ctx, j, repo, set, name, source, profilesFor(layout, req, 0, 0), hardwareExists(layout))
return
}
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.
// Asked for the CPU, and slow there: said, not refused — the request
// named it — so the job's log says why its replies take minutes.
if len(layout) > 0 && layout[0].CPU && layout[0].TokensPerSecond > 0 &&
layout[0].TokensPerSecond < modelplan.CPUSlowTPS {
tps := layout[0].TokensPerSecond
s.step(j, "check", Running, fmt.Sprintf("on the CPU about %.1f tokens a second for one reply: "+
"a reply of 500 tokens takes about %.0f seconds; a card would be many times faster", tps, 500/tps))
}
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 {
onStrata := o.Engine == "strata"
if o.Artifact == artifact && (req.Engine == "" || (req.Engine == "strata") == onStrata) {
opts = append(opts, o)
}
}
if req.Engine == "strata" {
// Strata's way only, on the class asked for if one was.
i := slices.IndexFunc(opts, func(o modelplan.Option) bool { return req.GPUClass == "" || o.Class == req.GPUClass })
if i < 0 {
return nil, fmt.Errorf("Strata cannot run %s here: it runs a few families only (`llmgw ctl model plan %s` "+
"says whether this is one), on a node with Strata installed, a card of 12 GB or more and the RAM for its experts",
set.Name, plan.Repo)
}
return []modelplan.Option{opts[i]}, nil
}
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 && out[0].Engine != "strata" {
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.Engine == "strata":
p.Name = "strata-" + o.Class
p.GPUClass, p.GPUsPerNode = o.Class, 1
p.Args = []string{"--family", o.StrataFamily, "--size", o.StrataSize,
"--context", strconv.Itoa(req.ContextLength), "--vision", "no"}
p.MinMemoryGB = int(o.NeedGB)
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
}
// DefaultCtxPool is llama.cpp's pool when a request names none: a context
// asked for longer than it needs a pool of its own.
const DefaultCtxPool = 16384
// 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
}
// strataContext is a Strata model's context when its caller did not say:
// Strata's own default.
const strataContext = 32768
// 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,
// The mirror deletes a copy that is not what the registry
// published, rather than holding it for the nodes.
SHA256: f.SHA256,
})
// 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
}