fix(gallery): coalesce Hugging Face artifact progress (#11117)

* fix(gallery): coalesce artifact download progress Buffer high-frequency downloading events and forward only the latest event on a periodic tick. Flush progress synchronously at phase boundaries and shutdown to preserve ordering and final state. Assisted-by: Codex:gpt-5 * fix(gallery): wire progress coalescing into model installs Route artifact progress through the 250 ms coalescer and flush it on every model operation exit. Keep the legacy download callback unchanged. Assisted-by: Codex:gpt-5 --------- Co-authored-by: Ettore Di Giacinto <mudler@localai.io>
2026-07-31 02:18:50 -04:00 · 2026-07-25 08:39:07 +02:00
parent 2d889e61a6
commit 0d82efde2b
3 changed files with 323 additions and 1 deletions
--- a/core/services/galleryop/artifact_progress_coalescer.go
+++ b/core/services/galleryop/artifact_progress_coalescer.go
@@ -0,0 +1,129 @@
+package galleryop
+
+import (
+	"sync"
+	"time"
+
+	"github.com/mudler/LocalAI/pkg/modelartifacts"
+)
+
+type artifactProgressTicker interface {
+	Chan() <-chan time.Time
+	Stop()
+}
+
+type realArtifactProgressTicker struct {
+	ticker *time.Ticker
+}
+
+func (t *realArtifactProgressTicker) Chan() <-chan time.Time { return t.ticker.C }
+
+func (t *realArtifactProgressTicker) Stop() { t.ticker.Stop() }
+
+func newRealArtifactProgressTicker(interval time.Duration) artifactProgressTicker {
+	return &realArtifactProgressTicker{ticker: time.NewTicker(interval)}
+}
+
+var newArtifactProgressTicker = newRealArtifactProgressTicker
+
+type artifactProgressCoalescer struct {
+	mu        sync.Mutex
+	forwardMu sync.Mutex
+	closed    bool
+	pending   *modelartifacts.ProgressEvent
+	ticker    artifactProgressTicker
+	done      chan struct{}
+	forward   modelartifacts.ProgressSink
+}
+
+func newArtifactProgressCoalescer(interval time.Duration, forward modelartifacts.ProgressSink) *artifactProgressCoalescer {
+	c := &artifactProgressCoalescer{
+		ticker:  newArtifactProgressTicker(interval),
+		done:    make(chan struct{}),
+		forward: forward,
+	}
+	go c.run()
+	return c
+}
+
+func (c *artifactProgressCoalescer) Sink(event modelartifacts.ProgressEvent) {
+	if event.Phase == modelartifacts.PhaseDownloading {
+		c.mu.Lock()
+		defer c.mu.Unlock()
+		if c.closed {
+			return
+		}
+		c.pending = &event
+		return
+	}
+
+	c.forwardMu.Lock()
+	defer c.forwardMu.Unlock()
+
+	c.mu.Lock()
+	if c.closed {
+		c.mu.Unlock()
+		return
+	}
+	pending := c.takePendingLocked()
+	c.mu.Unlock()
+
+	c.forwardEvent(pending)
+	c.forwardEvent(&event)
+}
+
+func (c *artifactProgressCoalescer) Close() {
+	c.forwardMu.Lock()
+	defer c.forwardMu.Unlock()
+
+	c.mu.Lock()
+	if c.closed {
+		c.mu.Unlock()
+		return
+	}
+	c.closed = true
+	pending := c.takePendingLocked()
+	close(c.done)
+	c.ticker.Stop()
+	c.mu.Unlock()
+
+	c.forwardEvent(pending)
+}
+
+func (c *artifactProgressCoalescer) run() {
+	for {
+		select {
+		case <-c.ticker.Chan():
+			c.flush()
+		case <-c.done:
+			return
+		}
+	}
+}
+
+func (c *artifactProgressCoalescer) flush() {
+	c.forwardMu.Lock()
+	defer c.forwardMu.Unlock()
+
+	c.mu.Lock()
+	if c.closed {
+		c.mu.Unlock()
+		return
+	}
+	pending := c.takePendingLocked()
+	c.mu.Unlock()
+
+	c.forwardEvent(pending)
+}
+
+func (c *artifactProgressCoalescer) takePendingLocked() *modelartifacts.ProgressEvent {
+	pending := c.pending
+	c.pending = nil
+	return pending
+}
+
+func (c *artifactProgressCoalescer) forwardEvent(event *modelartifacts.ProgressEvent) {
+	if event != nil && c.forward != nil {
+		c.forward(*event)
+	}
+}
--- a/core/services/galleryop/artifact_progress_coalescer_test.go
+++ b/core/services/galleryop/artifact_progress_coalescer_test.go
@@ -0,0 +1,190 @@
+package galleryop
+
+import (
+	"context"
+	"errors"
+	"sync"
+	"time"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+
+	"github.com/mudler/LocalAI/core/config"
+	"github.com/mudler/LocalAI/core/gallery"
+	"github.com/mudler/LocalAI/core/services/messaging"
+	"github.com/mudler/LocalAI/pkg/modelartifacts"
+)
+
+type manualArtifactProgressTicker struct {
+	channel chan time.Time
+	stopped bool
+}
+
+func newManualArtifactProgressTicker() *manualArtifactProgressTicker {
+	return &manualArtifactProgressTicker{channel: make(chan time.Time, 1)}
+}
+
+func (t *manualArtifactProgressTicker) Chan() <-chan time.Time { return t.channel }
+
+func (t *manualArtifactProgressTicker) Stop() { t.stopped = true }
+
+type artifactProgressRecorder struct {
+	mu     sync.Mutex
+	events []modelartifacts.ProgressEvent
+}
+
+func (r *artifactProgressRecorder) Sink(event modelartifacts.ProgressEvent) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.events = append(r.events, event)
+}
+
+func (r *artifactProgressRecorder) Events() []modelartifacts.ProgressEvent {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	return append([]modelartifacts.ProgressEvent(nil), r.events...)
+}
+
+type modelOperationProgressManager struct {
+	err error
+}
+
+func (m *modelOperationProgressManager) InstallModel(ctx context.Context, _ *ManagementOp[gallery.GalleryModel, gallery.ModelConfig], _ ProgressCallback) error {
+	modelartifacts.ReportProgress(ctx, modelartifacts.ProgressEvent{
+		Phase: modelartifacts.PhaseDownloading, File: "model.bin", CurrentBytes: 32, TotalBytes: 64,
+	})
+	modelartifacts.ReportProgress(ctx, modelartifacts.ProgressEvent{
+		Phase: modelartifacts.PhaseDownloading, File: "model.bin", CurrentBytes: 64, TotalBytes: 64,
+	})
+	return m.err
+}
+
+func (m *modelOperationProgressManager) DeleteModel(string) error { return nil }
+
+type recordingProgressClient struct {
+	mu      sync.Mutex
+	updates []*OpStatus
+}
+
+func (c *recordingProgressClient) Publish(_ string, data any) error {
+	event, ok := data.(GalleryProgressEvent)
+	if !ok || event.Status == nil || event.Status.Phase != string(modelartifacts.PhaseDownloading) {
+		return nil
+	}
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.updates = append(c.updates, event.Status)
+	return nil
+}
+
+func (c *recordingProgressClient) Subscribe(string, func([]byte)) (messaging.Subscription, error) {
+	return nil, nil
+}
+
+func (c *recordingProgressClient) QueueSubscribe(string, string, func([]byte)) (messaging.Subscription, error) {
+	return nil, nil
+}
+
+func (c *recordingProgressClient) QueueSubscribeReply(string, string, func([]byte, func([]byte))) (messaging.Subscription, error) {
+	return nil, nil
+}
+
+func (c *recordingProgressClient) SubscribeReply(string, func([]byte, func([]byte))) (messaging.Subscription, error) {
+	return nil, nil
+}
+
+func (c *recordingProgressClient) Request(string, []byte, time.Duration) ([]byte, error) {
+	return nil, nil
+}
+
+func (c *recordingProgressClient) IsConnected() bool { return true }
+
+func (c *recordingProgressClient) Close() {}
+
+func (c *recordingProgressClient) Updates() []*OpStatus {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	return append([]*OpStatus(nil), c.updates...)
+}
+
+var _ = Describe("artifact progress coalescer", func() {
+	var (
+		ticker   *manualArtifactProgressTicker
+		recorder *artifactProgressRecorder
+	)
+
+	BeforeEach(func() {
+		ticker = newManualArtifactProgressTicker()
+		recorder = &artifactProgressRecorder{}
+		newArtifactProgressTicker = func(time.Duration) artifactProgressTicker {
+			return ticker
+		}
+	})
+
+	AfterEach(func() {
+		newArtifactProgressTicker = newRealArtifactProgressTicker
+	})
+
+	It("coalesces downloading events and forwards the latest event on a tick", func() {
+		coalescer := newArtifactProgressCoalescer(250*time.Millisecond, recorder.Sink)
+		DeferCleanup(coalescer.Close)
+
+		coalescer.Sink(modelartifacts.ProgressEvent{Phase: modelartifacts.PhaseDownloading, CurrentBytes: 32})
+		coalescer.Sink(modelartifacts.ProgressEvent{Phase: modelartifacts.PhaseDownloading, CurrentBytes: 64})
+		Expect(recorder.Events()).To(BeEmpty())
+
+		ticker.channel <- time.Now()
+		Eventually(recorder.Events).Should(Equal([]modelartifacts.ProgressEvent{{Phase: modelartifacts.PhaseDownloading, CurrentBytes: 64}}))
+	})
+
+	It("flushes pending downloading progress before a phase event", func() {
+		coalescer := newArtifactProgressCoalescer(250*time.Millisecond, recorder.Sink)
+		DeferCleanup(coalescer.Close)
+
+		download := modelartifacts.ProgressEvent{Phase: modelartifacts.PhaseDownloading, CurrentBytes: 64}
+		verifying := modelartifacts.ProgressEvent{Phase: modelartifacts.PhaseVerifying, CurrentBytes: 64}
+		coalescer.Sink(download)
+		coalescer.Sink(verifying)
+
+		Expect(recorder.Events()).To(Equal([]modelartifacts.ProgressEvent{download, verifying}))
+	})
+
+	It("flushes pending progress and disables forwarding when closed", func() {
+		coalescer := newArtifactProgressCoalescer(250*time.Millisecond, recorder.Sink)
+		download := modelartifacts.ProgressEvent{Phase: modelartifacts.PhaseDownloading, CurrentBytes: 64}
+		coalescer.Sink(download)
+
+		coalescer.Close()
+		Expect(recorder.Events()).To(Equal([]modelartifacts.ProgressEvent{download}))
+		Expect(ticker.stopped).To(BeTrue())
+
+		coalescer.Sink(modelartifacts.ProgressEvent{Phase: modelartifacts.PhaseCommitting})
+		ticker.channel <- time.Now()
+		Consistently(recorder.Events).Should(Equal([]modelartifacts.ProgressEvent{download}))
+	})
+
+	It("coalesces model operation progress through the existing bridge", func() {
+		installErr := errors.New("stop after progress")
+		progressClient := &recordingProgressClient{}
+		service := NewGalleryService(&config.ApplicationConfig{}, nil)
+		service.modelManager = &modelOperationProgressManager{err: installErr}
+		service.natsClient = progressClient
+		op := &ManagementOp[gallery.GalleryModel, gallery.ModelConfig]{
+			ID:                 "model-operation",
+			GalleryElementName: "model",
+			Context:            context.Background(),
+		}
+
+		Expect(service.modelHandler(op, nil, nil)).To(MatchError(installErr))
+		Expect(progressClient.Updates()).To(ConsistOf(&OpStatus{
+			Phase:              string(modelartifacts.PhaseDownloading),
+			Message:            "Downloading model file: model.bin",
+			FileName:           "model.bin",
+			Progress:           90,
+			CurrentBytes:       64,
+			TotalBytes:         64,
+			GalleryElementName: "model",
+			Cancellable:        true,
+		}))
+	})
+})
--- a/core/services/galleryop/models.go
+++ b/core/services/galleryop/models.go
@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"os"
 	"slices"
+	"time"

 	"github.com/mudler/LocalAI/core/config"
 	"github.com/mudler/LocalAI/core/gallery"
@@ -58,11 +59,13 @@ func (g *GalleryService) modelHandler(op *ManagementOp[gallery.GalleryModel, gal
 		status.GalleryElementName = op.GalleryElementName
 		g.UpdateStatus(op.ID, status)
 	})
+	coalescer := newArtifactProgressCoalescer(250*time.Millisecond, bridge.Sink)
+	defer coalescer.Close()
 	operationCtx := op.Context
 	if operationCtx == nil {
 		operationCtx = context.Background()
 	}
-	operationCtx = modelartifacts.WithProgressSink(operationCtx, bridge.Sink)
+	operationCtx = modelartifacts.WithProgressSink(operationCtx, coalescer.Sink)

 	// displayDownload displays the download progress
 	progressCallback := func(fileName string, current string, total string, percentage float64) {