diff --git a/core/services/galleryop/artifact_progress_coalescer.go b/core/services/galleryop/artifact_progress_coalescer.go new file mode 100644 index 000000000..5bc18e72d --- /dev/null +++ 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) + } +} diff --git a/core/services/galleryop/artifact_progress_coalescer_test.go b/core/services/galleryop/artifact_progress_coalescer_test.go new file mode 100644 index 000000000..2e8953072 --- /dev/null +++ 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, + })) + }) +}) diff --git a/core/services/galleryop/models.go b/core/services/galleryop/models.go index e6a965b28..609299385 100644 --- 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) {