feat: bound artifact download concurrency

Assisted-by: Codex:gpt-5

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 {

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