// 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" ) // 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 } // 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"` // 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) } if req.CtxPool <= 0 { req.CtxPool = 16384 } if req.ContextLength <= 0 { req.ContextLength = 4096 } s.mu.Lock() s.seq++ j := &Job{ID: fmt.Sprintf("serve-%d-%d", time.Now().Unix(), s.seq), Request: req, State: Running, Started: time.Now()} s.jobs[j.ID] = j s.mu.Unlock() go s.run(j) return s.snapshot(j), nil } // Jobs is every job since the control plane started, 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)}, } 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 }