llmgw/internal/agent/weights/fetch.go

// Package weights makes a model's files present on a node.
//
// The simple arrangement, and the default: every node has its own copy. A
// source that is already a local path is used as it stands; one that names a
// remote file is fetched into a cache directory once and reused. Nothing here
// manages storage, decides a retention policy, or knows what a filesystem is —
// it answers one question, which is where the engine should be pointed.
//
// The exotic arrangements are all handled by not acting. Weights on a shared
// filer, an image exported over the fabric, or anything else an operator has
// mounted are simply present, and present means there is nothing to do. That is
// why the interesting case needs no code: the escape hatch from a distribution
// mechanism is a distribution mechanism that checks first.
package weights

import (
	"context"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/trace"

	"gitea.mynodes.xyz/ilker/llmgw/internal/obs"
)

// Progress reports how a fetch is going, for anything that wants to say so.
type Progress struct {
	Source     string
	BytesDone  int64
	BytesTotal int64
	// Phase is empty while the bytes arrive, PhaseChecking while a set's
	// files are checked against the mirror's checksums.
	Phase string
}

// PhaseChecking is a set's files being checked, after they arrived.
const PhaseChecking = "checking"

// Percent is how far along, or -1 when the total is unknown.
func (p Progress) Percent() int {
	if p.BytesTotal <= 0 {
		return -1
	}
	return int(p.BytesDone * 100 / p.BytesTotal)
}

// Fetcher makes sources local.
type Fetcher struct {
	CacheDir string
	Client   *http.Client
	Log      *slog.Logger
	// OnProgress is called occasionally during a fetch. A missing model is the
	// difference between a replica that is slow to start and one that is stuck,
	// and without this the two are indistinguishable from outside.
	OnProgress func(Progress)

	// inflight makes concurrent requests for one source share a single fetch.
	//
	// Eight shards of a tensor-parallel replica start at the same instant and
	// want the same file. Without this they would each open a connection and
	// each write the same bytes, turning one download into eight and, on a
	// shared origin, into a self-inflicted denial of service.
	mu       sync.Mutex
	inflight map[string]*fetchCall
}

type fetchCall struct {
	done chan struct{}
	path string
	err  error
}

// ErrNotLocal means a source names a file that is not here and cannot be
// fetched — a bare path with nothing to fetch it from.
var ErrNotLocal = errors.New("the weights are not on this node and the source says nothing about where to get them")

// Ensure returns a local path for a source, fetching it if it is not here.
//
// A plain path is never fetched, only checked. That is deliberate: a path is
// what an operator writes when the file is already there, whether it arrived on
// local disk, a Lustre mount or an image exported over the fabric, and a
// downloader that second-guessed that would copy a hundred gigabytes that were
// already present.
//
// Traced as "fetch weights", with what it came to — the file was local,
// already cached, downloaded, or downloaded by another shard this one waited
// for — since that is most of the difference between a replica that starts
// in a minute and one that takes an hour.
//
// A source in a pickle format, or a set with such a file in it, is refused
// before anything is fetched unless allowPickle; see Pickled.
func (f *Fetcher) Ensure(ctx context.Context, source string, allowPickle bool) (path string, err error) {
	return f.EnsureChecked(ctx, source, allowPickle, "")
}

// EnsureChecked is Ensure for a file whose sha256 the registry published: a
// copy that is not it, downloaded or cached, is refused rather than served.
// An empty sha256 checks nothing.
func (f *Fetcher) EnsureChecked(ctx context.Context, source string, allowPickle bool, sha256 string) (path string, err error) {
	ctx, span := obs.Tracer(tracerName).Start(ctx, "fetch weights",
		trace.WithAttributes(attribute.String("llmgw.weights.source", redact(source))))
	defer func() {
		span.SetAttributes(attribute.String("llmgw.weights.path", path))
		obs.End(span, err)
	}()
	if !allowPickle && !IsSet(source) && Pickled(source) {
		return "", refusePickle(source)
	}
	return f.ensure(ctx, source, allowPickle, strings.ToLower(sha256))
}

// tracerName names the spans this package opens.
const tracerName = "gitea.mynodes.xyz/ilker/llmgw/internal/agent/weights"

