llmgw/internal/control/modelsvc/pin.go

package modelsvc

import (
	"context"
	"fmt"
	"regexp"
	"slices"
	"strings"

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

// A model added with weights from the registry is pinned to the commit they
// were at: a file named by hf: or by a resolve URL on a branch becomes that
// commit's URL, and a repository an engine downloads itself is given
// --revision. The registry's sha256 of each file goes in the model's origin,
// for the nodes to check what they fetch against. Without it, a new upload
// to the repository reached the nodes that fetched after it, unseen, beside
// nodes that still had the old one.

// Registry is where a repository's commit and files are looked up.
type Registry interface {
	RepoAt(ctx context.Context, id, rev string) (hf.Repo, error)
	Base() string
}

// SetRegistry has every model added from now on pinned against it.
func (s *Service) SetRegistry(r Registry) { s.registry = r }

// engineDownloads are the engines that download a repository named by its id
// themselves, and take --revision.
var engineDownloads = map[string]bool{"vllm": true, "sglang": true}

var repoID = regexp.MustCompile(`^[A-Za-z0-9][\w.-]*/[\w.-]+$`)

// pin pins a model's registry sources, in place. One already given an origin
// — served through the mirror, which records it — is left as it is.
func pin(ctx context.Context, reg Registry, m *catalog.Model) error {
	if m.Origin != nil {
		return nil
	}
	repos := map[string]hf.Repo{}
	lookup := func(id, rev string) (hf.Repo, error) {
		key := id + "@" + rev
		if r, ok := repos[key]; ok {
			return r, nil
		}
		r, err := reg.RepoAt(ctx, id, rev)
		if err != nil {
			return r, fmt.Errorf("model %q: looking up %s in the registry, to pin it to a commit: %w", m.Name, id, err)
		}
		repos[key] = r
		return r, nil
	}
	record := func(r hf.Repo, files ...hf.File) {
		if m.Origin == nil {
			m.Origin = &catalog.Origin{Repo: r.ID, Revision: r.Revision}
		}
		if m.Origin.Repo != r.ID || m.Origin.Revision != r.Revision {
			return // one origin a model; a second repository is pinned, not recorded
		}
		for _, f := range files {
			m.Origin.Files = append(m.Origin.Files, catalog.OriginFile{Path: f.Path, SHA256: f.SHA256, Bytes: f.Bytes})
		}
	}

	// One file: its commit's URL.
	pinned := map[string]string{}
	for _, src := range m.WeightSources() {
		id, rev, file, ok := hf.FileSource(src, reg.Base())
		if !ok {
			continue
		}
		r, err := lookup(id, rev)
		if err != nil {
			return err
		}
		i := slices.IndexFunc(r.Files, func(f hf.File) bool { return f.Path == file })
		if i < 0 {
			return fmt.Errorf("model %q: %s has no file %s at %s", m.Name, id, file, r.ShortRevision())
		}
		pinned[src] = r.DownloadURL(reg.Base(), file)
		record(r, r.Files[i])
	}
	if to, ok := pinned[m.Source]; ok {
		m.Source = to
	}
	for i := range m.Profiles {
		if to, ok := pinned[m.Profiles[i].Source]; ok {
			m.Profiles[i].Source = to
		}
	}

	// A repository an engine downloads: --revision.
	for i := range m.Profiles {
		p := &m.Profiles[i]
		src := m.SourceOf(*p)
		if !engineDownloads[p.Engine] || !repoID.MatchString(src) {
			continue
		}
		r, err := lookup(src, argValue(p.Args, "--revision"))
		if err != nil {
			return err
		}
		if argValue(p.Args, "--revision") == "" {
			p.Args = append(slices.Clone(p.Args), "--revision", r.Revision)
		}
		var weights []hf.File
		for _, f := range r.Files {
			if f.IsWeights() {
				weights = append(weights, f)
			}
		}
		if m.Origin == nil || m.Origin.Repo != r.ID || len(m.Origin.Files) == 0 {
			record(r, weights...)
		}
	}
	return nil
}

// argValue is a flag's value in arguments, "" when absent.
func argValue(args []string, flag string) string {
	for i, a := range args {
		if a == flag && i+1 < len(args) {
			return args[i+1]
		}
		if v, ok := strings.CutPrefix(a, flag+"="); ok {
			return v
		}
	}
	return ""
}