diff --git a/core/cli/models.go b/core/cli/models.go index e2a832b0c..11f63e6ae 100644 --- a/core/cli/models.go +++ b/core/cli/models.go @@ -21,13 +21,14 @@ import ( ) type ModelsCMDFlags struct { - Galleries string `env:"LOCALAI_GALLERIES,GALLERIES" help:"JSON list of galleries" group:"models" default:"${galleries}"` - BackendGalleries string `env:"LOCALAI_BACKEND_GALLERIES,BACKEND_GALLERIES" help:"JSON list of backend galleries" group:"backends" default:"${backends}"` - ModelsPath string `env:"LOCALAI_MODELS_PATH,MODELS_PATH" type:"path" default:"${basepath}/models" help:"Path containing models used for inferencing" group:"storage"` - BackendsPath string `env:"LOCALAI_BACKENDS_PATH,BACKENDS_PATH" type:"path" default:"${basepath}/backends" help:"Path containing backends used for inferencing" group:"storage"` - Color string `env:"COLOR" hidden:""` - NoColor string `env:"NO_COLOR" hidden:""` - HFToken string `env:"HF_TOKEN" hidden:""` + Galleries string `env:"LOCALAI_GALLERIES,GALLERIES" help:"JSON list of galleries" group:"models" default:"${galleries}"` + BackendGalleries string `env:"LOCALAI_BACKEND_GALLERIES,BACKEND_GALLERIES" help:"JSON list of backend galleries" group:"backends" default:"${backends}"` + ModelsPath string `env:"LOCALAI_MODELS_PATH,MODELS_PATH" type:"path" default:"${basepath}/models" help:"Path containing models used for inferencing" group:"storage"` + BackendsPath string `env:"LOCALAI_BACKENDS_PATH,BACKENDS_PATH" type:"path" default:"${basepath}/backends" help:"Path containing backends used for inferencing" group:"storage"` + Color string `env:"COLOR" hidden:""` + NoColor string `env:"NO_COLOR" hidden:""` + HFToken string `env:"HF_TOKEN" hidden:""` + ArtifactDownloadConcurrency int `env:"LOCALAI_ARTIFACT_DOWNLOAD_CONCURRENCY" name:"artifact-download-concurrency" default:"1" help:"Maximum number of model artifact files downloaded concurrently" group:"models"` } type ModelsList struct { @@ -87,6 +88,7 @@ func (mi *ModelsInstall) Run(ctx *cliContext.Context) error { artifactMaterializer := modelartifacts.NewDefaultManager( modelartifacts.WithHuggingFaceToken(mi.HFToken), + modelartifacts.WithDownloadConcurrency(mi.ArtifactDownloadConcurrency), ) galleryService := galleryop.NewGalleryService(&config.ApplicationConfig{ SystemState: systemState, diff --git a/core/cli/run.go b/core/cli/run.go index 749dd16f0..b63db0e20 100644 --- a/core/cli/run.go +++ b/core/cli/run.go @@ -30,10 +30,11 @@ import ( // and document the deprecation in the help text. type RunCMD struct { - ModelArgs []string `arg:"" optional:"" name:"models" help:"Model configuration URLs to load"` - Color string `env:"COLOR" hidden:""` - NoColor string `env:"NO_COLOR" hidden:""` - HFToken string `env:"HF_TOKEN" hidden:""` + ModelArgs []string `arg:"" optional:"" name:"models" help:"Model configuration URLs to load"` + Color string `env:"COLOR" hidden:""` + NoColor string `env:"NO_COLOR" hidden:""` + HFToken string `env:"HF_TOKEN" hidden:""` + ArtifactDownloadConcurrency int `env:"LOCALAI_ARTIFACT_DOWNLOAD_CONCURRENCY" name:"artifact-download-concurrency" default:"1" help:"Maximum number of model artifact files downloaded concurrently" group:"models"` ExternalBackends []string `env:"LOCALAI_EXTERNAL_BACKENDS,EXTERNAL_BACKENDS" help:"A list of external backends to load from gallery on boot" group:"backends"` WebRTCNAT1To1IPs []string `env:"LOCALAI_WEBRTC_NAT_1TO1_IPS,WEBRTC_NAT_1TO1_IPS" help:"IPs advertised as the host ICE candidates for /v1/realtime WebRTC instead of every local interface. Set to the reachable host/LAN IP when running under Docker host networking or NAT, where pion otherwise offers unreachable bridge addresses and the connection drops after ICE consent checks fail." group:"api"` @@ -280,6 +281,7 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error { config.WithContext(context.Background()), config.WithModelArtifactMaterializer(modelartifacts.NewDefaultManager( modelartifacts.WithHuggingFaceToken(r.HFToken), + modelartifacts.WithDownloadConcurrency(r.ArtifactDownloadConcurrency), )), config.WithModelPreloadDisplay(r.Color, r.NoColor != ""), config.WithConfigFile(r.ModelsConfigFile), diff --git a/docs/content/advanced/model-configuration.md b/docs/content/advanced/model-configuration.md index a162a6225..9c2caa069 100644 --- a/docs/content/advanced/model-configuration.md +++ b/docs/content/advanced/model-configuration.md @@ -139,6 +139,13 @@ locally. `parameters.model` remains the logical repository ID. Once Configurations without `artifacts` keep the existing lazy repository-ID behavior. +Artifact files download sequentially by default. Set +`--artifact-download-concurrency` or +`LOCALAI_ARTIFACT_DOWNLOAD_CONCURRENCY` to increase the bounded concurrency. +Start conservatively with `2` or `4`; higher values increase bandwidth use, +open file descriptors, and pressure on the remote server. Values below `1` +are treated as `1`. + The initially migrated backend families are `transformers` and its aliases, `diffusers`, `qwen-asr`, `fish-speech`, `nemo`, `voxcpm`, `qwen-tts`, `liquid-audio`, `vllm`, `vllm-omni`, and `sglang`. Automatic imports add diff --git a/pkg/downloader/download_plan.go b/pkg/downloader/download_plan.go index bc007ed00..13341b38c 100644 --- a/pkg/downloader/download_plan.go +++ b/pkg/downloader/download_plan.go @@ -4,6 +4,7 @@ import ( "context" "github.com/mudler/xlog" + "golang.org/x/sync/errgroup" ) // FileTask describes one download operation and an optional post-download @@ -19,27 +20,34 @@ type FileTask struct { Options []DownloadOption } -// DownloadFilesWithContext executes a set of file downloads sequentially. +// DownloadFilesWithContext executes a set of file downloads with bounded concurrency. // 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 { + limit := applyDownloadOptions(opts).fileConcurrency + if limit < 1 { + limit = 1 + } + group, groupCtx := errgroup.WithContext(ctx) + group.SetLimit(limit) 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 := downloadTaskWithRetry(ctx, task, status, taskOpts); err != nil { - return err - } - if task.AfterDownload != nil { - if err := task.AfterDownload(task.Destination); err != nil { + group.Go(func() error { + if err := groupCtx.Err(); err != nil { return err } - } + taskOpts := append([]DownloadOption{}, opts...) + taskOpts = append(taskOpts, task.Options...) + if err := downloadTaskWithRetry(groupCtx, task, status, taskOpts); err != nil { + return err + } + if task.AfterDownload != nil { + return task.AfterDownload(task.Destination) + } + return nil + }) } - return nil + return group.Wait() } // downloadTaskWithRetry fetches one file, retrying transient failures. Without diff --git a/pkg/downloader/plan_test.go b/pkg/downloader/plan_test.go index 9483664d2..acc50b363 100644 --- a/pkg/downloader/plan_test.go +++ b/pkg/downloader/plan_test.go @@ -6,6 +6,8 @@ import ( "net/http/httptest" "os" "path/filepath" + "sync/atomic" + "time" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" @@ -41,4 +43,67 @@ var _ = Describe("DownloadFilesWithContext", func() { Expect(err).NotTo(HaveOccurred()) Expect(hookCalled).To(BeTrue()) }) + + It("reaches but never exceeds the configured concurrency limit", func() { + var active atomic.Int32 + var maximum atomic.Int32 + started := make(chan struct{}, 5) + release := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + current := active.Add(1) + defer active.Add(-1) + for { + previous := maximum.Load() + if current <= previous || maximum.CompareAndSwap(previous, current) { + break + } + } + started <- struct{}{} + <-release + _, _ = w.Write([]byte(r.URL.Path)) + })) + DeferCleanup(server.Close) + + tasks := make([]downloader.FileTask, 5) + for i := range tasks { + tasks[i] = downloader.FileTask{ + URI: downloader.URI(server.URL + "/file"), Destination: filepath.Join(GinkgoT().TempDir(), "file"), + FileIndex: i, TotalFiles: len(tasks), + } + } + done := make(chan error, 1) + go func() { + done <- downloader.DownloadFilesWithContext(context.Background(), tasks, nil, downloader.WithFileConcurrency(2)) + }() + Eventually(started).Should(Receive()) + Eventually(started).Should(Receive()) + Consistently(active.Load, 100*time.Millisecond).Should(Equal(int32(2))) + close(release) + Expect(<-done).NotTo(HaveOccurred()) + Expect(maximum.Load()).To(Equal(int32(2))) + }) + + It("cancels a blocked sibling and returns the permanent task error", func() { + blockedStarted := make(chan struct{}) + blockedCanceled := make(chan struct{}) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/blocked" { + close(blockedStarted) + <-r.Context().Done() + close(blockedCanceled) + return + } + <-blockedStarted + w.WriteHeader(http.StatusNotFound) + })) + DeferCleanup(server.Close) + + err := downloader.DownloadFilesWithContext(context.Background(), []downloader.FileTask{ + {URI: downloader.URI(server.URL + "/blocked"), Destination: filepath.Join(GinkgoT().TempDir(), "blocked"), TotalFiles: 2}, + {URI: downloader.URI(server.URL + "/missing"), Destination: filepath.Join(GinkgoT().TempDir(), "missing"), FileIndex: 1, TotalFiles: 2}, + }, nil, downloader.WithFileConcurrency(2)) + + Expect(err).To(MatchError(ContainSubstring("404"))) + Eventually(blockedCanceled).Should(BeClosed()) + }) }) diff --git a/pkg/downloader/uri.go b/pkg/downloader/uri.go index 1c985720c..644959202 100644 --- a/pkg/downloader/uri.go +++ b/pkg/downloader/uri.go @@ -69,6 +69,7 @@ type downloadOptions struct { verifier ImageVerifier bearerToken string transferProgress TransferProgressSink + fileConcurrency int } // DownloadOption configures DownloadFileWithContext / DownloadFile. @@ -96,6 +97,17 @@ func WithTransferProgress(sink TransferProgressSink) DownloadOption { return func(o *downloadOptions) { o.transferProgress = sink } } +// WithFileConcurrency bounds concurrent work when an option is passed to +// DownloadFilesWithContext. Individual file downloads safely ignore it. +func WithFileConcurrency(limit int) DownloadOption { + return func(o *downloadOptions) { + if limit < 1 { + limit = 1 + } + o.fileConcurrency = limit + } +} + func applyDownloadOptions(opts []DownloadOption) downloadOptions { var o downloadOptions for _, fn := range opts { diff --git a/pkg/modelartifacts/materializer.go b/pkg/modelartifacts/materializer.go index 8f1f03192..77cb64e76 100644 --- a/pkg/modelartifacts/materializer.go +++ b/pkg/modelartifacts/materializer.go @@ -13,6 +13,7 @@ import ( "path" "path/filepath" "strings" + "sync" "syscall" "time" @@ -56,10 +57,11 @@ const ( ) type Manager struct { - resolver SnapshotResolver - huggingFaceToken string - newLocker func(string) Locker - lockWait time.Duration + resolver SnapshotResolver + huggingFaceToken string + newLocker func(string) Locker + lockWait time.Duration + downloadConcurrency int // writerID names this manager's staging trees. It is drawn once, at // construction, and deliberately never persisted: a partial tree belongs to // the process run that created it, and outliving that run is precisely what @@ -100,12 +102,23 @@ func WithLockWait(wait time.Duration) ManagerOption { } } +// WithDownloadConcurrency bounds the number of artifact files downloaded at once. +func WithDownloadConcurrency(limit int) ManagerOption { + return func(manager *Manager) { + if limit < 1 { + limit = 1 + } + manager.downloadConcurrency = limit + } +} + func NewManager(resolver SnapshotResolver, options ...ManagerOption) *Manager { manager := &Manager{ - resolver: resolver, - newLocker: func(path string) Locker { return flock.New(path) }, - lockWait: DefaultLockWait, - writerID: newWriterID(), + resolver: resolver, + newLocker: func(path string) Locker { return flock.New(path) }, + lockWait: DefaultLockWait, + downloadConcurrency: 1, + writerID: newWriterID(), } for _, option := range options { option(manager) @@ -377,6 +390,9 @@ func (m *Manager) materializeLocked(ctx context.Context, modelsPath string, spec // corrupt tree look valid. manifest := Manifest{Version: ManifestVersion, Artifact: spec, Files: make([]ManifestFile, len(snapshot.Files))} completedBytes := int64(0) + completedFiles := 0 + writtenByFile := make([]int64, len(snapshot.Files)) + var bookkeeping sync.Mutex skippedFiles := 0 skippedBytes := int64(0) tasks := make([]downloader.FileTask, 0, len(snapshot.Files)) @@ -399,6 +415,7 @@ func (m *Manager) materializeLocked(ctx context.Context, modelsPath string, spec if entry, ok := reuseMaterializedFile(snapshotAbs, file); ok { manifest.Files[taskIndex] = entry completedBytes += file.Size + completedFiles++ skippedFiles++ skippedBytes += file.Size continue @@ -415,27 +432,44 @@ func (m *Manager) materializeLocked(ctx context.Context, modelsPath string, spec Options: []downloader.DownloadOption{ downloader.WithBearerToken(token), downloader.WithTransferProgress(func(event downloader.TransferProgress) { + bookkeeping.Lock() + defer bookkeeping.Unlock() + if event.Written > writtenByFile[taskIndex] { + writtenByFile[taskIndex] = event.Written + } + currentBytes := completedBytes + for _, written := range writtenByFile { + currentBytes += written + } + currentBytes = min(currentBytes, totalBytes) ReportProgress(ctx, ProgressEvent{ Phase: PhaseDownloading, Artifact: spec.Name, File: file.Path, - CurrentBytes: completedBytes + event.Written, + CurrentBytes: currentBytes, TotalBytes: totalBytes, - CompletedFiles: taskIndex, + CompletedFiles: completedFiles, TotalFiles: len(snapshot.Files), }) }), }, AfterDownload: func(string) error { + bookkeeping.Lock() + currentBytes := completedBytes + for _, written := range writtenByFile { + currentBytes += written + } + currentBytes = min(currentBytes, totalBytes) ReportProgress(ctx, ProgressEvent{ Phase: PhaseVerifying, Artifact: spec.Name, File: file.Path, - CurrentBytes: completedBytes + file.Size, + CurrentBytes: currentBytes, TotalBytes: totalBytes, - CompletedFiles: taskIndex, + CompletedFiles: completedFiles, TotalFiles: len(snapshot.Files), }) + bookkeeping.Unlock() entry, err := verifyDownloadedFile(blobAbs, file) if err != nil { _ = root.Remove(blobRel) @@ -453,8 +487,12 @@ func (m *Manager) materializeLocked(ctx context.Context, modelsPath string, spec if err := root.Rename(blobRel, destination); err != nil { return err } + bookkeeping.Lock() manifest.Files[taskIndex] = entry + writtenByFile[taskIndex] = 0 completedBytes += file.Size + completedFiles++ + bookkeeping.Unlock() return nil }, } @@ -471,7 +509,7 @@ func (m *Manager) materializeLocked(ctx context.Context, modelsPath string, spec "remaining_files", len(tasks), "total_files", len(snapshot.Files)) } - if err := downloader.DownloadFilesWithContext(ctx, tasks, nil); err != nil { + if err := downloader.DownloadFilesWithContext(ctx, tasks, nil, downloader.WithFileConcurrency(m.downloadConcurrency)); err != nil { return Result{}, err } if err := root.RemoveAll(".downloads"); err != nil { diff --git a/pkg/modelartifacts/materializer_test.go b/pkg/modelartifacts/materializer_test.go index b891960ae..3090821b0 100644 --- a/pkg/modelartifacts/materializer_test.go +++ b/pkg/modelartifacts/materializer_test.go @@ -1,6 +1,7 @@ package modelartifacts_test import ( + "bytes" "context" "crypto/sha256" "encoding/hex" @@ -12,6 +13,7 @@ import ( "strings" "sync" "sync/atomic" + "time" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" @@ -112,6 +114,77 @@ var _ = Describe("controller artifact materializer", func() { Expect(result.Spec.Resolved.PrimaryFile).To(BeEmpty()) }) + It("materializes a snapshot concurrently with ordered manifests and monotonic aggregate progress", func() { + contents := map[string][]byte{ + "slow.bin": bytes.Repeat([]byte("s"), 64*1024), + "fast.bin": bytes.Repeat([]byte("f"), 8*1024), + "mid.bin": bytes.Repeat([]byte("m"), 24*1024), + } + order := []string{"slow.bin", "fast.bin", "mid.bin"} + var active atomic.Int32 + var maximum atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + current := active.Add(1) + defer active.Add(-1) + for { + previous := maximum.Load() + if current <= previous || maximum.CompareAndSwap(previous, current) { + break + } + } + name := strings.TrimPrefix(r.URL.Path, "/") + body := contents[name] + w.Header().Set("Content-Length", strconv.Itoa(len(body))) + for offset := 0; offset < len(body); offset += 4096 { + end := min(offset+4096, len(body)) + _, _ = w.Write(body[offset:end]) + w.(http.Flusher).Flush() + time.Sleep(time.Millisecond) + } + })) + DeferCleanup(server.Close) + + files := make([]hfapi.SnapshotFile, 0, len(order)) + var total int64 + for _, name := range order { + sum := sha256.Sum256(contents[name]) + files = append(files, hfapi.SnapshotFile{Path: name, Size: int64(len(contents[name])), LFSOID: hex.EncodeToString(sum[:]), URL: server.URL + "/" + name}) + total += int64(len(contents[name])) + } + resolver := &fakeSnapshotResolver{snapshot: hfapi.Snapshot{ + Endpoint: "https://huggingface.co", Repo: "owner/repo", + ResolvedRevision: "0123456789abcdef0123456789abcdef01234567", Files: files, + }} + var progressMu sync.Mutex + var progress []int64 + ctx := modelartifacts.WithProgressSink(context.Background(), func(event modelartifacts.ProgressEvent) { + if event.Phase == modelartifacts.PhaseDownloading || event.Phase == modelartifacts.PhaseVerifying { + progressMu.Lock() + progress = append(progress, event.CurrentBytes) + progressMu.Unlock() + } + }) + result, err := modelartifacts.NewManager(resolver, modelartifacts.WithDownloadConcurrency(2)).Ensure(ctx, GinkgoT().TempDir(), + modelartifacts.Spec{Source: modelartifacts.Source{Type: "huggingface", Repo: "owner/repo"}}) + Expect(err).NotTo(HaveOccurred()) + Expect(maximum.Load()).To(Equal(int32(2))) + Expect(result.Manifest.Files).To(HaveLen(len(order))) + for index, entry := range result.Manifest.Files { + Expect(entry.Path).To(Equal(order[index])) + sum := sha256.Sum256(contents[entry.Path]) + Expect(entry.SHA256).To(Equal(hex.EncodeToString(sum[:]))) + } + progressMu.Lock() + defer progressMu.Unlock() + Expect(progress).NotTo(BeEmpty()) + for index, current := range progress { + Expect(current).To(BeNumerically("<=", total)) + if index > 0 { + Expect(current).To(BeNumerically(">=", progress[index-1])) + } + } + }) + It("rejects a path escape before opening a destination", func() { resolver := &fakeSnapshotResolver{snapshot: hfapi.Snapshot{ Endpoint: "https://huggingface.co", Repo: "owner/repo",