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]
2026-07-15 01:09:33 +02:00
committed by GitHub
parent d82c38ee77
commit bcc41219f7
83 changed files with 4315 additions and 339 deletions

View File

@@ -0,0 +1,173 @@
package downloader_test
import (
"context"
"errors"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"sync"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"github.com/mudler/LocalAI/pkg/downloader"
)
var _ = Describe("authenticated HTTP downloads", func() {
It("sends bearer auth to the origin and strips it at a cross-host redirect", func() {
var originAuth, cdnAuth string
cdn := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
cdnAuth = r.Header.Get("Authorization")
w.Header().Set("Content-Length", "5")
_, _ = w.Write([]byte("model"))
}))
DeferCleanup(cdn.Close)
origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
originAuth = r.Header.Get("Authorization")
http.Redirect(w, r, cdn.URL+"/weights", http.StatusTemporaryRedirect)
}))
DeferCleanup(origin.Close)
target := filepath.Join(GinkgoT().TempDir(), "weights.bin")
err := downloader.URI(origin.URL+"/resolve").DownloadFileWithContext(
context.Background(),
target,
"",
0,
1,
nil,
downloader.WithBearerToken("secret-token"),
)
Expect(err).NotTo(HaveOccurred())
Expect(originAuth).To(Equal("Bearer secret-token"))
Expect(cdnAuth).To(BeEmpty())
})
It("authenticates range probes and reports resumed bytes", func() {
var mu sync.Mutex
seen := make(map[string]string)
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
mu.Lock()
seen[r.Method] = r.Header.Get("Authorization")
mu.Unlock()
if r.Method == http.MethodHead {
w.Header().Set("Accept-Ranges", "bytes")
w.Header().Set("Content-Length", "6")
return
}
Expect(r.Header.Get("Range")).To(Equal("bytes=3-"))
w.Header().Set("Content-Length", "3")
w.Header().Set("Content-Range", "bytes 3-5/6")
w.WriteHeader(http.StatusPartialContent)
_, _ = w.Write([]byte("def"))
}))
DeferCleanup(server.Close)
target := filepath.Join(GinkgoT().TempDir(), "resume.bin")
Expect(os.WriteFile(target+".partial", []byte("abc"), 0o600)).To(Succeed())
events := []downloader.TransferProgress{}
err := downloader.URI(server.URL+"/weights").DownloadFileWithContext(
context.Background(),
target,
"",
0,
1,
nil,
downloader.WithBearerToken("resume-token"),
downloader.WithTransferProgress(func(event downloader.TransferProgress) {
events = append(events, event)
}),
)
Expect(err).NotTo(HaveOccurred())
Expect(seen).To(HaveKeyWithValue(http.MethodHead, "Bearer resume-token"))
Expect(seen).To(HaveKeyWithValue(http.MethodGet, "Bearer resume-token"))
Expect(events).NotTo(BeEmpty())
Expect(events[len(events)-1].Written).To(Equal(int64(6)))
Expect(events[len(events)-1].Total).To(Equal(int64(6)))
Expect(os.ReadFile(target)).To(Equal([]byte("abcdef")))
})
It("rejects an origin that ignores a resume range", func() {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method == http.MethodHead {
w.Header().Set("Accept-Ranges", "bytes")
w.Header().Set("Content-Length", "6")
return
}
Expect(r.Header.Get("Range")).To(Equal("bytes=3-"))
w.Header().Set("Content-Length", "6")
_, _ = w.Write([]byte("abcdef"))
}))
DeferCleanup(server.Close)
target := filepath.Join(GinkgoT().TempDir(), "ignored-range.bin")
Expect(os.WriteFile(target+".partial", []byte("abc"), 0o600)).To(Succeed())
err := downloader.URI(server.URL).DownloadFileWithContext(
context.Background(),
target,
"",
0,
1,
nil,
)
Expect(err).To(MatchError(ContainSubstring("status 200 instead of 206")))
Expect(target).NotTo(BeAnExistingFile())
Expect(target + ".partial").NotTo(BeAnExistingFile())
})
It("creates resumable partial files with owner-only permissions", func() {
ctx, cancel := context.WithCancel(context.Background())
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Content-Length", "5")
_, _ = w.Write([]byte("model"))
}))
DeferCleanup(server.Close)
target := filepath.Join(GinkgoT().TempDir(), "private.bin")
err := downloader.URI(server.URL).DownloadFileWithContext(
ctx,
target,
"",
0,
1,
nil,
downloader.WithTransferProgress(func(downloader.TransferProgress) {
cancel()
}),
)
Expect(errors.Is(err, context.Canceled)).To(BeTrue())
info, statErr := os.Stat(target + ".partial")
Expect(statErr).NotTo(HaveOccurred())
Expect(info.Mode().Perm()).To(Equal(os.FileMode(0o600)))
})
It("keeps the legacy total empty when the response length is unknown", func() {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.WriteHeader(http.StatusOK)
if flusher, ok := w.(http.Flusher); ok {
flusher.Flush()
}
_, _ = w.Write([]byte("model"))
}))
DeferCleanup(server.Close)
target := filepath.Join(GinkgoT().TempDir(), "unknown-size.bin")
totals := []string{}
err := downloader.URI(server.URL).DownloadFileWithContext(
context.Background(),
target,
"",
0,
1,
func(_, _, total string, _ float64) {
totals = append(totals, total)
},
)
Expect(err).NotTo(HaveOccurred())
Expect(totals).NotTo(BeEmpty())
Expect(totals[len(totals)-1]).To(BeEmpty())
})
})

View File

@@ -0,0 +1,39 @@
package downloader
import "context"
// FileTask describes one download operation and an optional post-download
// hook that runs after the bytes are present on disk. Callers keep any
// higher-level commit logic outside this helper.
type FileTask struct {
URI URI
Destination string
SHA256 string
FileIndex int
TotalFiles int
AfterDownload func(string) error
Options []DownloadOption
}
// DownloadFilesWithContext executes a set of file downloads sequentially.
// The helper centralizes the shared download path so callers only provide
// source/destination metadata and any post-download hook they need.
func DownloadFilesWithContext(ctx context.Context, tasks []FileTask, status func(string, string, string, float64), opts ...DownloadOption) error {
for i := range tasks {
task := tasks[i]
if err := ctx.Err(); err != nil {
return err
}
taskOpts := append([]DownloadOption{}, opts...)
taskOpts = append(taskOpts, task.Options...)
if err := task.URI.DownloadFileWithContext(ctx, task.Destination, task.SHA256, task.FileIndex, task.TotalFiles, status, taskOpts...); err != nil {
return err
}
if task.AfterDownload != nil {
if err := task.AfterDownload(task.Destination); err != nil {
return err
}
}
}
return nil
}

View File

@@ -0,0 +1,44 @@
package downloader_test
import (
"context"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"github.com/mudler/LocalAI/pkg/downloader"
)
var _ = Describe("DownloadFilesWithContext", func() {
It("runs the post-download hook after fetching each file", func() {
payload := []byte("downloaded-bytes")
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, _ = w.Write(payload)
}))
DeferCleanup(server.Close)
dest := filepath.Join(GinkgoT().TempDir(), "model.bin")
hookCalled := false
err := downloader.DownloadFilesWithContext(context.Background(), []downloader.FileTask{{
URI: downloader.URI(server.URL),
Destination: dest,
FileIndex: 0,
TotalFiles: 1,
AfterDownload: func(path string) error {
hookCalled = true
got, err := os.ReadFile(path)
Expect(err).NotTo(HaveOccurred())
Expect(string(got)).To(Equal(string(payload)))
return nil
},
}}, nil)
Expect(err).NotTo(HaveOccurred())
Expect(hookCalled).To(BeTrue())
})
})

View File

