mirror of
https://github.com/mudler/LocalAI.git
synced 2026-07-30 09:57:57 -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:
173
pkg/downloader/auth_progress_test.go
Normal file
173
pkg/downloader/auth_progress_test.go
Normal 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())
|
||||
})
|
||||
})
|
||||
39
pkg/downloader/download_plan.go
Normal file
39
pkg/downloader/download_plan.go
Normal 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
|
||||
}
|
||||
44
pkg/downloader/plan_test.go
Normal file
44
pkg/downloader/plan_test.go
Normal 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())
|
||||
})
|
||||
})
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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++ {
|
||||
|
||||
239
pkg/huggingface-api/snapshot.go
Normal file
239
pkg/huggingface-api/snapshot.go
Normal 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
|
||||
}
|
||||
149
pkg/huggingface-api/snapshot_test.go
Normal file
149
pkg/huggingface-api/snapshot_test.go
Normal 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))
|
||||
})
|
||||
})
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
47
pkg/modelartifacts/manifest.go
Normal file
47
pkg/modelartifacts/manifest.go
Normal 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
|
||||
}
|
||||
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
|
||||
}
|
||||
153
pkg/modelartifacts/materializer_test.go
Normal file
153
pkg/modelartifacts/materializer_test.go
Normal 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,
|
||||
}))
|
||||
})
|
||||
})
|
||||
13
pkg/modelartifacts/modelartifacts_suite_test.go
Normal file
13
pkg/modelartifacts/modelartifacts_suite_test.go
Normal 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")
|
||||
}
|
||||
86
pkg/modelartifacts/path.go
Normal file
86
pkg/modelartifacts/path.go
Normal 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
|
||||
}
|
||||
109
pkg/modelartifacts/path_test.go
Normal file
109
pkg/modelartifacts/path_test.go
Normal 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}))
|
||||
})
|
||||
})
|
||||
40
pkg/modelartifacts/progress.go
Normal file
40
pkg/modelartifacts/progress.go
Normal 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)
|
||||
}
|
||||
}
|
||||
82
pkg/modelartifacts/reference.go
Normal file
82
pkg/modelartifacts/reference.go
Normal 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
160
pkg/modelartifacts/types.go
Normal 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
|
||||
}
|
||||
70
pkg/modelartifacts/types_test.go
Normal file
70
pkg/modelartifacts/types_test.go
Normal 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())
|
||||
})
|
||||
})
|
||||
Reference in New Issue
Block a user