Files
llmgw/internal/control/modelserve/modelserve.go
T
Ilker ManapandClaude Opus 5.5 1b71c38f89 Notifications: webhooks for a parked model, a lost node, a spent budget and new weights; Alertmanager receiver templates
notify.webhooks sends signed JSON or Slack-compatible messages, queued,
retried, de-duplicated hourly; ctl notify test checks each webhook.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2026-10-04 16:58:11 +02:00

937 lines
30 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/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"
"gitea.mynodes.xyz/ilker/llmgw/internal/modelplan"
"gitea.mynodes.xyz/ilker/llmgw/internal/notify"
)
// 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; nil
// tells no one.
Notify *notify.Notifier
}
// 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, shieldgemma or json. Empty serves clients.
GuardFormat string `json:"guard_format,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.
CtxPool int `json:"ctx_pool,omitempty"`
ContextLength int `json:"context_length,omitempty"`
// CPUCores is what a replica computes on; zero takes half the smallest
// node it fits on.
CPUCores int `json:"cpu_cores,omitempty"`
// By is who asked, for the catalog's record.
By string `json:"-"`
}
// The states a job and its steps can be in.
const (
Running = "running"
Done = "done"
Failed = "failed"
Skipped = "skipped"
)
// Step is one stage of a job.
type Step struct {
Name string `json:"name"`
State string `json:"state"`
Detail string `json:"detail,omitempty"`
At time.Time `json:"at"`
}
// Job is one model being served, and where it stands.
type Job struct {
ID string `json:"id"`
// Kind is "update" for a model moved to new weights, empty for one
// served.
Kind string `json:"kind,omitempty"`
Request Request `json:"request"`
State string `json:"state"`
Error string `json:"error,omitempty"`
Model string `json:"model,omitempty"`
Steps []Step `json:"steps"`
Started time.Time `json:"started"`
// Reply is what the model said to the smoke test, and TokensPerSecond
// how fast it said it.
Reply string `json:"reply,omitempty"`
TokensPerSecond float64 `json:"tokens_per_second,omitempty"`
}
// 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)
}
switch req.GuardFormat {
case "", "llamaguard3", "shieldgemma", "json":
default:
return Job{}, fmt.Errorf("the guard format is llamaguard3, shieldgemma or json, not %q", req.GuardFormat)
}
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")
}
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.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()
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
}
if req.TP > 0 && req.GPUsPerNode == 0 {
req.GPUsPerNode = req.TP
}
// The plan: which file suits the hardware, and how it spreads over it.
planCtx := req.CtxPool
if req.Format == "safetensors" {
planCtx = req.ContextLength
}
var plan Plan
planned := false
if s.deps.Pool != nil {
if plan, err = s.plan(ctx, repo, PlanRequest{Repo: repo.ID, Format: req.Format, Context: planCtx}); err == nil {
planned = true
} else {
s.deps.Log.Warn("planning a model", "repo", repo.ID, "err", err)
}
}
if req.Format == "gguf" && req.Quant == "" && req.File == "" {
if rec, ok := plan.Recommended(); planned && ok {
req.File = rec.Artifact
} else {
req.Quant = "Q4_K_M"
}
}
pickSet := pick
if req.Format == "safetensors" {
pickSet = pickSafetensors
}
set, err := pickSet(repo, req)
if err != nil {
s.fail(j, "registry", err)
return
}
layout, err := chooseLayout(req, plan, planned, set)
if err != nil {
s.fail(j, "registry", err)
return
}
if req.Projector != "" {
i := slices.IndexFunc(repo.Files, func(f hf.File) bool {
return (f.Path == req.Projector || mirror.LastPathElement(f.Path) == req.Projector) &&
strings.Contains(strings.ToLower(f.Path), "mmproj")
})
if i < 0 {
s.fail(j, "registry", fmt.Errorf("%s has no vision projector %q; it has: %s", repo.ID, req.Projector,
strings.Join(Projectors(repo), ", ")))
return
}
set.Files = append(set.Files, repo.Files[i])
set.Bytes += repo.Files[i].Bytes
}
name := req.Name
if name == "" {
name = hf.ModelName(repo.ID)
}
s.mu.Lock()
j.Model = name
s.mu.Unlock()
where := make([]string, 0, len(layout))
for _, o := range layout {
where = append(where, o.Label())
}
s.step(j, "registry", Done, fmt.Sprintf("%s, %.1f GB in %d file(s), revision %s; runs on %s",
set.Name, gb(set.Bytes), len(set.Files), repo.ShortRevision(), strings.Join(where, ", else ")))
// 2. The mirror: fetch, and check what arrived against what was published.
source, err := s.mirror(ctx, j, repo, set, name)
if err != nil {
s.fail(j, "mirror", err)
return
}
// 3. Whether the nodes' engine loads it, and how much memory it needs.
if req.Format == "safetensors" {
// vLLM reads the model's shape from its config.json, which came
// with the weights.
s.step(j, "check", Done, fmt.Sprintf("%.1f GB of weights for vLLM", gb(set.Bytes)))
s.addAndStart(ctx, j, repo, set, name, source, profilesFor(layout, req, 0, 0), hardwareExists(layout))
return
}
s.step(j, "check", Running, "reading its header from the mirror")
var header gguf.Header
if u := headerURL(source, set); u != "" {
v := gguf.Check(ctx, u, s.deps.LlamaVersion)
if !v.Servable {
s.fail(j, "check", fmt.Errorf("llama.cpp %s on the nodes cannot run this: %s", s.deps.LlamaVersion, v.Reason))
return
}
header = v.Header
}
memGB := int((header.MemoryBytes(set.Bytes, req.CtxPool) + 1<<30 - 1) >> 30)
basis := fmt.Sprintf("weights %.1f GB + KV cache %.1f GB for %d tokens + engine",
gb(set.Bytes), gb(int64(req.CtxPool)*header.KVBytesPerToken()), req.CtxPool)
if memGB == 0 {
// A set of shards, or a header that did not say: weights plus half
// again plus two, the rough figure from before shapes were read.
memGB = int(gb(set.Bytes)*1.5) + 2
basis = "weights plus half again plus two; the header gave no shape"
}
s.step(j, "check", Done, fmt.Sprintf("%s, needs about %d GB (%s)", orUnknown(header.Architecture), memGB, basis))
cores := req.CPUCores
if cores <= 0 {
if nc, ok := s.deps.Fleet.(interface{ NodeCPUs(int) []int }); ok {
cores = catalog.SuggestCores(nc.NodeCPUs(memGB))
}
}
s.addAndStart(ctx, j, repo, set, name, source, profilesFor(layout, req, memGB, cores), hardwareExists(layout))
}
// chooseLayout is where a model runs: what was asked, else the plan's best
// for the chosen file; with Fallbacks, the next best of other shapes after.
func chooseLayout(req Request, plan Plan, planned bool, set hf.FileSet) ([]modelplan.Option, error) {
artifact := set.Name
engine := "llamacpp"
if req.Format == "safetensors" {
artifact, engine = "original", "vllm"
}
var opts []modelplan.Option
for _, o := range plan.Options {
if o.Artifact == artifact {
opts = append(opts, o)
}
}
var out []modelplan.Option
switch {
case req.CPU || req.GPUClass != "":
o := modelplan.Option{Artifact: artifact, Format: req.Format, Engine: engine, CPU: req.CPU}
if !req.CPU {
o.Class, o.GPUsPerNode, o.Nodes = req.GPUClass, max(req.GPUsPerNode, 1), max(req.Nodes, 1)
}
if o.Nodes > 1 && engine == "llamacpp" {
return nil, errors.New("llama.cpp does not span nodes; serve the original weights with vLLM for that")
}
// The plan's own reckoning of that way, if it has one: what it holds
// and on how many hosts.
for _, p := range opts {
if p.CPU == o.CPU && p.Class == o.Class && p.GPUsPerNode == o.GPUsPerNode && max(p.Nodes, 1) == o.Nodes {
o = p
}
}
out = append(out, o)
case len(opts) > 0:
out = append(out, opts[0])
case planned:
return nil, fmt.Errorf("%s fits nowhere in this cluster with a %d-token context: no card holds it, "+
"nor does any node's memory; `llmgw ctl model plan %s` shows what each file needs", set.Name, plan.Context, plan.Repo)
case engine == "vllm":
return nil, errors.New("original weights run on vLLM, on GPUs: name the GPU class")
default:
out = append(out, modelplan.Option{Artifact: artifact, Format: req.Format, Engine: engine, CPU: true})
}
if req.Fallbacks {
shape := func(o modelplan.Option) string {
return fmt.Sprintf("%v/%s/%d/%d", o.CPU, o.Class, o.GPUsPerNode, o.Nodes)
}
seen := map[string]bool{shape(out[0]): true}
for _, o := range opts {
if len(out) == 3 {
break
}
if !seen[shape(o)] {
seen[shape(o)] = true
out = append(out, o)
}
}
}
return out, nil
}
// profilesFor turns the ways chosen into catalog profiles, the first
// preferred.
func profilesFor(layout []modelplan.Option, req Request, memGB, cores int) []catalog.Profile {
var out []catalog.Profile
for i, o := range layout {
p := catalog.Profile{Engine: o.Engine, Weight: 1, Priority: 100 - 10*i}
switch {
case o.CPU:
p.Name = "cpu"
p.Args = []string{"--ctx-size", strconv.Itoa(req.CtxPool)}
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 = []string{"--ctx-size", strconv.Itoa(req.CtxPool)}
default:
p.Name = fmt.Sprintf("gpu-%dx%s", o.GPUs(), o.Class)
p.GPUClass, p.TP = o.Class, o.GPUsPerNode
if o.Nodes > 1 {
p.PP, p.Nodes, p.GPUsPerNode = o.Nodes, o.Nodes, o.GPUsPerNode
}
p.Args = []string{"--max-model-len", strconv.Itoa(req.ContextLength)}
}
out = append(out, p)
}
return out
}
// 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: 1024,
Profiles: profs,
Scaling: catalog.Scaling{Min: req.Replicas, Max: max(req.Replicas, 1)},
Capabilities: catalog.Capabilities{Vision: req.Projector != ""},
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}, ""
}
s.step(j, "catalog", Running, "adding "+name)
if err := s.deps.Catalog.Add(ctx, m, req.By); err != nil {
s.fail(j, "catalog", err)
return
}
if req.Replicas == 0 {
s.step(j, "catalog", Done, name+" added on demand: it loads when first asked for")
s.finish(j)
return
}
s.step(j, "catalog", Done, fmt.Sprintf("%s added, %d replica(s)", name, req.Replicas))
// 5. A replica, ready.
rep, err := s.waitReady(ctx, j, name, hardware)
if errors.Is(err, errBusy) {
// In the catalog, and the hardware for it exists: it starts when
// the cards are free, with nothing more for anyone to do.
s.step(j, "placement", Skipped, err.Error())
s.finish(j)
return
}
if err != nil {
s.fail(j, "placement", err)
return
}
// 6. Does it answer.
s.step(j, "smoke test", Running, "asking "+rep.ID+" to say hello")
reply, tps, err := s.smoke(ctx, rep, name)
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
// 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.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
}
// 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
}