@@ -12,10 +12,34 @@ type progressWriter struct {
totalFiles int
written int64
downloadStatus func(string, string, string, float64)
transferSink TransferProgressSink
hash hash.Hash
ctx context.Context
}
func (pw *progressWriter) report() {
if pw.transferSink != nil {
pw.transferSink(TransferProgress{
FileName: pw.fileName,
Written: pw.written,
Total: pw.total,
})
}
if pw.downloadStatus == nil {
return
}
percentage := float64(0)
total := ""
if pw.total > 0 {
percentage = float64(pw.written) / float64(pw.total) * 100
total = formatBytes(pw.total)
if pw.totalFiles > 1 {
percentage = percentage/float64(pw.totalFiles) + float64(pw.fileNo)*100/float64(pw.totalFiles)
}
}
pw.downloadStatus(pw.fileName, formatBytes(pw.written), total, percentage)
}
func (pw *progressWriter) Write(p []byte) (n int, err error) {
// Check for cancellation before writing
if pw.ctx != nil {
@@ -41,24 +65,7 @@ func (pw *progressWriter) Write(p []byte) (n int, err error) {
}
}
if pw.total > 0 {
percentage := float64(pw.written) / float64(pw.total) * 100
if pw.totalFiles > 1 {
// This is a multi-file download
// so we need to adjust the percentage
// to reflect the progress of the whole download
// This is the file pw.fileNo (0-indexed) of pw.totalFiles files. We assume that
// the files before successfully downloaded.
percentage = percentage / float64(pw.totalFiles)
if pw.fileNo > 0 {
percentage += float64(pw.fileNo) * 100 / float64(pw.totalFiles)
}
}
//log.Debug().Msgf("Downloading %s: %s/%s (%.2f%%)", pw.fileName, formatBytes(pw.written), formatBytes(pw.total), percentage)
pw.downloadStatus(pw.fileName, formatBytes(pw.written), formatBytes(pw.total), percentage)
} else {
pw.downloadStatus(pw.fileName, formatBytes(pw.written), "", 0)
}
pw.report()
return
}

View File

