llmgw/internal/agent/discovery/weights.go

package discovery

import (
	"io/fs"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"syscall"
)

// LocalWeights is one model file or directory a node can see.
type LocalWeights struct {
	Path      string `json:"path"`
	SizeBytes int64  `json:"size_bytes"`
	// Shared marks weights on a filesystem this node did not have to fill.
	//
	// Worth distinguishing: twenty nodes each holding their own copy and twenty
	// nodes seeing one filer are the same list of paths and completely
	// different situations — one has twenty copies of a 141 GB model and the
	// other has one.
	Shared bool `json:"shared,omitempty"`
}

// ScanWeights lists the model files under the given directories.
//
// Storage is the part of this system that will differ at every installation —
// a Lustre mount, an image exported over the fabric, local NVMe, an rsync from
// a laptop — so llmgw does not try to own it. It observes the result instead:
// whatever put the file there, the question placement needs answered is only
// whether it is present.
//
// Scanning nothing is the default and constrains nothing. A node with no
// directories configured behaves exactly as every node did before this existed,
// which is what makes the whole feature optional.
func ScanWeights(dirs []string) []LocalWeights {
	var out []LocalWeights
	seen := map[string]struct{}{}

	for _, dir := range dirs {
		if dir == "" {
			continue
		}
		shared := isNetworkFS(dir)
		entries, err := os.ReadDir(dir)
		if err != nil {
			// A directory that is not there is not an error to report upward:
			// a mount that has not come back yet looks exactly like this, and
			// the node is simply one without those weights until it does.
			continue
		}
		for _, e := range entries {
			path := filepath.Join(dir, e.Name())
			if _, dup := seen[path]; dup {
				continue
			}
			size, ok := weightSize(path, e)
			if !ok {
				continue
			}
			seen[path] = struct{}{}
			out = append(out, LocalWeights{Path: path, SizeBytes: size, Shared: shared})
		}
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Path < out[j].Path })
	return out
}

// weightSize measures a candidate and says whether it is one at all.
//
// A model is either a single file — a .gguf, usually — or a directory of
// shards and a config, which is what every HuggingFace checkout looks like.
// Anything else in the directory is somebody's notes and is not reported.
func weightSize(path string, e fs.DirEntry) (int64, bool) {
	if e.IsDir() {
		// A directory counts only if it actually holds weights. An empty one
		// is a half-finished download, and reporting it as present would have
		// placement choose a node whose engine then fails to load.
		var total int64
		var found bool
		_ = filepath.WalkDir(path, func(p string, d fs.DirEntry, err error) error {
			if err != nil || d.IsDir() {
				return nil //nolint:nilerr // an unreadable subtree is not fatal
			}
			if isWeightFile(p) {
				found = true
			}
			if info, err := d.Info(); err == nil {
				total += info.Size()
			}
			return nil
		})
		return total, found
	}
	if !isWeightFile(path) {
		return 0, false
	}
	info, err := e.Info()
	if err != nil {
		return 0, false
	}
	return info.Size(), true
}

func isWeightFile(path string) bool {
	switch strings.ToLower(filepath.Ext(path)) {
	case ".gguf", ".safetensors", ".bin", ".pt", ".pth":
		return true
	}
	return false
}

// isNetworkFS reports whether a path is on a filesystem the node did not fill.
//
// By magic number rather than by parsing mounts, because the mount table says
// what was asked for and this says what the kernel actually has. The list is
// the network filesystems that turn up in front of GPUs: NFS, Lustre, GPFS,
// CephFS, GlusterFS and SMB. Anything unrecognised is treated as local, which
// is the conservative direction — a shared filer mistaken for local storage
// costs an unnecessary copy, while local storage mistaken for shared would have
// the cluster believe a model is everywhere when it is on one node.
func isNetworkFS(path string) bool {
	var st syscall.Statfs_t
	if err := syscall.Statfs(path, &st); err != nil {
		return false
	}
	switch uint32(st.Type) { //nolint:unconvert // Type is int64 on some platforms
	case 0x6969, // NFS
		0x0BD00BD0, // Lustre
		0x47504653, // GPFS
		0x00C36400, // CephFS
		0x65735546, // FUSE, which is how GlusterFS and many object mounts appear
		0xFF534D42: // SMB/CIFS
		return true
	}
	return false
}