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
}

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

// 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) (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")
		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 {
		outcome(ctx, "cached")
		return target, nil
	}

	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()

	outcome(ctx, "downloaded")
	call.path, call.err = f.fetch(ctx, source, target)
	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)
}

// 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) (_ 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(Progress{Source: reportAs, BytesDone: existing, BytesTotal: total})
	written, err := f.copyWithProgress(ctx, out, resp.Body, reportAs, existing, total)
	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,
) (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(Progress{Source: source, BytesDone: done, BytesTotal: total})
				last = time.Now()
			}
		}
		if errors.Is(err, io.EOF) {
			f.report(Progress{Source: source, BytesDone: done, BytesTotal: total})
			return done, nil
		}
		if err != nil {
			return done, err
		}
	}
}

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
}