llmgw/internal/control/modelserve/modelserve.go

// Package modelserve takes a model from a registry repository to a replica
// answering requests, as one job.
//
// Adding a model was four commands and a hand-written YAML file: plan it,
// fetch it into the mirror, write its catalog entry with a memory figure
// somebody guessed, scale it — and then find out whether it loads. This runs
// those steps in the control plane, where the mirror, the registry, the
// catalog and the fleet already are, and reports each one as it goes, so
// whatever asks — ctl, a page — only has to watch.
package modelserve

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"regexp"
	"slices"
	"strconv"
	"strings"
	"sync"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/catalog"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/reconciler"
	"gitea.mynodes.xyz/ilker/llmgw/internal/gguf"
	"gitea.mynodes.xyz/ilker/llmgw/internal/hf"
	"gitea.mynodes.xyz/ilker/llmgw/internal/mirror"
)

// Registry reads a repository.
type Registry interface {
	Repo(ctx context.Context, id string) (hf.Repo, error)
}

// Mirror fetches files and reports on them.
type Mirror interface {
	Fetch(ctx context.Context, req mirror.FetchRequest) (mirror.FetchResponse, error)
	List(ctx context.Context) (mirror.ModelsResponse, error)
}

// Catalog adds a model.
type Catalog interface {
	Add(ctx context.Context, m catalog.Model, by string) error
}

// Fleet is what the job watches once the model is in the catalog.
type Fleet interface {
	Replicas() []reconciler.Replica
	ModelHealth() []reconciler.ModelHealth
	EngineKey(replicaID string) string
}