// outcome says on the current span how a source was made present.
func outcome(ctx context.Context, how string) {
	trace.SpanFromContext(ctx).SetAttributes(attribute.String("llmgw.weights.outcome", how))
}

// redact is a source fit for a trace: a URL without its credentials or
// query, which is where a token would be.
func redact(source string) string {
	u, err := url.Parse(source)
	if err != nil || u.Scheme == "" || u.Host == "" {
		return source
	}
	u.User, u.RawQuery, u.Fragment = nil, "", ""
	return u.String()
}

func (f *Fetcher) ensure(ctx context.Context, source string, allowPickle bool, want string) (string, error) {
	if source == "" {
		return "", fmt.Errorf("no source given")
	}
	// A version rather than a file. Everything below assumes one URL is one
	// file, and a set has to be handled before any of it.
	if IsSet(source) {
		outcome(ctx, "set")
		// Shared like a single file: four replicas placed on one node's
		// four cards want the same set at the same instant, and each
		// fetching it wrote the same .part files over one another — three
		// failed renaming a file another had already renamed.
		return f.shared(ctx, source, func() (string, error) { return f.ensureSet(ctx, source, allowPickle) })
	}
	if !isRemote(source) {
		if !isPath(source) {
			// A bare model id — "Qwen/Qwen2.5-7B-Instruct" — is not a file and
			// was never meant to be one. vLLM and SGLang resolve those
			// themselves, and doing it for them would be both a duplicate
			// downloader and a change in documented behaviour.
			outcome(ctx, "left to the engine")
			return source, nil
		}
		if _, err := os.Stat(source); err == nil {
			outcome(ctx, "local")
			return source, nil
		}
		return "", fmt.Errorf("%w: %s", ErrNotLocal, source)
	}

	target, err := f.cachePath(source)
	if err != nil {
		return "", err
	}
	if _, err := os.Stat(target); err == nil {
		if ok, err := f.cachedIs(target, want); err != nil {
			return "", err
		} else if ok {
			outcome(ctx, "cached")
			return target, nil
		}
		// Not the file the catalog names: fetched again below, and checked.
	}

	return f.shared(ctx, source, func() (string, error) {
		outcome(ctx, "downloaded")
		return f.fetchChecked(ctx, source, target, want)
	})
}

// shared runs fetch once for every caller asking for source at the same
// time; the others wait for its answer.
func (f *Fetcher) shared(ctx context.Context, source string, fetch func() (string, error)) (string, error) {
	f.mu.Lock()
	if f.inflight == nil {
		f.inflight = map[string]*fetchCall{}
	}
	if call, running := f.inflight[source]; running {
		f.mu.Unlock()
		outcome(ctx, "fetched by another shard")
		select {
		case <-call.done:
			return call.path, call.err
		case <-ctx.Done():
			return "", ctx.Err()
		}
	}
	call := &fetchCall{done: make(chan struct{})}
	f.inflight[source] = call
	f.mu.Unlock()

	call.path, call.err = fetch()
	close(call.done)

	f.mu.Lock()
	delete(f.inflight, source)
	f.mu.Unlock()
	return call.path, call.err
}

// fetch downloads a source to a target, resuming and verifying.
func (f *Fetcher) fetch(ctx context.Context, source, target string) (string, error) {
	return f.fetchAs(ctx, source, source, target, setShare{})
}

// fetchChecked is fetch for a file whose sha256 is known: a download that
// is not it is deleted and refused, naming both, and one that is has its
// checksum kept beside it so a cached copy is not read again at each start.
func (f *Fetcher) fetchChecked(ctx context.Context, source, target, want string) (string, error) {
	got, err := f.fetch(ctx, source, target)
	if err != nil || want == "" {
		return got, err
	}
	sum, err := fileSHA256(target)
	if err != nil {
		return "", err
	}
	if sum != want {
		_ = os.Remove(target)
		return "", fmt.Errorf("%s arrived with sha256 %s, and the registry published %s for it: refused and deleted",
			redact(source), sum, want)
	}
	_ = writeChecksum(target, sum)
	return target, nil
}

