feat(modelartifacts): bound Hugging Face file downloads

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,

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),

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

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

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())
+	})
 })

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 {

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 {