// Deps are what a job needs.
type Deps struct {
	Registry Registry
	Mirror   Mirror
	Catalog  Catalog
	Fleet    Fleet
	// LlamaVersion is the llama.cpp build the nodes install, which decides
	// which architectures load.
	LlamaVersion string
	HTTP         *http.Client
	Log          *slog.Logger
	// Poll is how often the job looks at the mirror and the fleet.
	Poll time.Duration
	// ReadyTimeout bounds the wait for a replica: the node fetching the
	// weights from the mirror and the engine loading them.
	ReadyTimeout time.Duration

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

// 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"`
	GPUClass string `json:"gpu_class,omitempty"`
	TP       int    `json:"tp,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.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.GPUClass == "" {
			return Job{}, errors.New("safetensors weights run on vLLM, on GPUs: name the GPU class, as `llmgw ctl status` lists it")
		}
		if req.Projector != "" {
			return Job{}, errors.New("a projector is for a GGUF model; vLLM reads a vision model's own files")
		}
		req.TP = max(req.TP, 1)
	default:
		return Job{}, fmt.Errorf("the format is gguf or safetensors, not %q", req.Format)
	}
	if req.Quant == "" && req.File == "" {
		req.Quant = "Q4_K_M"
	}
	if req.CtxPool <= 0 {
		req.CtxPool = 16384
	}
	if req.ContextLength <= 0 {
		req.ContextLength = 4096
	}
	s.mu.Lock()
	s.seq++
	j := &Job{ID: fmt.Sprintf("serve-%d-%d", time.Now().Unix(), s.seq), Request: req,
		State: Running, Started: time.Now()}
	s.jobs[j.ID] = j
	s.mu.Unlock()
	go s.run(j)
	return s.snapshot(j), nil
}

// Get returns a job as it stands.
func (s *Server) Get(id string) (Job, bool) {
	s.mu.Lock()
	j, ok := s.jobs[id]
	s.mu.Unlock()
	if !ok {
		return Job{}, false
	}
	return s.snapshot(j), true
}

func (s *Server) snapshot(j *Job) Job {
	s.mu.Lock()
	defer s.mu.Unlock()
	out := *j
	out.Steps = slices.Clone(j.Steps)
	return out
}

// step starts a new step, or updates the detail of the current one.
func (s *Server) step(j *Job, name, state, detail string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if n := len(j.Steps); n > 0 && j.Steps[n-1].Name == name {
		j.Steps[n-1].State, j.Steps[n-1].Detail, j.Steps[n-1].At = state, detail, time.Now()
		return
	}
	j.Steps = append(j.Steps, Step{Name: name, State: state, Detail: detail, At: time.Now()})
}

func (s *Server) fail(j *Job, name string, err error) {
	s.step(j, name, Failed, err.Error())
	s.mu.Lock()
	j.State, j.Error = Failed, err.Error()
	s.mu.Unlock()
	s.deps.Log.Warn("serving a model failed", "job", j.ID, "repo", j.Request.Repo, "step", name, "err", err)
}

func (s *Server) run(j *Job) {
	ctx := context.Background()
	req := j.Request

	// 1. The registry: which files.
	s.step(j, "registry", Running, "reading "+req.Repo)
	repo, err := s.deps.Registry.Repo(ctx, req.Repo)
	if err != nil {
		s.fail(j, "registry", err)
		return
	}
	if repo.Gated.Is() {
		s.fail(j, "registry", fmt.Errorf("%s is gated: it needs an account's approval the mirror does not have", req.Repo))
		return
	}
	pickSet := pick
	if req.Format == "safetensors" {
		pickSet = pickSafetensors
	}
	set, err := pickSet(repo, req)
	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()
	s.step(j, "registry", Done, fmt.Sprintf("%s, %.1f GB in %d file(s), revision %s",
		set.Name, gb(set.Bytes), len(set.Files), repo.ShortRevision()))

	// 2. The mirror: fetch, and check what arrived against what was published.
	source, err := s.mirror(ctx, j, repo, set, name)
	if err != nil {
		s.fail(j, "mirror", err)
		return
	}

	// 3. Whether the nodes' llama.cpp loads it, and how much memory it needs.
	if req.Format == "safetensors" {
		// vLLM reads the model's shape from its config.json, which came
		// with the weights; the GPUs it needs are the class and width asked.
		s.step(j, "check", Done, fmt.Sprintf("%.1f GB of weights for vLLM on %d × %s", gb(set.Bytes), req.TP, req.GPUClass))
		s.addAndStart(ctx, j, repo, set, name, source, catalog.Profile{
			Name: "gpu", Engine: "vllm", GPUClass: req.GPUClass, TP: req.TP, Weight: 1, Priority: 100,
			Args: []string{"--max-model-len", strconv.Itoa(req.ContextLength)},
		})
		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, catalog.Profile{
		Name: "cpu", Engine: "llamacpp", Weight: 1, Priority: 100,
		Args: []string{"--ctx-size", strconv.Itoa(req.CtxPool)}, MinMemoryGB: memGB, CPUCores: cores,
	})
}

// addAndStart puts the model in the catalog with its profile, waits for a
// replica and asks it to answer.
func (s *Server) addAndStart(ctx context.Context, j *Job, repo hf.Repo, set hf.FileSet, name, source string, prof catalog.Profile) {
	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:     []catalog.Profile{prof},
		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)
	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) (reconciler.Replica, error) {
	s.step(j, "placement", Running, "waiting for the control plane to place it")
	started := time.Now()
	deadline := started.Add(s.deps.ReadyTimeout)
	for {
		for _, h := range s.deps.Fleet.ModelHealth() {
			if h.Model != name {
				continue
			}
			if h.Parked {
				return reconciler.Replica{}, fmt.Errorf("it failed to start %d times and is parked: %s", h.Failures, h.Reason)
			}
			// A minute's grace from when this step began, not from the job:
			// the download before it can take hours, and the first pass may
			// not have run yet.
			if h.Unplaceable && time.Since(started) > time.Minute {
				return reconciler.Replica{}, fmt.Errorf("no node can host it:\n%s", h.Placement)
			}
		}
		for _, r := range s.deps.Fleet.Replicas() {
			if r.Model != name {
				continue
			}
			if r.State == reconciler.Ready {
				s.step(j, "placement", Done, fmt.Sprintf("%s ready on %s at %s", r.ID,
					strings.Join(r.NodeIDs(), ","), r.Endpoint))
				return r, nil
			}
			s.step(j, "placement", Running, fmt.Sprintf("%s on %s is %s", r.ID,
				strings.Join(r.NodeIDs(), ","), r.State))
		}
		if time.Now().After(deadline) {
			return reconciler.Replica{}, fmt.Errorf("no replica was ready after %s", s.deps.ReadyTimeout)
		}
		select {
		case <-ctx.Done():
			return reconciler.Replica{}, ctx.Err()
		case <-time.After(s.deps.Poll):
		}
	}
}

// smoke sends one short request straight to the replica, with its engine key.
func (s *Server) smoke(ctx context.Context, r reconciler.Replica, name string) (string, float64, error) {
	body, _ := json.Marshal(map[string]any{
		"model":      name,
		"max_tokens": 32,
		"messages":   []map[string]string{{"role": "user", "content": "Say hello in one short sentence."}},
	})
	req, err := http.NewRequestWithContext(ctx, http.MethodPost,
		strings.TrimRight(r.Endpoint, "/")+"/v1/chat/completions", bytes.NewReader(body))
	if err != nil {
		return "", 0, err
	}
	req.Header.Set("Content-Type", "application/json")
	if k := s.deps.Fleet.EngineKey(r.ID); k != "" {
		req.Header.Set("Authorization", "Bearer "+k)
	}
	start := time.Now()
	resp, err := s.deps.HTTP.Do(req)
	if err != nil {
		return "", 0, err
	}
	defer resp.Body.Close()
	b, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if resp.StatusCode != http.StatusOK {
		return "", 0, fmt.Errorf("the replica answered %s: %s", resp.Status, strings.TrimSpace(string(b)))
	}
	var out struct {
		Choices []struct {
			Message struct {
				Content string `json:"content"`
			} `json:"message"`
		} `json:"choices"`
		Usage struct {
			CompletionTokens int `json:"completion_tokens"`
		} `json:"usage"`
	}
	if err := json.Unmarshal(b, &out); err != nil || len(out.Choices) == 0 {
		return "", 0, fmt.Errorf("the replica's answer was not a completion: %s", strings.TrimSpace(string(b)))
	}
	secs := time.Since(start).Seconds()
	tps := 0.0
	if secs > 0 {
		tps = float64(out.Usage.CompletionTokens) / secs
	}
	return strings.TrimSpace(out.Choices[0].Message.Content), tps, nil
}

func gb(b int64) float64 { return float64(b) / (1 << 30) }

func or(a, b string) string {
	if a != "" {
		return a
	}
	return b
}

func orUnknown(arch string) string {
	if arch == "" {
		return "architecture not read"
	}
	return arch
}

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