// cachedIs says whether a cached file is the one whose sha256 is want: by
// the checksum kept beside it, or worked out once and then kept. One that
// is not is deleted, to be fetched again. Any file is, when want is empty.
func (f *Fetcher) cachedIs(target, want string) (bool, error) {
	kept := readChecksum(target)
	if want == "" || kept == want {
		return true, nil
	}
	if kept != "" {
		// Checked when it was fetched, as another file: the same address
		// gives the same bytes again, and another model may be serving
		// these. Refused, and left.
		return false, fmt.Errorf("the copy of this file here has sha256 %s, as fetched, and the catalog names %s: refused",
			kept, want)
	}
	sum, err := fileSHA256(target)
	if err != nil {
		return false, err
	}
	if sum == want {
		_ = writeChecksum(target, sum)
		return true, nil
	}
	if f.Log != nil {
		f.Log.Warn("a cached file is not the one the catalog names; fetching it again",
			"path", target, "sha256", sum, "want", want)
	}
	_ = os.Remove(checksumPath(target))
	return false, os.Remove(target)
}

// checksumPath is where a cached file's sha256 is kept: beside it, hidden,
// so no scan takes it for weights.
func checksumPath(file string) string {
	return filepath.Join(filepath.Dir(file), "."+filepath.Base(file)+".sha256")
}

func writeChecksum(file, sum string) error {
	return os.WriteFile(checksumPath(file), []byte(sum+"\n"), 0o644) //nolint:gosec // a checksum, for anyone to read
}

func readChecksum(file string) string {
	b, err := os.ReadFile(checksumPath(file)) //nolint:gosec // beside a file this fetcher wrote
	if err != nil {
		return ""
	}
	return strings.TrimSpace(string(b))
}

// ChecksumPath is where the sha256 of a cached file is kept, for whoever
// deletes the file to delete with it.
func ChecksumPath(file string) string { return checksumPath(file) }

// fetchAs downloads one URL but reports progress under another name.
//
// The two differ for a set: the shards waiting on it asked for the version, not
// for the fourth file in it, and progress reported under the file's own URL
// reaches nobody. What they see is one number climbing across the whole set.
//
// Each file is a "download" span, with how many bytes it took and where a
// resumed one started from.
func (f *Fetcher) fetchAs(ctx context.Context, source, reportAs, target string, share setShare) (_ string, err error) {
	ctx, span := obs.Tracer(tracerName).Start(ctx, "download",
		trace.WithSpanKind(trace.SpanKindClient),
		trace.WithAttributes(attribute.String("llmgw.weights.source", redact(source))))
	defer func() { obs.End(span, err) }()

	url, err := resolveURL(source)
	if err != nil {
		return "", err
	}
	if err := os.MkdirAll(filepath.Dir(target), 0o755); err != nil {
		return "", err
	}

	// Written beside the target and renamed at the end. A partial file under
	// the real name would be found by the next scan, reported as present, and
	// handed to an engine that fails on it minutes later — the exact failure
	// this whole area exists to prevent.
	part := target + ".part"
	existing := int64(0)
	if st, err := os.Stat(part); err == nil {
		existing = st.Size()
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return "", err
	}
	if existing > 0 {
		// Resume. A forty gigabyte fetch that fails at ninety percent must not
		// start again from nothing.
		req.Header.Set("Range", fmt.Sprintf("bytes=%d-", existing))
	}

	client := f.Client
	if client == nil {
		client = &http.Client{}
	}
	resp, err := client.Do(req)
	if err != nil {
		return "", fmt.Errorf("fetching %s: %w", source, err)
	}
	defer resp.Body.Close()

	switch resp.StatusCode {
	case http.StatusOK:
		existing = 0 // the server ignored the range; start over
	case http.StatusPartialContent:
	default:
		return "", fmt.Errorf("fetching %s: %s returned %s", source, url, resp.Status)
	}

	flags := os.O_CREATE | os.O_WRONLY
	if existing > 0 {
		flags |= os.O_APPEND
	} else {
		flags |= os.O_TRUNC
	}
	out, err := os.OpenFile(part, flags, 0o644)
	if err != nil {
		return "", err
	}

	total := existing + resp.ContentLength
	f.report(share.of(Progress{Source: reportAs, BytesDone: existing, BytesTotal: total}))
	written, err := f.copyWithProgress(ctx, out, resp.Body, reportAs, existing, total, share)
	span.SetAttributes(
		attribute.Int64("llmgw.weights.resumed_from", existing),
		attribute.Int64("llmgw.weights.bytes", written),
	)
	closeErr := out.Close()
	if err != nil {
		// The part file is kept on purpose, so the next attempt resumes.
		return "", fmt.Errorf("fetching %s after %d bytes: %w", source, written, err)
	}
	if closeErr != nil {
		return "", closeErr
	}

	if err := os.Rename(part, target); err != nil {
		return "", err
	}
	if f.Log != nil {
		f.Log.Info("weights fetched", "source", source, "path", target, "bytes", written)
	}
	return target, nil
}

