llmgw/internal/mirror/manager.go

package mirror

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

// State is where a model stands in the mirror.
type State string

// The states a model can be in.
const (
	// Fetching means a download is running. It is the state that makes this
	// service worth an API: without it the control plane cannot tell a model
	// that is arriving from one that will never arrive.
	Fetching State = "fetching"
	Present  State = "present"
	Failed   State = "failed"
)

// Status is one model's state, and how far a download has got.
type Status struct {
	Model   string `json:"model"`
	Version string `json:"version"`
	File    string `json:"file"`
	State   State  `json:"state"`
	Source  string `json:"source,omitempty"`

	BytesDone int64 `json:"bytes_done"`
	// BytesTotal is 0 when the upstream server did not say. A percentage would
	// then be invented, and an invented one that stops moving is worse than a
	// byte count that keeps rising.
	BytesTotal int64  `json:"bytes_total,omitempty"`
	SHA256     string `json:"sha256,omitempty"`
	Err        string `json:"error,omitempty"`

	StartedAt  time.Time `json:"started_at,omitzero"`
	FinishedAt time.Time `json:"finished_at,omitzero"`
}

// Percent reports progress, and false when the total is unknown.
func (s Status) Percent() (int, bool) {
	if s.BytesTotal <= 0 {
		return 0, false
	}
	return int(s.BytesDone * 100 / s.BytesTotal), true
}

// Manager runs downloads in the background and remembers how they went.
//
// Background because a model is tens of gigabytes: a request that waited for
// one would time out in every client between here and the operator, and the
// operator would learn nothing about why. So a fetch is started, and its
// progress is asked for.
type Manager struct {
	root string
	log  *slog.Logger

	mu sync.Mutex
	// running is keyed by model/version/file, which is how a second request for
	// a download already in flight joins it rather than starting a rival one
	// writing to the same file.
	//
	// By file, not by version. A model is routinely several files — sharded
	// safetensors, or two quantisations somebody wants side by side — and
	// keyed by version the second shard was refused as "already fetching"
	// until the first had finished. A fifteen-shard model had to be pulled one
	// shard at a time, each waiting on the last.
	running map[string]*Status

	// client fetches sources. It will not reach private addresses unless
	// AllowPrivateSources was called.
	client *http.Client
}

// AllowPrivateSources lets downloads reach loopback, private and link-local
// addresses — a model server on the same network, say. Off by default: the
// file server hands whatever was fetched to anyone, so a fetch that can reach
// the private network is a way to read it.
func (m *Manager) AllowPrivateSources() { m.client = fetchClient(true) }

// runKey identifies one download.
func runKey(e Entry) string { return e.Model + "/" + e.Version + "/" + e.File }

// NewManager prepares a manager over a mirror root.
func NewManager(root string, log *slog.Logger) *Manager {
	if log == nil {
		log = slog.Default()
	}
	return &Manager{root: root, log: log, running: map[string]*Status{}, client: fetchClient(false)}
}

// ErrAlreadyFetching means a download for that model and version is running.
var ErrAlreadyFetching = errors.New("a download for that version is already running")

// Start begins a download and returns once it is under way.
//
// The context bounds the caller's request, not the download: a fetch outlives
// the API call that asked for it, which is the point of doing it here rather
// than making the caller hold a connection open for an hour.
func (m *Manager) Start(source string, e Entry) (Status, error) {
	if source == "" {
		return Status{}, fmt.Errorf("a source URL is required")
	}
	if err := validSource(source); err != nil {
		return Status{}, err
	}
	key := runKey(e)

	m.mu.Lock()
	if st, ok := m.running[key]; ok && st.State == Fetching {
		// Copied under the lock. The download goroutine writes progress into
		// this struct, so dereferencing it after unlocking would be a read
		// racing a write — of a byte count nobody would notice was wrong.
		out := *st
		m.mu.Unlock()
		return out, ErrAlreadyFetching
	}
	st := &Status{
		Model: e.Model, Version: e.Version, File: e.File,
		State: Fetching, Source: source, StartedAt: time.Now(),
	}
	m.running[key] = st
	out := *st
	m.mu.Unlock()

	go m.run(source, e, key)
	return out, nil
}

func (m *Manager) run(source string, e Entry, key string) {
	// Its own context, deliberately unbounded by any request. A download is
	// minutes to hours and the thing that asked for it is long gone.
	ctx := context.Background()

	added, err := addWith(ctx, m.client, m.root, source, e, func(done, total int64) {
		m.mu.Lock()
		if st, ok := m.running[key]; ok {
			st.BytesDone, st.BytesTotal = done, total
		}
		m.mu.Unlock()
	})

	m.mu.Lock()
	defer m.mu.Unlock()
	st, ok := m.running[key]
	if !ok {
		return
	}
	st.FinishedAt = time.Now()
	if err != nil {
		st.State, st.Err = Failed, err.Error()
		// Kept rather than forgotten: a failed fetch that vanished would leave
		// the control plane waiting for a model that is never coming, with
		// nothing to say why.
		m.log.Error("mirror fetch failed", "model", e.Model, "version", e.Version,
			"source", source, "err", err)
		return
	}
	st.State = Present
	st.BytesDone, st.BytesTotal = added.Entry.Bytes, added.Entry.Bytes
	st.SHA256 = added.SHA256
	m.log.Info("mirror fetched a model", "model", e.Model, "version", e.Version,
		"bytes", added.Entry.Bytes, "sha256", added.SHA256)
}