@@ -52,8 +52,21 @@ type ImageVerifier interface {
VerifyImage(ctx context.Context, imageRef string) error
}
// TransferProgress reports the raw byte counts for an HTTP download.
// Total is negative when the server does not advertise a response length.
type TransferProgress struct {
FileName string
Written int64
Total int64
}
// TransferProgressSink receives raw byte progress updates for an HTTP download.
type TransferProgressSink func(TransferProgress)
type downloadOptions struct {
verifier ImageVerifier
verifier ImageVerifier
bearerToken string
transferProgress TransferProgressSink
}
// DownloadOption configures DownloadFileWithContext / DownloadFile.
@@ -70,6 +83,17 @@ func WithImageVerifier(v ImageVerifier) DownloadOption {
return func(o *downloadOptions) { o.verifier = v }
}
// WithBearerToken authenticates HTTP download requests with a bearer token.
// The token is stripped if a request redirects to a different origin.
func WithBearerToken(token string) DownloadOption {
return func(o *downloadOptions) { o.bearerToken = token }
}
// WithTransferProgress attaches a sink for raw HTTP download byte progress.
func WithTransferProgress(sink TransferProgressSink) DownloadOption {
return func(o *downloadOptions) { o.transferProgress = sink }
}
func applyDownloadOptions(opts []DownloadOption) downloadOptions {
var o downloadOptions
for _, fn := range opts {
@@ -367,9 +391,28 @@ func calculateHashForPartialFile(file *os.File) (hash.Hash, error) {
// deadline so large downloads are not truncated.
var downloadClient = httpclient.New(httpclient.WithFollowRedirects())
func (uri URI) checkSeverSupportsRangeHeader() (bool, error) {
url := uri.ResolveURL()
resp, err := downloadClient.Head(url)
func newDownloadRequest(
ctx context.Context,
method string,
rawURL string,
bearerToken string,
) (*http.Request, error) {
req, err := http.NewRequestWithContext(ctx, method, rawURL, nil)
if err != nil {
return nil, err
}
if bearerToken != "" {
req.Header.Set("Authorization", "Bearer "+bearerToken)
}
return req, nil
}
func (uri URI) checkServerSupportsRangeHeader(ctx context.Context, bearerToken string) (bool, error) {
req, err := newDownloadRequest(ctx, http.MethodHead, uri.ResolveURL(), bearerToken)
if err != nil {
return false, err
}
resp, err := downloadClient.Do(req)
if err != nil {
return false, err
}
@@ -563,21 +606,22 @@ func (uri URI) DownloadFileWithContext(ctx context.Context, filePath, sha string
xlog.Info("Downloading", "url", url)
req, err := http.NewRequestWithContext(ctx, "GET", url, nil)
req, err := newDownloadRequest(ctx, http.MethodGet, url, dopts.bearerToken)
if err != nil {
return fmt.Errorf("failed to create request for %q: %v", filePath, err)
}
// save partial download to dedicated file
tmpFilePath := filePath + ".partial"
var startPos int64
tmpFileInfo, err := os.Stat(tmpFilePath)
if err == nil && uri.LooksLikeHTTPURL() {
support, err := uri.checkSeverSupportsRangeHeader()
support, err := uri.checkServerSupportsRangeHeader(ctx, dopts.bearerToken)
if err != nil {
return fmt.Errorf("failed to check if uri server supports range header: %v", err)
}
if support {
startPos := tmpFileInfo.Size()
startPos = tmpFileInfo.Size()
req.Header.Set("Range", fmt.Sprintf("bytes=%d-", startPos))
} else {
err := removePartialFile(tmpFilePath)
@@ -622,6 +666,15 @@ func (uri URI) DownloadFileWithContext(ctx context.Context, filePath, sha string
}
//defer resp.Body.Close()
if startPos > 0 && resp.StatusCode != http.StatusPartialContent {
_ = resp.Body.Close()
_ = removePartialFile(tmpFilePath)
return fmt.Errorf(
"resume request for %q returned status %d instead of 206",
filePath,
resp.StatusCode,
)
}
if resp.StatusCode >= 400 {
return fmt.Errorf("failed to download url %q, invalid status code %d", url, resp.StatusCode)
}
@@ -633,7 +686,7 @@ func (uri URI) DownloadFileWithContext(ctx context.Context, filePath, sha string
if DownloadStallTimeout > 0 {
source = newIdleTimeoutReader(resp.Body, DownloadStallTimeout)
}
contentLength = resp.ContentLength
contentLength = resp.ContentLength + startPos
}
defer source.Close()
@@ -644,11 +697,14 @@ func (uri URI) DownloadFileWithContext(ctx context.Context, filePath, sha string
}
// Create and write file
outFile, err := os.OpenFile(tmpFilePath, os.O_APPEND|os.O_RDWR|os.O_CREATE, 0644)
outFile, err := os.OpenFile(tmpFilePath, os.O_APPEND|os.O_RDWR|os.O_CREATE, 0600)
if err != nil {
return fmt.Errorf("failed to create / open file %q: %v", tmpFilePath, err)
}
defer outFile.Close()
if err := outFile.Chmod(0600); err != nil {
return fmt.Errorf("failed to restrict partial file %q permissions: %v", tmpFilePath, err)
}
hash, err := calculateHashForPartialFile(outFile)
if err != nil {
return fmt.Errorf("failed to calculate hash for partial file")
@@ -659,7 +715,9 @@ func (uri URI) DownloadFileWithContext(ctx context.Context, filePath, sha string
hash: hash,
fileNo: fileN,
totalFiles: total,
written: startPos,
downloadStatus: downloadStatus,
transferSink: dopts.transferProgress,
ctx: ctx,
}

View File

@@ -238,8 +238,17 @@ var _ = Describe("Download Test", func() {
}
}
respData = mockData[startPos:endPos]
w.WriteHeader(http.StatusOK)
w.Write(respData)
w.Header().Set("Content-Length", strconv.Itoa(len(respData)))
if rangeString != "" {
w.Header().Set(
"Content-Range",
fmt.Sprintf("bytes %d-%d/%d", startPos, endPos-1, len(mockData)),
)
w.WriteHeader(http.StatusPartialContent)
} else {
w.WriteHeader(http.StatusOK)
}
_, _ = w.Write(respData)
}))
mockServer.EnableHTTP2 = true
mockServer.Start()

View File

@@ -1,6 +1,7 @@
package hfapi
import (
"context"
"encoding/json"
"errors"
"fmt"
@@ -112,6 +113,17 @@ func NewClient() *Client {
}
}
func (c *Client) newRequest(ctx context.Context, method, rawURL, token string) (*http.Request, error) {
req, err := http.NewRequestWithContext(ctx, method, rawURL, nil)
if err != nil {
return nil, err
}
if token != "" {
req.Header.Set("Authorization", "Bearer "+token)
}
return req, nil
}
// SearchModels searches for models using the Hugging Face API
func (c *Client) SearchModels(params SearchParams) ([]Model, error) {
for attempt := 1; attempt <= c.maxRetries; attempt++ {

View File

@@ -0,0 +1,239 @@
package hfapi
import (
"context"
"encoding/json"
"fmt"
"net/http"
"net/url"
"path"
"regexp"
"sort"
"strings"
)
const defaultMaxSnapshotFiles = 100_000
var immutableRevisionPattern = regexp.MustCompile(`^[0-9a-fA-F]{40}$`)
type SnapshotRequest struct {
Repo string
Revision string
Token string
AllowPatterns []string
IgnorePatterns []string
MaxFiles int
}
type Snapshot struct {
Endpoint string
Repo string
RequestedRevision string
ResolvedRevision string
Files []SnapshotFile
}
type SnapshotFile struct {
Path string
Size int64
BlobOID string
LFSOID string
XetHash string
URL string
}
func FilterSnapshotFiles(files []SnapshotFile, allow, ignore []string) ([]SnapshotFile, error) {
selected := make([]SnapshotFile, 0, len(files))
for _, file := range files {
allowed, err := matchesAny(file.Path, allow)
if err != nil {
return nil, err
}
if len(allow) > 0 && !allowed {
continue
}
ignored, err := matchesAny(file.Path, ignore)
if err != nil {
return nil, err
}
if ignored {
continue
}
selected = append(selected, file)
}
sort.Slice(selected, func(i, j int) bool { return selected[i].Path < selected[j].Path })
return selected, nil
}
func matchesAny(filePath string, patterns []string) (bool, error) {
for _, pattern := range patterns {
matched, err := path.Match(pattern, filePath)
if err != nil {
return false, fmt.Errorf("invalid Hub file pattern %q: %w", pattern, err)
}
if !matched && !strings.Contains(pattern, "/") {
matched, err = path.Match(pattern, path.Base(filePath))
if err != nil {
return false, fmt.Errorf("invalid Hub file pattern %q: %w", pattern, err)
}
}
if matched {
return true, nil
}
}
return false, nil
}
type revisionResponse struct {
SHA string `json:"sha"`
}
type snapshotTreeItem struct {
Type string `json:"type"`
OID string `json:"oid"`
Size int64 `json:"size"`
Path string `json:"path"`
LFS *LFSInfo `json:"lfs,omitempty"`
XetHash string `json:"xetHash,omitempty"`
}
func (c *Client) ResolveSnapshot(ctx context.Context, req SnapshotRequest) (Snapshot, error) {
if err := ctx.Err(); err != nil {
return Snapshot{}, err
}
if req.MaxFiles <= 0 {
req.MaxFiles = defaultMaxSnapshotFiles
}
endpoint := strings.TrimSuffix(c.baseURL, "/api/models")
parsed, err := url.Parse(endpoint)
if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") || parsed.Host == "" || parsed.User != nil || parsed.RawQuery != "" || parsed.Fragment != "" {
return Snapshot{}, fmt.Errorf("invalid Hugging Face endpoint")
}
endpoint = strings.TrimRight(parsed.String(), "/")
owner, repo, err := splitRepo(req.Repo)
if err != nil {
return Snapshot{}, err
}
revision := req.Revision
if revision == "" {
revision = "main"
}
metadataURL := fmt.Sprintf("%s/api/models/%s/%s/revision/%s", endpoint, url.PathEscape(owner), url.PathEscape(repo), url.PathEscape(revision))
metadataReq, err := c.newRequest(ctx, http.MethodGet, metadataURL, req.Token)
if err != nil {
return Snapshot{}, err
}
var metadata revisionResponse
if err := c.decodeJSON(metadataReq, &metadata); err != nil {
return Snapshot{}, fmt.Errorf("resolve Hugging Face revision: %w", err)
}
if !immutableRevisionPattern.MatchString(metadata.SHA) {
return Snapshot{}, fmt.Errorf("Hub returned invalid immutable revision")
}
resolvedRevision := strings.ToLower(metadata.SHA)
treeURL := fmt.Sprintf("%s/api/models/%s/%s/tree/%s?recursive=true&expand=true", endpoint, url.PathEscape(owner), url.PathEscape(repo), resolvedRevision)
items := make([]snapshotTreeItem, 0)
for treeURL != "" {
treeReq, err := c.newRequest(ctx, http.MethodGet, treeURL, req.Token)
if err != nil {
return Snapshot{}, err
}
page, next, err := c.decodeTreePage(treeReq)
if err != nil {
return Snapshot{}, err
}
if len(items)+len(page) > req.MaxFiles {
return Snapshot{}, fmt.Errorf("snapshot exceeds maximum file count %d", req.MaxFiles)
}
items = append(items, page...)
treeURL = next
}
files := make([]SnapshotFile, 0, len(items))
for _, item := range items {
if item.Type != "file" {
continue
}
file := SnapshotFile{Path: item.Path, Size: item.Size, BlobOID: item.OID, XetHash: item.XetHash}
if item.LFS != nil {
file.LFSOID = item.LFS.Oid
}
file.URL = fmt.Sprintf("%s/%s/%s/resolve/%s/%s", endpoint, url.PathEscape(owner), url.PathEscape(repo), resolvedRevision, escapeFilePath(item.Path))
files = append(files, file)
}
files, err = FilterSnapshotFiles(files, req.AllowPatterns, req.IgnorePatterns)
if err != nil {
return Snapshot{}, err
}
return Snapshot{Endpoint: endpoint, Repo: req.Repo, RequestedRevision: revision, ResolvedRevision: resolvedRevision, Files: files}, nil
}
func splitRepo(repo string) (string, string, error) {
parts := strings.Split(repo, "/")
if len(parts) != 2 || parts[0] == "" || parts[1] == "" {
return "", "", fmt.Errorf("Hugging Face repo must be owner/repo")
}
return parts[0], parts[1], nil
}
func escapeFilePath(filePath string) string {
parts := strings.Split(filePath, "/")
for i := range parts {
parts[i] = url.PathEscape(parts[i])
}
return strings.Join(parts, "/")
}
func (c *Client) decodeJSON(req *http.Request, dst any) error {
resp, err := c.client.Do(req)
if err != nil {
return err
}
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode != http.StatusOK {
return fmt.Errorf("Hub request returned status %d", resp.StatusCode)
}
return json.NewDecoder(http.MaxBytesReader(nil, resp.Body, 8<<20)).Decode(dst)
}
func (c *Client) decodeTreePage(req *http.Request) ([]snapshotTreeItem, string, error) {
resp, err := c.client.Do(req)
if err != nil {
return nil, "", err
}
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode != http.StatusOK {
return nil, "", fmt.Errorf("Hub tree request returned status %d", resp.StatusCode)
}
var page []snapshotTreeItem
if err := json.NewDecoder(http.MaxBytesReader(nil, resp.Body, 32<<20)).Decode(&page); err != nil {
return nil, "", err
}
next, err := nextPage(resp.Header.Get("Link"), req.URL)
if err != nil {
return nil, "", err
}
return page, next, nil
}
func nextPage(link string, current *url.URL) (string, error) {
for _, entry := range strings.Split(link, ",") {
parts := strings.Split(entry, ";")
if len(parts) >= 2 && strings.TrimSpace(parts[1]) == `rel="next"` {
raw := strings.Trim(strings.TrimSpace(parts[0]), "<>")
next, err := url.Parse(raw)
if err != nil {
return "", fmt.Errorf("invalid Hub pagination link: %w", err)
}
if !next.IsAbs() {
next = current.ResolveReference(next)
}
if next.Scheme != current.Scheme || next.Host != current.Host {
return "", fmt.Errorf("cross-origin pagination link rejected")
}
return next.String(), nil
}
}
return "", nil
}

View File

@@ -0,0 +1,149 @@
package hfapi_test
import (
"context"
"fmt"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"strings"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"github.com/mudler/LocalAI/pkg/downloader"
hfapi "github.com/mudler/LocalAI/pkg/huggingface-api"
)
var _ = Describe("immutable snapshot resolution", func() {
It("pins the metadata sha and follows tree pagination", func() {
const commit = "0123456789abcdef0123456789abcdef01234567"
var server *httptest.Server
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
Expect(r.Header.Get("Authorization")).To(Equal("Bearer test-token"))
switch {
case strings.Contains(r.URL.Path, "/revision/main"):
_, _ = w.Write([]byte(`{"sha":"` + commit + `"}`))
case strings.Contains(r.URL.Path, "/tree/"+commit) && r.URL.Query().Get("cursor") == "page-2":
_, _ = w.Write([]byte(`[{"type":"file","path":"tokenizer/tokenizer.json","size":7,"oid":"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"}]`))
case strings.Contains(r.URL.Path, "/tree/"+commit):
Expect(r.URL.Query().Get("recursive")).To(Equal("true"))
w.Header().Set("Link", fmt.Sprintf(`<%s%s?recursive=true&expand=true&cursor=page-2>; rel="next"`, server.URL, r.URL.Path))
_, _ = w.Write([]byte(`[{"type":"file","path":"model.safetensors","size":11,"oid":"bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb","lfs":{"oid":"cccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccccc","size":11,"pointerSize":130}}]`))
default:
w.WriteHeader(http.StatusNotFound)
}
}))
DeferCleanup(server.Close)
client := hfapi.NewClient()
client.SetBaseURL(server.URL + "/api/models")
snapshot, err := client.ResolveSnapshot(context.Background(), hfapi.SnapshotRequest{
Repo: "owner/repo",
Revision: "main",
Token: "test-token",
MaxFiles: 10,
})
Expect(err).NotTo(HaveOccurred())
Expect(snapshot.Endpoint).To(Equal(server.URL))
Expect(snapshot.ResolvedRevision).To(Equal(commit))
Expect(snapshot.Files).To(HaveLen(2))
Expect(snapshot.Files[0].Path).To(Equal("model.safetensors"))
Expect(snapshot.Files[0].LFSOID).To(HaveLen(64))
Expect(snapshot.Files[1].Path).To(Equal("tokenizer/tokenizer.json"))
Expect(snapshot.Files[0].URL).To(ContainSubstring("/resolve/" + commit + "/model.safetensors"))
})
It("applies deterministic allow and ignore filters", func() {
files := []hfapi.SnapshotFile{
{Path: "config.json"},
{Path: "nested/tokenizer.json"},
{Path: "weights.bin"},
{Path: "weights.safetensors"},
}
filtered, err := hfapi.FilterSnapshotFiles(files, []string{"*.json", "*.safetensors"}, []string{"nested/*"})
Expect(err).NotTo(HaveOccurred())
Expect(filtered).To(ConsistOf(
hfapi.SnapshotFile{Path: "config.json"},
hfapi.SnapshotFile{Path: "weights.safetensors"},
))
})
It("rejects a manifest larger than MaxFiles", func() {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if strings.Contains(r.URL.Path, "/revision/") {
_, _ = w.Write([]byte(`{"sha":"0123456789abcdef0123456789abcdef01234567"}`))
return
}
_, _ = w.Write([]byte(`[{"type":"file","path":"a","size":1},{"type":"file","path":"b","size":1}]`))
}))
DeferCleanup(server.Close)
client := hfapi.NewClient()
client.SetBaseURL(server.URL + "/api/models")
_, err := client.ResolveSnapshot(context.Background(), hfapi.SnapshotRequest{Repo: "owner/repo", Revision: "main", MaxFiles: 1})
Expect(err).To(MatchError(ContainSubstring("maximum file count")))
})
It("honors a cancelled context before metadata access", func() {
ctx, cancel := context.WithCancel(context.Background())
cancel()
_, err := hfapi.NewClient().ResolveSnapshot(ctx, hfapi.SnapshotRequest{Repo: "owner/repo", Revision: "main"})
Expect(err).To(MatchError(context.Canceled))
})
It("rejects a cross-origin pagination link before sending credentials", func() {
const commit = "0123456789abcdef0123456789abcdef01234567"
requestedCrossOrigin := false
evil := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
requestedCrossOrigin = true
w.WriteHeader(http.StatusNoContent)
}))
DeferCleanup(evil.Close)
origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if strings.Contains(r.URL.Path, "/revision/") {
_, _ = w.Write([]byte(`{"sha":"` + commit + `"}`))
return
}
w.Header().Set("Link", `<`+evil.URL+`/steal>; rel="next"`)
_, _ = w.Write([]byte(`[]`))
}))
DeferCleanup(origin.Close)
client := hfapi.NewClient()
client.SetBaseURL(origin.URL + "/api/models")
_, err := client.ResolveSnapshot(context.Background(), hfapi.SnapshotRequest{
Repo: "owner/repo", Revision: "main", Token: "test-token",
})
Expect(err).To(MatchError(ContainSubstring("cross-origin pagination")))
Expect(requestedCrossOrigin).To(BeFalse())
})
It("downloads a pinned public Xet fixture through the ordinary resolve URL", Label("network"), func() {
if os.Getenv("LOCALAI_HF_XET_SMOKE") != "1" {
Skip("set LOCALAI_HF_XET_SMOKE=1 to run the public Hub compatibility gate")
}
client := hfapi.NewClient()
snapshot, err := client.ResolveSnapshot(context.Background(), hfapi.SnapshotRequest{
Repo: "hf-internal-testing/tiny-random-t5-v1.1",
Revision: "main",
AllowPatterns: []string{"model.safetensors"},
})
Expect(err).NotTo(HaveOccurred())
Expect(snapshot.ResolvedRevision).To(MatchRegexp(`^[0-9a-f]{40}$`))
Expect(snapshot.Files).To(HaveLen(1))
file := snapshot.Files[0]
Expect(file.Path).To(Equal("model.safetensors"))
Expect(file.XetHash).NotTo(BeEmpty())
target := filepath.Join(GinkgoT().TempDir(), file.Path)
var final downloader.TransferProgress
err = downloader.URI(file.URL).DownloadFileWithContext(
context.Background(), target, file.LFSOID, 0, 1, nil,
downloader.WithTransferProgress(func(event downloader.TransferProgress) { final = event }),
)
Expect(err).NotTo(HaveOccurred())
info, err := os.Stat(target)
Expect(err).NotTo(HaveOccurred())
Expect(info.Size()).To(Equal(file.Size))
Expect(final.Written).To(Equal(file.Size))
})
})

