mirror of
https://github.com/mudler/LocalAI.git
synced 2026-07-31 18:38:23 -04:00
feat: materialize Hugging Face model artifacts (#10825)
* feat(config): add model artifact source contract Assisted-by: Codex:GPT-5 [Codex] * feat(downloader): add authenticated raw-byte progress Assisted-by: Codex:GPT-5 [Codex] * feat(huggingface): resolve immutable snapshot manifests Assisted-by: Codex:GPT-5 [Codex] * feat(models): add artifact storage primitives Assisted-by: Codex:GPT-5 [Codex] * feat(models): materialize pinned Hugging Face snapshots Assisted-by: Codex:GPT-5 [Codex] * feat(models): bind managed snapshots at runtime Assisted-by: Codex:GPT-5 [Codex] * feat(gallery): materialize model artifacts during install Assisted-by: Codex:GPT-5 [Codex] * feat(gallery): declare managed Hugging Face artifacts Assisted-by: Codex:GPT-5 [Codex] * feat(models): preload managed model artifacts Assisted-by: Codex:GPT-5 [Codex] * fix(gallery): retain shared artifact caches on delete Assisted-by: Codex:GPT-5 [Codex] * feat(models): report artifact acquisition progress Assisted-by: Codex:GPT-5 [Codex] * refactor(backends): load managed models from ModelFile Assisted-by: Codex:GPT-5 [Codex] * refactor(backends): load staged speech model snapshots Assisted-by: Codex:GPT-5 [Codex] * refactor(backends): use staged snapshots in engine backends Assisted-by: Codex:GPT-5 [Codex] * test(distributed): cover staged artifact snapshots Assisted-by: Codex:GPT-5 [Codex] * docs: explain managed model artifacts Assisted-by: Codex:GPT-5 [Codex] * docs: add product design context Assisted-by: Codex:GPT-5 [Codex] * feat(ui): show model artifact download progress Assisted-by: Codex:GPT-5 [Codex] * Eagerly materialize Hugging Face artifacts Materialize HF-backed model references as managed GGUF artifacts during load, with lazy download retained only as fallback. Assisted-by: Codex:GPT-5 [shell] * Refactor HF downloads through a shared executor Assisted-by: Codex:GPT-5 [shell] * drop Signed-off-by: Ettore Di Giacinto <mudler@localai.io> --------- Signed-off-by: Ettore Di Giacinto <mudler@localai.io> Co-authored-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
LocalAI [bot]
359
pkg/modelartifacts/materializer.go
Normal file
359
pkg/modelartifacts/materializer.go
Normal file
@@ -0,0 +1,359 @@
|
||||
package modelartifacts
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha1"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"path"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gofrs/flock"
|
||||
"github.com/mudler/xlog"
|
||||
|
||||
"github.com/mudler/LocalAI/pkg/downloader"
|
||||
hfapi "github.com/mudler/LocalAI/pkg/huggingface-api"
|
||||
)
|
||||
|
||||
type SnapshotResolver interface {
|
||||
ResolveSnapshot(context.Context, hfapi.SnapshotRequest) (hfapi.Snapshot, error)
|
||||
}
|
||||
|
||||
type Manager struct {
|
||||
resolver SnapshotResolver
|
||||
huggingFaceToken string
|
||||
}
|
||||
|
||||
type ManagerOption func(*Manager)
|
||||
|
||||
type Result struct {
|
||||
Spec Spec
|
||||
RelativePath string
|
||||
Manifest Manifest
|
||||
CacheHit bool
|
||||
}
|
||||
|
||||
func WithHuggingFaceToken(token string) ManagerOption {
|
||||
return func(manager *Manager) { manager.huggingFaceToken = token }
|
||||
}
|
||||
|
||||
func NewManager(resolver SnapshotResolver, options ...ManagerOption) *Manager {
|
||||
manager := &Manager{resolver: resolver}
|
||||
for _, option := range options {
|
||||
option(manager)
|
||||
}
|
||||
return manager
|
||||
}
|
||||
|
||||
func NewDefaultManager(options ...ManagerOption) *Manager {
|
||||
client := hfapi.NewClient()
|
||||
client.SetBaseURL(strings.TrimRight(downloader.HF_ENDPOINT, "/") + "/api/models")
|
||||
return NewManager(client, options...)
|
||||
}
|
||||
|
||||
func committedResult(modelsPath string, spec Spec) (Result, bool) {
|
||||
if spec.Resolved == nil || spec.Resolved.CacheKey == "" {
|
||||
return Result{}, false
|
||||
}
|
||||
layout, err := LayoutFor(modelsPath, spec)
|
||||
if err != nil {
|
||||
return Result{}, false
|
||||
}
|
||||
manifest, err := ReadManifest(layout.Manifest)
|
||||
if err != nil || manifest.Artifact.Resolved == nil || manifest.Artifact.Resolved.CacheKey != spec.Resolved.CacheKey {
|
||||
return Result{}, false
|
||||
}
|
||||
specKey, err := CacheKey(spec)
|
||||
if err != nil || specKey != spec.Resolved.CacheKey {
|
||||
return Result{}, false
|
||||
}
|
||||
manifestKey, err := CacheKey(manifest.Artifact)
|
||||
if err != nil || manifestKey != spec.Resolved.CacheKey || len(manifest.Files) == 0 {
|
||||
return Result{}, false
|
||||
}
|
||||
for _, file := range manifest.Files {
|
||||
info, err := os.Stat(filepath.Join(layout.Snapshot, filepath.FromSlash(file.Path)))
|
||||
if err != nil || !info.Mode().IsRegular() || info.Size() != file.Size {
|
||||
return Result{}, false
|
||||
}
|
||||
}
|
||||
relative, err := RelativeSnapshotPath(spec.Resolved.CacheKey)
|
||||
if err != nil {
|
||||
return Result{}, false
|
||||
}
|
||||
return Result{Spec: spec, RelativePath: relative, Manifest: manifest, CacheHit: true}, true
|
||||
}
|
||||
|
||||
func (m *Manager) Ensure(ctx context.Context, modelsPath string, spec Spec) (Result, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
normalized, err := spec.Normalize()
|
||||
if err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
if cached, ok := committedResult(modelsPath, normalized); ok {
|
||||
return cached, nil
|
||||
}
|
||||
if m == nil || m.resolver == nil {
|
||||
return Result{}, fmt.Errorf("artifact materializer has no snapshot resolver")
|
||||
}
|
||||
token := ""
|
||||
if normalized.Source.TokenEnv == HuggingFaceTokenEnv {
|
||||
token = m.huggingFaceToken
|
||||
if token == "" {
|
||||
return Result{}, fmt.Errorf("artifact requires non-empty %s", HuggingFaceTokenEnv)
|
||||
}
|
||||
}
|
||||
revision := normalized.Source.Revision
|
||||
if normalized.Resolved != nil {
|
||||
revision = normalized.Resolved.Revision
|
||||
}
|
||||
ReportProgress(ctx, ProgressEvent{Phase: PhaseResolving, Artifact: normalized.Name})
|
||||
snapshot, err := m.resolver.ResolveSnapshot(ctx, hfapi.SnapshotRequest{
|
||||
Repo: normalized.Source.Repo, Revision: revision, Token: token,
|
||||
AllowPatterns: normalized.Source.AllowPatterns, IgnorePatterns: normalized.Source.IgnorePatterns,
|
||||
})
|
||||
if err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
if normalized.Resolved != nil && (snapshot.Endpoint != normalized.Resolved.Endpoint || snapshot.ResolvedRevision != normalized.Resolved.Revision) {
|
||||
return Result{}, fmt.Errorf("resolved artifact identity changed; reinstall the model")
|
||||
}
|
||||
normalized.Resolved = &Resolved{Endpoint: snapshot.Endpoint, Revision: snapshot.ResolvedRevision}
|
||||
cacheKey, err := CacheKey(normalized)
|
||||
if err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
normalized.Resolved.CacheKey = cacheKey
|
||||
layout, err := LayoutFor(modelsPath, normalized)
|
||||
if err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
if err := os.MkdirAll(filepath.Dir(layout.Lock), 0o750); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
artifactLock := flock.New(layout.Lock)
|
||||
locked, err := artifactLock.TryLockContext(ctx, 100*time.Millisecond)
|
||||
if err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
if !locked {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
return Result{}, fmt.Errorf("artifact lock was not acquired")
|
||||
}
|
||||
defer func() {
|
||||
if err := artifactLock.Unlock(); err != nil {
|
||||
xlog.Warn("failed to unlock model artifact", "lock", layout.Lock, "error", err)
|
||||
}
|
||||
}()
|
||||
if err := os.Chmod(layout.Lock, 0o600); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
if cached, ok := committedResult(modelsPath, normalized); ok {
|
||||
return cached, nil
|
||||
}
|
||||
if err := removeInvalidFinal(layout); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
return m.materializeLocked(ctx, modelsPath, normalized, snapshot, token, layout)
|
||||
}
|
||||
|
||||
func removeInvalidFinal(layout Layout) error {
|
||||
root, err := os.OpenRoot(layout.Root)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() { _ = root.Close() }()
|
||||
relative, err := filepath.Rel(layout.Root, layout.Final)
|
||||
if err != nil || filepath.Dir(relative) != "huggingface" || !cacheKeyPattern.MatchString(filepath.Base(relative)) {
|
||||
return fmt.Errorf("refusing to remove invalid artifact path %q", layout.Final)
|
||||
}
|
||||
return root.RemoveAll(relative)
|
||||
}
|
||||
|
||||
func (m *Manager) materializeLocked(ctx context.Context, modelsPath string, spec Spec, snapshot hfapi.Snapshot, token string, layout Layout) (Result, error) {
|
||||
if err := os.MkdirAll(layout.Partial, 0o750); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
root, err := os.OpenRoot(layout.Partial)
|
||||
if err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
rootClosed := false
|
||||
defer func() {
|
||||
if !rootClosed {
|
||||
_ = root.Close()
|
||||
}
|
||||
}()
|
||||
if err := root.MkdirAll(".downloads", 0o750); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
if err := root.MkdirAll("snapshot", 0o750); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
totalBytes := int64(0)
|
||||
if len(snapshot.Files) == 0 {
|
||||
return Result{}, fmt.Errorf("resolved snapshot contains no selected files")
|
||||
}
|
||||
seenPaths := make(map[string]struct{}, len(snapshot.Files))
|
||||
for _, file := range snapshot.Files {
|
||||
if err := ValidateRelativeHubPath(file.Path); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
if file.Size < 0 || totalBytes > int64(^uint64(0)>>1)-file.Size {
|
||||
return Result{}, fmt.Errorf("invalid aggregate snapshot size")
|
||||
}
|
||||
if _, exists := seenPaths[file.Path]; exists {
|
||||
return Result{}, fmt.Errorf("duplicate Hub path %q", file.Path)
|
||||
}
|
||||
seenPaths[file.Path] = struct{}{}
|
||||
totalBytes += file.Size
|
||||
}
|
||||
|
||||
manifest := Manifest{Version: ManifestVersion, Artifact: spec, Files: make([]ManifestFile, 0, len(snapshot.Files))}
|
||||
completedBytes := int64(0)
|
||||
tasks := make([]downloader.FileTask, 0, len(snapshot.Files))
|
||||
for index, file := range snapshot.Files {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
nameSum := sha256.Sum256([]byte(file.Path))
|
||||
blobRel := path.Join(".downloads", hex.EncodeToString(nameSum[:]))
|
||||
blobAbs := filepath.Join(layout.Partial, filepath.FromSlash(blobRel))
|
||||
file := file
|
||||
taskIndex := index
|
||||
task := downloader.FileTask{
|
||||
URI: downloader.URI(file.URL),
|
||||
Destination: blobAbs,
|
||||
SHA256: file.LFSOID,
|
||||
FileIndex: taskIndex,
|
||||
TotalFiles: len(snapshot.Files),
|
||||
Options: []downloader.DownloadOption{
|
||||
downloader.WithBearerToken(token),
|
||||
downloader.WithTransferProgress(func(event downloader.TransferProgress) {
|
||||
ReportProgress(ctx, ProgressEvent{
|
||||
Phase: PhaseDownloading,
|
||||
Artifact: spec.Name,
|
||||
File: file.Path,
|
||||
CurrentBytes: completedBytes + event.Written,
|
||||
TotalBytes: totalBytes,
|
||||
CompletedFiles: taskIndex,
|
||||
TotalFiles: len(snapshot.Files),
|
||||
})
|
||||
}),
|
||||
},
|
||||
AfterDownload: func(string) error {
|
||||
ReportProgress(ctx, ProgressEvent{
|
||||
Phase: PhaseVerifying,
|
||||
Artifact: spec.Name,
|
||||
File: file.Path,
|
||||
CurrentBytes: completedBytes + file.Size,
|
||||
TotalBytes: totalBytes,
|
||||
CompletedFiles: taskIndex,
|
||||
TotalFiles: len(snapshot.Files),
|
||||
})
|
||||
entry, err := verifyDownloadedFile(blobAbs, file)
|
||||
if err != nil {
|
||||
_ = root.Remove(blobRel)
|
||||
return err
|
||||
}
|
||||
destination := path.Join("snapshot", file.Path)
|
||||
if err := root.MkdirAll(path.Dir(destination), 0o750); err != nil {
|
||||
return err
|
||||
}
|
||||
_ = root.Remove(destination)
|
||||
if err := root.Rename(blobRel, destination); err != nil {
|
||||
return err
|
||||
}
|
||||
manifest.Files = append(manifest.Files, entry)
|
||||
completedBytes += file.Size
|
||||
return nil
|
||||
},
|
||||
}
|
||||
tasks = append(tasks, task)
|
||||
}
|
||||
if err := downloader.DownloadFilesWithContext(ctx, tasks, nil); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
if err := root.RemoveAll(".downloads"); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
encoded, err := json.MarshalIndent(manifest, "", " ")
|
||||
if err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
if err := root.WriteFile("manifest.json.tmp", append(encoded, '\n'), 0o644); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
if err := root.Rename("manifest.json.tmp", "manifest.json"); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
if err := root.Close(); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
rootClosed = true
|
||||
if err := os.MkdirAll(filepath.Dir(layout.Final), 0o750); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
ReportProgress(ctx, ProgressEvent{Phase: PhaseCommitting, Artifact: spec.Name, CurrentBytes: totalBytes, TotalBytes: totalBytes, CompletedFiles: len(snapshot.Files), TotalFiles: len(snapshot.Files)})
|
||||
if err := os.Rename(layout.Partial, layout.Final); err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
relative, err := RelativeSnapshotPath(spec.Resolved.CacheKey)
|
||||
if err != nil {
|
||||
return Result{}, err
|
||||
}
|
||||
return Result{Spec: spec, RelativePath: relative, Manifest: manifest}, nil
|
||||
}
|
||||
|
||||
func verifyDownloadedFile(fileName string, source hfapi.SnapshotFile) (ManifestFile, error) {
|
||||
file, err := os.Open(fileName)
|
||||
if err != nil {
|
||||
return ManifestFile{}, err
|
||||
}
|
||||
defer func() { _ = file.Close() }()
|
||||
sha256Hash := sha256.New()
|
||||
gitHash := sha1.New()
|
||||
if _, err := fmt.Fprintf(gitHash, "blob %d%c", source.Size, byte(0)); err != nil {
|
||||
return ManifestFile{}, err
|
||||
}
|
||||
if _, err := io.Copy(io.MultiWriter(sha256Hash, gitHash), file); err != nil {
|
||||
return ManifestFile{}, err
|
||||
}
|
||||
info, err := file.Stat()
|
||||
if err != nil {
|
||||
return ManifestFile{}, err
|
||||
}
|
||||
if info.Size() != source.Size {
|
||||
return ManifestFile{}, fmt.Errorf("size mismatch for %q", source.Path)
|
||||
}
|
||||
rawSHA256 := hex.EncodeToString(sha256Hash.Sum(nil))
|
||||
if source.LFSOID != "" {
|
||||
if decoded, err := hex.DecodeString(source.LFSOID); err != nil || len(decoded) != sha256.Size {
|
||||
return ManifestFile{}, fmt.Errorf("invalid LFS SHA-256 for %q", source.Path)
|
||||
}
|
||||
if !strings.EqualFold(rawSHA256, source.LFSOID) {
|
||||
return ManifestFile{}, fmt.Errorf("LFS SHA-256 mismatch for %q", source.Path)
|
||||
}
|
||||
} else if source.BlobOID != "" {
|
||||
if decoded, err := hex.DecodeString(source.BlobOID); err != nil || len(decoded) != sha1.Size {
|
||||
return ManifestFile{}, fmt.Errorf("invalid Git blob OID for %q", source.Path)
|
||||
}
|
||||
if !strings.EqualFold(hex.EncodeToString(gitHash.Sum(nil)), source.BlobOID) {
|
||||
return ManifestFile{}, fmt.Errorf("Git blob OID mismatch for %q", source.Path)
|
||||
}
|
||||
}
|
||||
return ManifestFile{
|
||||
Path: source.Path, Size: source.Size, SHA256: rawSHA256,
|
||||
BlobOID: source.BlobOID, LFSOID: source.LFSOID, XetHash: source.XetHash,
|
||||
}, nil
|
||||
}
|
||||
Reference in New Issue
Block a user