func (f *Fetcher) copyWithProgress(ctx context.Context, dst io.Writer, src io.Reader,
	source string, done, total int64, share setShare,
) (int64, error) {
	buf := make([]byte, 4<<20)
	last := time.Now()
	for {
		if err := ctx.Err(); err != nil {
			return done, err
		}
		n, err := src.Read(buf)
		if n > 0 {
			if _, werr := dst.Write(buf[:n]); werr != nil {
				return done, werr
			}
			done += int64(n)
			// Occasionally, not per read: a four-megabyte buffer over a
			// hundred gigabytes is twenty-five thousand callbacks, and nobody
			// needs to be told that often.
			if time.Since(last) > 2*time.Second {
				f.report(share.of(Progress{Source: source, BytesDone: done, BytesTotal: total}))
				last = time.Now()
			}
		}
		if errors.Is(err, io.EOF) {
			f.report(share.of(Progress{Source: source, BytesDone: done, BytesTotal: total}))
			return done, nil
		}
		if err != nil {
			return done, err
		}
	}
}

// setShare places one file's download within its set's: the bytes of the
// files before it, and the whole set's. Zero for a file on its own.
type setShare struct{ before, total int64 }

// of is a file's progress as the set's.
func (s setShare) of(p Progress) Progress {
	if s.total > 0 {
		p.BytesDone, p.BytesTotal = s.before+p.BytesDone, s.total
	}
	return p
}

func (f *Fetcher) report(p Progress) {
	if f.OnProgress != nil {
		f.OnProgress(p)
	}
}

// cachePath is where a remote source lands.
//
// Named by a hash of the source with the original filename kept on the end, so
// two models with the same filename from different places cannot collide while
// a human looking at the directory can still tell what is in it.
func (f *Fetcher) cachePath(source string) (string, error) {
	if f.CacheDir == "" {
		return "", fmt.Errorf("no cache directory configured, so %s cannot be fetched", source)
	}
	return CachePathFor(f.CacheDir, source), nil
}

// isPath reports whether a source names a file on this machine.
//
// By its leading character, because that is the only honest signal available:
// "Qwen/Qwen2.5-7B-Instruct" and "models/qwen.gguf" are indistinguishable as
// strings, and an operator writing a local file writes an absolute path. The
// cost of being wrong in this direction is that a relative path is passed to
// the engine untouched, which is what happened before this code existed.
func isPath(source string) bool {
	return strings.HasPrefix(source, "/") ||
		strings.HasPrefix(source, "./") ||
		strings.HasPrefix(source, "../")
}

// isRemote reports whether a source has to be fetched at all.
func isRemote(source string) bool {
	return strings.HasPrefix(source, "hf:") ||
		strings.HasPrefix(source, "http://") ||
		strings.HasPrefix(source, "https://")
}

// resolveURL turns a source into something to GET.
//
// hf: names one file in a repository, which is what a GGUF is. A whole
// repository of shards is deliberately not handled here: vLLM and SGLang
// download those themselves and do it well, and reimplementing a snapshot
// downloader to compete with them would be work spent to arrive back where we
// started.
func resolveURL(source string) (string, error) {
	if rest, ok := strings.CutPrefix(source, "hf:"); ok {
		parts := strings.Split(rest, "/")
		if len(parts) < 3 {
			return "", fmt.Errorf(
				"%s names a repository, not a file; give the file as hf:org/repo/model.gguf, "+
					"or use the repository id directly and let the engine download it", source)
		}
		repo := strings.Join(parts[:2], "/")
		file := strings.Join(parts[2:], "/")
		return "https://huggingface.co/" + repo + "/resolve/main/" + file, nil
	}
	if _, err := url.Parse(source); err != nil {
		return "", fmt.Errorf("%s is not a usable source: %w", source, err)
	}
	return source, nil
}