// Status reports one model's state, reading the disk when no download for it
// is remembered.
//
// The disk is the truth: a restart forgets every download, and a model that is
// present should not be reported as unknown merely because this process did not
// fetch it.
// A version of several files is reported as one: fetching while any file still
// is, failed if any failed, and the byte counts summed. The caller asked about
// a version, and a version with one shard still arriving is not present.
func (m *Manager) Status(model, version string) (Status, bool) {
	var out Status
	found := false

	m.mu.Lock()
	for _, st := range m.running {
		if st.Model == model && st.Version == version {
			out, found = merge(out, *st, found)
		}
	}
	m.mu.Unlock()

	entries, err := List(m.root)
	if err != nil {
		return out, found
	}
	for _, e := range entries {
		if e.Model != model || e.Version != version {
			continue
		}
		if m.isRunning(model, version, e.File) {
			continue // already counted, and the live entry knows more
		}
		out, found = merge(out, Status{
			Model: e.Model, Version: e.Version, File: e.File,
			State: Present, BytesDone: e.Bytes, BytesTotal: e.Bytes,
		}, found)
	}
	return out, found
}

func (m *Manager) isRunning(model, version, file string) bool {
	m.mu.Lock()
	defer m.mu.Unlock()
	_, ok := m.running[model+"/"+version+"/"+file]
	return ok
}

// merge folds one file's status into a version's.
//
// The worst state wins: a version is only present once every file of it is, and
// a shard that failed makes the version unusable however many others arrived.
func merge(acc, one Status, have bool) (Status, bool) {
	if !have {
		return one, true
	}
	acc.BytesDone += one.BytesDone
	acc.BytesTotal += one.BytesTotal
	if worse(one.State) > worse(acc.State) {
		acc.State, acc.Err, acc.File = one.State, one.Err, one.File
	}
	if acc.SHA256 != "" && one.SHA256 != "" {
		// Two files, two checksums, and no single value for the version. Left
		// empty rather than showing one of them, which would look like the
		// version's and be compared against the wrong thing.
		acc.SHA256 = ""
	} else if one.SHA256 != "" {
		acc.SHA256 = one.SHA256
	}
	if one.FinishedAt.After(acc.FinishedAt) {
		acc.FinishedAt = one.FinishedAt
	}
	return acc, true
}

func worse(s State) int {
	switch s {
	case Failed:
		return 2
	case Fetching:
		return 1
	default:
		return 0
	}
}

// All reports everything the mirror holds or is fetching.
//
// Both, because a control plane deciding whether to place a model needs to
// know the difference between a file that is here and one that is on its way —
// and to see that the one it asked for is on its way at all.
func (m *Manager) All() ([]Status, error) {
	entries, err := List(m.root)
	if err != nil {
		return nil, err
	}
	// Taken before the disk entries are built so each one can be told what this
	// process knows about it that the filesystem cannot say.
	m.mu.Lock()
	known := make(map[string]Status, len(m.running))
	for key, st := range m.running {
		known[key] = *st
	}
	m.mu.Unlock()

	out := make([]Status, 0, len(entries))
	seen := map[string]struct{}{}
	for _, e := range entries {
		key := e.Model + "/" + e.Version + "/" + e.File
		st := Status{
			Model: e.Model, Version: e.Version, File: e.File,
			State: Present, BytesDone: e.Bytes, BytesTotal: e.Bytes,
		}
		// The checksum and the source are not recoverable from a directory
		// listing, and the checksum is the one field worth having: it is what
		// an operator compares against what upstream published. Dropping it in
		// favour of the disk entry lost it at the exact moment the file became
		// verifiable — the download finishing was what made it disappear.
		if k, ok := known[key]; ok {
			st.SHA256, st.Source = k.SHA256, k.Source
		} else {
			// Fetched before this process started: what was kept beside it.
			path := filepath.Join(m.root, e.Model, e.Version, e.File)
			st.SHA256, st.Source = readChecksum(path), readSource(path)
		}
		out = append(out, st)
		seen[key] = struct{}{}
	}

	for key, st := range known {
		if _, onDisk := seen[key]; onDisk {
			continue
		}
		out = append(out, st)
	}
	return out, nil
}

// Forget drops a remembered download, so a failed one can be retried.
//
// Without it a fetch that failed once would refuse every later attempt at the
// same version, and the only way through would be restarting the mirror.
func (m *Manager) Forget(model, version string) {
	m.mu.Lock()
	defer m.mu.Unlock()
	// Every file of it. The version is the unit an operator retries, and
	// leaving one shard's failure remembered would refuse the retry it was
	// meant to allow.
	for key, st := range m.running {
		if st.Model == model && st.Version == version {
			delete(m.running, key)
		}
	}
}

// RemoveVersion deletes every file under one model version and reports what
// went.
//
// The version is the unit, not the file: a model split across several files is
// useless with one of them missing, so removing part of it leaves a directory
// that looks present and cannot be loaded.
func (m *Manager) RemoveVersion(model, version string) ([]string, error) {
	entries, err := List(m.root)
	if err != nil {
		return nil, err
	}
	var removed []string
	for _, e := range entries {
		if e.Model != model || e.Version != version {
			continue
		}
		if err := Remove(m.root, e); err != nil {
			return removed, err
		}
		removed = append(removed, e.File)
	}
	if len(removed) == 0 {
		return nil, fmt.Errorf("%s/%s: %w", model, version, os.ErrNotExist)
	}
	m.Forget(model, version)
	m.log.Info("mirror removed a version", "model", model, "version", version, "files", len(removed))
	return removed, nil
}