View File

@@ -251,7 +251,11 @@ func (ml *ModelLoader) backendLoader(opts ...Option) (client grpc.Backend, err e
backend = realBackend
}
model, err := ml.LoadModel(o.modelID, o.model, ml.grpcModel(backend, o))
modelFileName := o.modelFile
if modelFileName == "" {
modelFileName = o.model
}
model, err := ml.LoadModelWithFile(o.modelID, o.model, modelFileName, ml.grpcModel(backend, o))
if err != nil {
// Defensive cleanup: the model usually wasn't registered yet (LoadModel
// failed before that), so StopGRPC reporting "model not found" is the

View File

@@ -397,14 +397,21 @@ func (ml *ModelLoader) ListLoadedModels() []*Model {
}
func (ml *ModelLoader) LoadModel(modelID, modelName string, loader func(string, string, string) (*Model, error)) (*Model, error) {
return ml.loadModel(modelID, modelName, loader, true)
return ml.LoadModelWithFile(modelID, modelName, modelName, loader)
}
// loadModel is the implementation behind LoadModel. checkCooldown gates fresh,
func (ml *ModelLoader) LoadModelWithFile(modelID, modelName, modelFileName string, loader func(string, string, string) (*Model, error)) (*Model, error) {
if modelFileName == "" {
modelFileName = modelName
}
return ml.loadModel(modelID, modelName, modelFileName, loader, true)
}
// loadModel is the implementation behind LoadModelWithFile. checkCooldown gates fresh,
// independent load triggers behind the per-model failure cooldown; it is set to
// false for the coalesced retry of an in-flight burst (a follower whose leader
// just failed), which is not a new trigger and should still get its one retry.
func (ml *ModelLoader) loadModel(modelID, modelName string, loader func(string, string, string) (*Model, error), checkCooldown bool) (*Model, error) {
func (ml *ModelLoader) loadModel(modelID, modelName, modelFileName string, loader func(string, string, string) (*Model, error), checkCooldown bool) (*Model, error) {
ml.mu.Lock()
distributed := ml.modelRouter != nil
ml.mu.Unlock()
@@ -430,7 +437,7 @@ func (ml *ModelLoader) loadModel(modelID, modelName string, loader func(string,
// per inference. Trade-off: cross-frontend in-flight visibility
// becomes eventually consistent, acceptable for 1-3 frontend
// deployments.
modelFile := filepath.Join(ml.ModelPath, modelName)
modelFile := filepath.Join(ml.ModelPath, modelFileName)
model, err := loader(modelID, modelName, modelFile)
if err != nil {
return nil, fmt.Errorf("failed to route model with internal loader: %s", err)
@@ -484,7 +491,7 @@ func (ml *ModelLoader) loadModel(modelID, modelName string, loader func(string,
}
// If still not loaded, the other goroutine failed. Retry once as part of
// this burst, bypassing the cooldown gate (we are not a new trigger).
return ml.loadModel(modelID, modelName, loader, false)
return ml.loadModel(modelID, modelName, modelFileName, loader, false)
}
// Mark this model as loading (create a channel that will be closed when done)
@@ -501,7 +508,7 @@ func (ml *ModelLoader) loadModel(modelID, modelName string, loader func(string,
}()
// Load the model (this can take a long time, no lock held)
modelFile := filepath.Join(ml.ModelPath, modelName)
modelFile := filepath.Join(ml.ModelPath, modelFileName)
xlog.Debug("Loading model in memory from file", "file", modelFile)
model, err := loader(modelID, modelName, modelFile)

View File

@@ -9,6 +9,7 @@ import (
type Options struct {
backendString string
model string
modelFile string
modelID string
context context.Context
@@ -73,6 +74,12 @@ func WithModel(modelFile string) Option {
}
}
func WithModelFile(modelFile string) Option {
return func(o *Options) {
o.modelFile = modelFile
}
}
func WithLoadGRPCLoadModelOpts(opts *pb.ModelOptions) Option {
return func(o *Options) {
o.gRPCOptions = opts

View File

@@ -73,6 +73,22 @@ var _ = Describe("ModelLoader", func() {
})
Context("LoadModel", func() {
It("passes a logical model and managed model file independently", func() {
const relative = ".artifacts/huggingface/0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef/snapshot"
var receivedName, receivedFile string
mockModel = model.NewModel("managed", "test.model", nil)
mockModel.MarkHealthy()
mockLoader := func(_ string, modelName, modelFile string) (*model.Model, error) {
receivedName, receivedFile = modelName, modelFile
return mockModel, nil
}
_, err := modelLoader.LoadModelWithFile("managed", "owner/repo", relative, mockLoader)
Expect(err).NotTo(HaveOccurred())
Expect(receivedName).To(Equal("owner/repo"))
Expect(receivedFile).To(Equal(filepath.Join(modelPath, filepath.FromSlash(relative))))
})
It("should load a model and keep it in memory", func() {
mockModel = model.NewModel("foo", "test.model", nil)
mockModel.MarkHealthy() // skip gRPC health check (no real server)

View File

@@ -0,0 +1,47 @@
package modelartifacts
import (
"encoding/json"
"fmt"
"os"
)
const ManifestVersion = 1
type Manifest struct {
Version int `json:"version"`
Artifact Spec `json:"artifact"`
Files []ManifestFile `json:"files"`
}
type ManifestFile struct {
Path string `json:"path"`
Size int64 `json:"size"`
SHA256 string `json:"sha256"`
BlobOID string `json:"blob_oid,omitempty"`
LFSOID string `json:"lfs_oid,omitempty"`
XetHash string `json:"xet_hash,omitempty"`
}
func ReadManifest(fileName string) (Manifest, error) {
data, err := os.ReadFile(fileName)
if err != nil {
return Manifest{}, err
}
var manifest Manifest
if err := json.Unmarshal(data, &manifest); err != nil {
return Manifest{}, err
}
if manifest.Version != ManifestVersion || manifest.Artifact.Resolved == nil {
return Manifest{}, fmt.Errorf("unsupported or incomplete artifact manifest")
}
for _, file := range manifest.Files {
if err := ValidateRelativeHubPath(file.Path); err != nil {
return Manifest{}, err
}
if file.Size < 0 || len(file.SHA256) != 64 {
return Manifest{}, fmt.Errorf("invalid manifest entry %q", file.Path)
}
}
return manifest, nil
}

View 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
}

View File

@@ -0,0 +1,153 @@
package modelartifacts_test
import (
"context"
"crypto/sha256"
"encoding/hex"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"sync"
"sync/atomic"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
hfapi "github.com/mudler/LocalAI/pkg/huggingface-api"
"github.com/mudler/LocalAI/pkg/modelartifacts"
)
type fakeSnapshotResolver struct {
mu sync.Mutex
snapshot hfapi.Snapshot
err error
calls int
}
func (f *fakeSnapshotResolver) ResolveSnapshot(context.Context, hfapi.SnapshotRequest) (hfapi.Snapshot, error) {
f.mu.Lock()
defer f.mu.Unlock()
f.calls++
return f.snapshot, f.err
}
func (f *fakeSnapshotResolver) callCount() int {
f.mu.Lock()
defer f.mu.Unlock()
return f.calls
}
var _ = Describe("controller artifact materializer", func() {
It("requires an injected token before resolving a private artifact", func() {
resolver := &fakeSnapshotResolver{}
_, err := modelartifacts.NewManager(resolver).Ensure(context.Background(), GinkgoT().TempDir(),
modelartifacts.Spec{Source: modelartifacts.Source{
Type: modelartifacts.SourceTypeHuggingFace, Repo: "owner/private", TokenEnv: modelartifacts.HuggingFaceTokenEnv,
}})
Expect(err).To(MatchError(ContainSubstring("non-empty HF_TOKEN")))
Expect(resolver.callCount()).To(BeZero())
})
It("downloads, commits, and reuses a pinned snapshot", func() {
weight := []byte("weight-bytes")
sum := sha256.Sum256(weight)
var requests atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
requests.Add(1)
Expect(r.Header.Get("Authorization")).To(Equal("Bearer hf-secret"))
w.Header().Set("Content-Length", "12")
_, _ = w.Write(weight)
}))
DeferCleanup(server.Close)
resolver := &fakeSnapshotResolver{snapshot: hfapi.Snapshot{
Endpoint: "https://huggingface.co", Repo: "owner/repo",
RequestedRevision: "main", ResolvedRevision: "0123456789abcdef0123456789abcdef01234567",
Files: []hfapi.SnapshotFile{{Path: "nested/model.safetensors", Size: 12, LFSOID: hex.EncodeToString(sum[:]), URL: server.URL + "/model"}},
}}
manager := modelartifacts.NewManager(resolver, modelartifacts.WithHuggingFaceToken("hf-secret"))
modelsPath := GinkgoT().TempDir()
spec := modelartifacts.Spec{Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo", TokenEnv: "HF_TOKEN"}}
result, err := manager.Ensure(context.Background(), modelsPath, spec)
Expect(err).NotTo(HaveOccurred())
Expect(result.CacheHit).To(BeFalse())
Expect(result.Spec.Resolved.Revision).To(Equal("0123456789abcdef0123456789abcdef01234567"))
Expect(result.RelativePath).To(HavePrefix(".artifacts/huggingface/"))
Expect(os.ReadFile(filepath.Join(modelsPath, filepath.FromSlash(result.RelativePath), "nested", "model.safetensors"))).To(Equal(weight))
manifestBytes, err := os.ReadFile(filepath.Join(modelsPath, filepath.Dir(filepath.FromSlash(result.RelativePath)), "manifest.json"))
Expect(err).NotTo(HaveOccurred())
Expect(string(manifestBytes)).NotTo(ContainSubstring("hf-secret"))
cached, err := manager.Ensure(context.Background(), modelsPath, result.Spec)
Expect(err).NotTo(HaveOccurred())
Expect(cached.CacheHit).To(BeTrue())
Expect(requests.Load()).To(Equal(int32(1)))
Expect(resolver.callCount()).To(Equal(1))
})
It("rejects a path escape before opening a destination", func() {
resolver := &fakeSnapshotResolver{snapshot: hfapi.Snapshot{
Endpoint: "https://huggingface.co", Repo: "owner/repo",
ResolvedRevision: "0123456789abcdef0123456789abcdef01234567",
Files: []hfapi.SnapshotFile{{Path: "../escape", Size: 1, URL: "https://example.invalid/file"}},
}}
modelsPath := GinkgoT().TempDir()
_, err := modelartifacts.NewManager(resolver).Ensure(context.Background(), modelsPath,
modelartifacts.Spec{Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}})
Expect(err).To(MatchError(ContainSubstring("unsafe Hub path")))
Expect(filepath.Join(modelsPath, "escape")).NotTo(BeAnExistingFile())
})
It("does not commit after cancellation", func() {
started := make(chan struct{})
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
close(started)
<-r.Context().Done()
}))
DeferCleanup(server.Close)
resolver := &fakeSnapshotResolver{snapshot: hfapi.Snapshot{
Endpoint: "https://huggingface.co", Repo: "owner/repo",
ResolvedRevision: "0123456789abcdef0123456789abcdef01234567",
Files: []hfapi.SnapshotFile{{Path: "model.bin", Size: 10, URL: server.URL}},
}}
ctx, cancel := context.WithCancel(context.Background())
done := make(chan error, 1)
modelsPath := GinkgoT().TempDir()
go func() {
_, err := modelartifacts.NewManager(resolver).Ensure(ctx, modelsPath,
modelartifacts.Spec{Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}})
done <- err
}()
<-started
cancel()
Expect(<-done).To(MatchError(context.Canceled))
entries, err := filepath.Glob(filepath.Join(modelsPath, ".artifacts", "huggingface", "*"))
Expect(err).NotTo(HaveOccurred())
Expect(entries).To(BeEmpty())
})
It("emits resolving, downloading, verifying, and committing phases", func() {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { _, _ = w.Write([]byte("x")) }))
DeferCleanup(server.Close)
resolver := &fakeSnapshotResolver{snapshot: hfapi.Snapshot{
Endpoint: "https://huggingface.co", Repo: "owner/repo",
ResolvedRevision: "0123456789abcdef0123456789abcdef01234567",
Files: []hfapi.SnapshotFile{{Path: "config.json", Size: 1, URL: server.URL}},
}}
var phases []modelartifacts.Phase
ctx := modelartifacts.WithProgressSink(context.Background(), func(event modelartifacts.ProgressEvent) {
if len(phases) == 0 || phases[len(phases)-1] != event.Phase {
phases = append(phases, event.Phase)
}
})
_, err := modelartifacts.NewManager(resolver).Ensure(ctx, GinkgoT().TempDir(),
modelartifacts.Spec{Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}})
Expect(err).NotTo(HaveOccurred())
Expect(phases).To(Equal([]modelartifacts.Phase{
modelartifacts.PhaseResolving, modelartifacts.PhaseDownloading, modelartifacts.PhaseVerifying, modelartifacts.PhaseCommitting,
}))
})
})

View File

@@ -0,0 +1,13 @@
package modelartifacts_test
import (
"testing"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
)
func TestModelArtifacts(t *testing.T) {
RegisterFailHandler(Fail)
RunSpecs(t, "Model Artifacts Suite")
}

View File

@@ -0,0 +1,86 @@
package modelartifacts
import (
"crypto/sha256"
"encoding/hex"
"encoding/json"
"fmt"
"path/filepath"
"strings"
)
type Layout struct {
Root string
Lock string
Partial string
PartialSnapshot string
Final string
Snapshot string
Manifest string
}
func CacheKey(spec Spec) (string, error) {
normalized, err := spec.Normalize()
if err != nil {
return "", err
}
if normalized.Resolved == nil {
return "", fmt.Errorf("cache key requires resolved artifact state")
}
identity := struct {
Type string `json:"type"`
Endpoint string `json:"endpoint"`
Repo string `json:"repo"`
Revision string `json:"revision"`
AllowPatterns []string `json:"allow_patterns,omitempty"`
IgnorePatterns []string `json:"ignore_patterns,omitempty"`
}{
Type: normalized.Source.Type, Endpoint: normalized.Resolved.Endpoint,
Repo: normalized.Source.Repo, Revision: normalized.Resolved.Revision,
AllowPatterns: normalized.Source.AllowPatterns, IgnorePatterns: normalized.Source.IgnorePatterns,
}
encoded, err := json.Marshal(identity)
if err != nil {
return "", err
}
sum := sha256.Sum256(encoded)
return hex.EncodeToString(sum[:]), nil
}
func RelativeSnapshotPath(cacheKey string) (string, error) {
if !cacheKeyPattern.MatchString(cacheKey) {
return "", fmt.Errorf("invalid artifact cache key")
}
return filepath.Join(".artifacts", "huggingface", cacheKey, "snapshot"), nil
}
func LayoutFor(modelsPath string, spec Spec) (Layout, error) {
if spec.Resolved == nil || !cacheKeyPattern.MatchString(spec.Resolved.CacheKey) {
return Layout{}, fmt.Errorf("layout requires a resolved cache key")
}
root := filepath.Join(modelsPath, ".artifacts")
partial := filepath.Join(root, ".partial", spec.Resolved.CacheKey)
final := filepath.Join(root, "huggingface", spec.Resolved.CacheKey)
return Layout{
Root: root,
Lock: filepath.Join(root, ".locks", spec.Resolved.CacheKey+".lock"),
Partial: partial,
PartialSnapshot: filepath.Join(partial, "snapshot"),
Final: final,
Snapshot: filepath.Join(final, "snapshot"),
Manifest: filepath.Join(final, "manifest.json"),
}, nil
}
func ValidateRelativeHubPath(candidate string) error {
if candidate == "" || filepath.IsAbs(candidate) || strings.ContainsAny(candidate, "\\\x00") {
return fmt.Errorf("unsafe Hub path %q", candidate)
}
parts := strings.Split(candidate, "/")
for _, part := range parts {
if part == "" || part == "." || part == ".." {
return fmt.Errorf("unsafe Hub path %q", candidate)
}
}
return nil
}

View File

@@ -0,0 +1,109 @@
package modelartifacts_test
import (
"context"
"encoding/json"
"os"
"path/filepath"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"github.com/mudler/LocalAI/pkg/modelartifacts"
)
var _ = Describe("artifact storage primitives", func() {
baseSpec := func() modelartifacts.Spec {
return modelartifacts.Spec{
Name: "model",
Target: "model",
Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo", Revision: "main"},
Resolved: &modelartifacts.Resolved{
Endpoint: "https://huggingface.co",
Revision: "0123456789abcdef0123456789abcdef01234567",
},
}
}
It("creates a stable path-safe cache identity", func() {
first, err := modelartifacts.CacheKey(baseSpec())
Expect(err).NotTo(HaveOccurred())
second := baseSpec()
second.Source.AllowPatterns = []string{"*.json", "*.safetensors"}
third := second
third.Source.AllowPatterns = []string{"*.safetensors", "*.json"}
secondKey, err := modelartifacts.CacheKey(second)
Expect(err).NotTo(HaveOccurred())
thirdKey, err := modelartifacts.CacheKey(third)
Expect(err).NotTo(HaveOccurred())
Expect(first).To(MatchRegexp(`^[0-9a-f]{64}$`))
Expect(secondKey).To(Equal(thirdKey))
Expect(secondKey).NotTo(Equal(first))
})
It("places all state below the models artifact root", func() {
spec := baseSpec()
key, err := modelartifacts.CacheKey(spec)
Expect(err).NotTo(HaveOccurred())
spec.Resolved.CacheKey = key
layout, err := modelartifacts.LayoutFor("/models", spec)
Expect(err).NotTo(HaveOccurred())
Expect(layout.Final).To(Equal(filepath.Join("/models", ".artifacts", "huggingface", key)))
Expect(layout.Snapshot).To(Equal(filepath.Join(layout.Final, "snapshot")))
relative, err := modelartifacts.RelativeSnapshotPath(key)
Expect(err).NotTo(HaveOccurred())
Expect(relative).To(Equal(filepath.Join(".artifacts", "huggingface", key, "snapshot")))
_, err = modelartifacts.RelativeSnapshotPath("sha256:" + key)
Expect(err).To(MatchError(ContainSubstring("invalid artifact cache key")))
})
DescribeTable("rejects hostile Hub paths",
func(candidate string) { Expect(modelartifacts.ValidateRelativeHubPath(candidate)).NotTo(Succeed()) },
Entry("absolute", "/etc/passwd"),
Entry("parent", "nested/../escape"),
Entry("backslash", `nested\escape`),
Entry("NUL", "bad\x00path"),
Entry("empty component", "nested//file"),
)
It("reads a complete versioned manifest", func() {
spec := baseSpec()
key, err := modelartifacts.CacheKey(spec)
Expect(err).NotTo(HaveOccurred())
spec.Resolved.CacheKey = key
encoded, err := json.Marshal(modelartifacts.Manifest{
Version: modelartifacts.ManifestVersion,
Artifact: spec,
Files: []modelartifacts.ManifestFile{{
Path: "nested/model.safetensors",
Size: 11,
SHA256: "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef",
}},
})
Expect(err).NotTo(HaveOccurred())
fileName := filepath.Join(GinkgoT().TempDir(), "manifest.json")
Expect(os.WriteFile(fileName, encoded, 0o600)).To(Succeed())
manifest, err := modelartifacts.ReadManifest(fileName)
Expect(err).NotTo(HaveOccurred())
Expect(manifest.Artifact.Resolved.CacheKey).To(Equal(key))
Expect(manifest.Files).To(HaveLen(1))
})
It("reports progress only through the request-scoped sink", func() {
event := modelartifacts.ProgressEvent{
Phase: modelartifacts.PhaseDownloading,
Artifact: "model",
File: "weights.safetensors",
CurrentBytes: 7,
TotalBytes: 11,
}
var received []modelartifacts.ProgressEvent
ctx := modelartifacts.WithProgressSink(context.Background(), func(update modelartifacts.ProgressEvent) {
received = append(received, update)
})
modelartifacts.ReportProgress(ctx, event)
modelartifacts.ReportProgress(context.Background(), event)
Expect(received).To(Equal([]modelartifacts.ProgressEvent{event}))
})
})

View File

@@ -0,0 +1,40 @@
package modelartifacts
import "context"
type Phase string
const (
PhaseResolving Phase = "resolving"
PhaseDownloading Phase = "downloading"
PhaseVerifying Phase = "verifying"
PhaseCommitting Phase = "committing"
PhasePersisting Phase = "persisting"
)
type ProgressEvent struct {
Phase Phase
Artifact string
File string
CurrentBytes int64
TotalBytes int64
CompletedFiles int
TotalFiles int
}
type ProgressSink func(ProgressEvent)
type progressSinkKey struct{}
func WithProgressSink(ctx context.Context, sink ProgressSink) context.Context {
if sink == nil {
return ctx
}
return context.WithValue(ctx, progressSinkKey{}, sink)
}
func ReportProgress(ctx context.Context, event ProgressEvent) {
if sink, ok := ctx.Value(progressSinkKey{}).(ProgressSink); ok && sink != nil {
sink(event)
}
}

View File

@@ -0,0 +1,82 @@
package modelartifacts
import (
"fmt"
"strings"
)
var hfReferencePrefixes = []string{
"https://huggingface.co/",
"huggingface://",
"hf://",
"hf.co/",
}
// ParsePrimaryReference converts a Hugging Face repository or file reference
// into a managed artifact spec. It accepts repo roots like "owner/repo" and
// direct file references like "huggingface://owner/repo/path/to/model.gguf".
// The boolean return is false when the reference is not Hugging Face-shaped.
func ParsePrimaryReference(raw string) (Spec, bool, error) {
source, ok, err := ParsePrimarySource(raw)
if err != nil || !ok {
return Spec{}, ok, err
}
return Spec{
Name: TargetModel,
Target: TargetModel,
Source: source,
}, true, nil
}
// ParsePrimarySource converts a Hugging Face repository or file reference into
// a managed source definition. Direct file references are translated into a
// repo plus a single allow pattern so the materializer downloads only the
// selected file.
func ParsePrimarySource(raw string) (Source, bool, error) {
ref := strings.TrimSpace(raw)
if ref == "" {
return Source{}, false, nil
}
stripped := ref
for _, prefix := range hfReferencePrefixes {
stripped = strings.TrimPrefix(stripped, prefix)
}
stripped = strings.TrimSuffix(stripped, "/")
parts := strings.Split(stripped, "/")
if len(parts) < 2 {
return Source{}, false, nil
}
hadPrefix := ref != stripped
repo, err := normalizeRepo(strings.Join(parts[:2], "/"))
if err != nil {
if ref != stripped {
return Source{}, false, fmt.Errorf("invalid Hugging Face reference %q: %w", raw, err)
}
return Source{}, false, nil
}
if !hadPrefix && len(parts) == 2 {
return Source{}, false, nil
}
source := Source{
Type: SourceTypeHuggingFace,
Repo: repo,
}
if len(parts) == 2 {
return source, true, nil
}
if parts[2] == "resolve" {
if len(parts) < 5 {
return Source{}, false, fmt.Errorf("invalid Hugging Face file reference %q", raw)
}
source.AllowPatterns = []string{strings.Join(parts[4:], "/")}
return source, true, nil
}
source.AllowPatterns = []string{strings.Join(parts[2:], "/")}
return source, true, nil
}

160
pkg/modelartifacts/types.go Normal file
View File

@@ -0,0 +1,160 @@
package modelartifacts
import (
"fmt"
"net/url"
"path"
"regexp"
"slices"
"strings"
)
const (
SourceTypeHuggingFace = "huggingface"
TargetModel = "model"
HuggingFaceTokenEnv = "HF_TOKEN"
)
var (
commitPattern = regexp.MustCompile(`^[0-9a-f]{40}$`)
cacheKeyPattern = regexp.MustCompile(`^[0-9a-f]{64}$`)
)
type Spec struct {
Name string `yaml:"name" json:"name"`
Target string `yaml:"target" json:"target"`
Source Source `yaml:"source" json:"source"`
Resolved *Resolved `yaml:"resolved,omitempty" json:"resolved,omitempty"`
}
type Source struct {
Type string `yaml:"type" json:"type"`
Repo string `yaml:"repo" json:"repo"`
Revision string `yaml:"revision,omitempty" json:"revision,omitempty"`
TokenEnv string `yaml:"token_env,omitempty" json:"token_env,omitempty"`
AllowPatterns []string `yaml:"allow_patterns,omitempty" json:"allow_patterns,omitempty"`
IgnorePatterns []string `yaml:"ignore_patterns,omitempty" json:"ignore_patterns,omitempty"`
}
type Resolved struct {
Endpoint string `yaml:"endpoint" json:"endpoint"`
Revision string `yaml:"revision" json:"revision"`
CacheKey string `yaml:"cache_key" json:"cache_key"`
}
func (s Spec) Normalize() (Spec, error) {
if strings.TrimSpace(s.Name) == "" {
s.Name = TargetModel
}
if strings.TrimSpace(s.Target) == "" {
s.Target = TargetModel
}
s.Name = strings.TrimSpace(s.Name)
s.Target = strings.TrimSpace(s.Target)
s.Source.Type = strings.TrimSpace(s.Source.Type)
s.Source.TokenEnv = strings.TrimSpace(s.Source.TokenEnv)
if s.Name != TargetModel || s.Target != TargetModel {
return Spec{}, fmt.Errorf("only artifact name/target %q is supported", TargetModel)
}
if s.Source.Type != SourceTypeHuggingFace {
return Spec{}, fmt.Errorf("unsupported artifact source type %q", s.Source.Type)
}
repo, err := normalizeRepo(s.Source.Repo)
if err != nil {
return Spec{}, err
}
s.Source.Repo = repo
if strings.TrimSpace(s.Source.Revision) == "" {
s.Source.Revision = "main"
} else {
s.Source.Revision = strings.TrimSpace(s.Source.Revision)
}
if strings.ContainsRune(s.Source.Revision, '\x00') {
return Spec{}, fmt.Errorf("revision contains NUL")
}
if s.Source.TokenEnv != "" && s.Source.TokenEnv != HuggingFaceTokenEnv {
return Spec{}, fmt.Errorf("token_env must be empty or %s", HuggingFaceTokenEnv)
}
for _, patterns := range [][]string{s.Source.AllowPatterns, s.Source.IgnorePatterns} {
for _, pattern := range patterns {
if err := validatePattern(pattern); err != nil {
return Spec{}, err
}
}
}
s.Source.AllowPatterns = slices.Clone(s.Source.AllowPatterns)
s.Source.IgnorePatterns = slices.Clone(s.Source.IgnorePatterns)
slices.Sort(s.Source.AllowPatterns)
slices.Sort(s.Source.IgnorePatterns)
if s.Resolved != nil {
resolved := *s.Resolved
resolved.Endpoint, err = normalizeEndpoint(resolved.Endpoint)
if err != nil {
return Spec{}, err
}
resolved.Revision = strings.ToLower(strings.TrimSpace(resolved.Revision))
if !commitPattern.MatchString(resolved.Revision) {
return Spec{}, fmt.Errorf("resolved revision must be 40 lowercase hexadecimal characters")
}
if resolved.CacheKey != "" && !cacheKeyPattern.MatchString(resolved.CacheKey) {
return Spec{}, fmt.Errorf("resolved cache key must be 64 lowercase hexadecimal characters")
}
s.Resolved = &resolved
}
return s, nil
}
func (s Spec) Validate() error {
normalized, err := s.Normalize()
if err != nil {
return err
}
if normalized.Resolved != nil && normalized.Resolved.CacheKey == "" {
return fmt.Errorf("resolved cache key is required in installed state")
}
return nil
}
func normalizeRepo(raw string) (string, error) {
repo := strings.TrimSpace(raw)
for _, prefix := range []string{"https://huggingface.co/", "huggingface://", "hf://"} {
repo = strings.TrimPrefix(repo, prefix)
}
repo = strings.TrimSuffix(repo, "/")
parts := strings.Split(repo, "/")
if len(parts) != 2 || parts[0] == "" || parts[1] == "" || strings.ContainsAny(repo, "?#\\\x00") {
return "", fmt.Errorf("Hugging Face repo must be exactly owner/repo")
}
return repo, nil
}
// CanonicalRepo normalizes a Hugging Face repository reference to owner/repo.
func CanonicalRepo(raw string) (string, error) {
return normalizeRepo(raw)
}
func validatePattern(pattern string) error {
if pattern == "" || path.IsAbs(pattern) || strings.ContainsAny(pattern, "\\\x00") {
return fmt.Errorf("invalid artifact pattern %q", pattern)
}
for _, part := range strings.Split(pattern, "/") {
if part == ".." {
return fmt.Errorf("invalid artifact pattern %q", pattern)
}
}
return nil
}
func normalizeEndpoint(raw string) (string, error) {
u, err := url.Parse(strings.TrimSpace(raw))
if err != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Host == "" || u.User != nil || u.RawQuery != "" || u.Fragment != "" {
return "", fmt.Errorf("invalid resolved Hugging Face endpoint")
}
u.Path = strings.TrimRight(u.Path, "/")
return u.String(), nil
}

View File

@@ -0,0 +1,70 @@
package modelartifacts_test
import (
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"github.com/mudler/LocalAI/pkg/modelartifacts"
)
var _ = Describe("artifact configuration", func() {
It("normalizes the supported primary Hugging Face source", func() {
spec, err := (modelartifacts.Spec{
Source: modelartifacts.Source{
Type: "huggingface",
Repo: "hf://Qwen/Qwen3-ASR-1.7B",
TokenEnv: "HF_TOKEN",
},
}).Normalize()
Expect(err).NotTo(HaveOccurred())
Expect(spec.Name).To(Equal("model"))
Expect(spec.Target).To(Equal("model"))
Expect(spec.Source.Repo).To(Equal("Qwen/Qwen3-ASR-1.7B"))
Expect(spec.Source.Revision).To(Equal("main"))
})
It("parses Hugging Face repo and file references into managed sources", func() {
repo, ok, err := modelartifacts.ParsePrimarySource("huggingface://Qwen/Qwen3-ASR-1.7B")
Expect(err).NotTo(HaveOccurred())
Expect(ok).To(BeTrue())
Expect(repo.Type).To(Equal(modelartifacts.SourceTypeHuggingFace))
Expect(repo.Repo).To(Equal("Qwen/Qwen3-ASR-1.7B"))
Expect(repo.AllowPatterns).To(BeEmpty())
file, ok, err := modelartifacts.ParsePrimarySource("https://huggingface.co/nomic-ai/nomic-embed-text-v1.5-GGUF/resolve/main/nomic-embed-text-v1.5.f16.gguf")
Expect(err).NotTo(HaveOccurred())
Expect(ok).To(BeTrue())
Expect(file.Repo).To(Equal("nomic-ai/nomic-embed-text-v1.5-GGUF"))
Expect(file.AllowPatterns).To(Equal([]string{"nomic-embed-text-v1.5.f16.gguf"}))
})
It("ignores non-Hugging Face references", func() {
_, ok, err := modelartifacts.ParsePrimarySource("models/local-model.gguf")
Expect(err).NotTo(HaveOccurred())
Expect(ok).To(BeFalse())
})
DescribeTable("rejects unsafe or unsupported declarations",
func(spec modelartifacts.Spec, message string) {
Expect(spec.Validate()).To(MatchError(ContainSubstring(message)))
},
Entry("unknown source", modelartifacts.Spec{Source: modelartifacts.Source{Type: "s3", Repo: "owner/repo"}}, "source type"),
Entry("secondary target", modelartifacts.Spec{Name: "controlnet", Target: "controlnet", Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}}, "target"),
Entry("malformed repo", modelartifacts.Spec{Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo/file"}}, "owner/repo"),
Entry("unrelated secret", modelartifacts.Spec{Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo", TokenEnv: "AWS_SECRET_ACCESS_KEY"}}, "HF_TOKEN"),
Entry("parent filter", modelartifacts.Spec{Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo", AllowPatterns: []string{"../*.json"}}}, "pattern"),
Entry("prefixed cache key", modelartifacts.Spec{Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}, Resolved: &modelartifacts.Resolved{Endpoint: "https://huggingface.co", Revision: "0123456789abcdef0123456789abcdef01234567", CacheKey: "sha256:bad"}}, "cache key"),
)
It("validates installed state", func() {
spec := modelartifacts.Spec{
Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"},
Resolved: &modelartifacts.Resolved{
Endpoint: "https://huggingface.co",
Revision: "0123456789abcdef0123456789abcdef01234567",
CacheKey: "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef",
},
}
Expect(spec.Validate()).To(Succeed())
})
})