Skip to content

Commit 735ca75

Browse files
mudlermudler
authored
feat(ui): allow to cancel ops (#7264)
* feat(ui): allow to cancel ops Signed-off-by: Ettore Di Giacinto <mudler@localai.io> * Improve progress text Signed-off-by: Ettore Di Giacinto <mudler@localai.io> * Cancel queued ops, don't show up message cancellation always Signed-off-by: Ettore Di Giacinto <mudler@localai.io> * fix: fixup displaying of total progress over multiple files Signed-off-by: Ettore Di Giacinto <mudler@localai.io> --------- Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
1 parent b1d1f2a commit 735ca75

29 files changed

+648
-199
lines changed

‎core/application/startup.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -62,12 +62,12 @@ func New(opts ...config.AppOption) (*Application, error) {
6262
}
6363
}
6464

65-
if err := coreStartup.InstallModels(application.GalleryService(), options.Galleries, options.BackendGalleries, options.SystemState, application.ModelLoader(), options.EnforcePredownloadScans, options.AutoloadBackendGalleries, nil, options.ModelsURL...); err != nil {
65+
if err := coreStartup.InstallModels(options.Context, application.GalleryService(), options.Galleries, options.BackendGalleries, options.SystemState, application.ModelLoader(), options.EnforcePredownloadScans, options.AutoloadBackendGalleries, nil, options.ModelsURL...); err != nil {
6666
log.Error().Err(err).Msg("error installing models")
6767
}
6868

6969
for _, backend := range options.ExternalBackends {
70-
if err := coreStartup.InstallExternalBackends(options.BackendGalleries, options.SystemState, application.ModelLoader(), nil, backend, "", ""); err != nil {
70+
if err := coreStartup.InstallExternalBackends(options.Context, options.BackendGalleries, options.SystemState, application.ModelLoader(), nil, backend, "", ""); err != nil {
7171
log.Error().Err(err).Msg("error installing external backend")
7272
}
7373
}

‎core/backend/llm.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ func ModelInference(ctx context.Context, s string, messages schema.Messages, ima
4545
if !slices.Contains(modelNames, c.Name) {
4646
utils.ResetDownloadTimers()
4747
// if we failed to load the model, we try to download it
48-
err := gallery.InstallModelFromGallery(o.Galleries, o.BackendGalleries, o.SystemState, loader, c.Name, gallery.GalleryModel{}, utils.DisplayDownloadFunction, o.EnforcePredownloadScans, o.AutoloadBackendGalleries)
48+
err := gallery.InstallModelFromGallery(ctx, o.Galleries, o.BackendGalleries, o.SystemState, loader, c.Name, gallery.GalleryModel{}, utils.DisplayDownloadFunction, o.EnforcePredownloadScans, o.AutoloadBackendGalleries)
4949
if err != nil {
5050
log.Error().Err(err).Msgf("failed to install model %q from gallery", modelFile)
5151
//return nil, err

‎core/cli/backends.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package cli
22

33
import (
4+
"context"
45
"encoding/json"
56
"fmt"
67

@@ -102,7 +103,7 @@ func (bi *BackendsInstall) Run(ctx *cliContext.Context) error {
102103
}
103104

104105
modelLoader := model.NewModelLoader(systemState, true)
105-
err = startup.InstallExternalBackends(galleries, systemState, modelLoader, progressCallback, bi.BackendArgs, bi.Name, bi.Alias)
106+
err = startup.InstallExternalBackends(context.Background(), galleries, systemState, modelLoader, progressCallback, bi.BackendArgs, bi.Name, bi.Alias)
106107
if err != nil {
107108
return err
108109
}

‎core/cli/models.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -135,7 +135,7 @@ func (mi *ModelsInstall) Run(ctx *cliContext.Context) error {
135135
}
136136

137137
modelLoader := model.NewModelLoader(systemState, true)
138-
err = startup.InstallModels(galleryService, galleries, backendGalleries, systemState, modelLoader, !mi.DisablePredownloadScan, mi.AutoloadBackendGalleries, progressCallback, modelName)
138+
err = startup.InstallModels(context.Background(), galleryService, galleries, backendGalleries, systemState, modelLoader, !mi.DisablePredownloadScan, mi.AutoloadBackendGalleries, progressCallback, modelName)
139139
if err != nil {
140140
return err
141141
}

‎core/cli/worker/worker_llamacpp.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package worker
22

33
import (
4+
"context"
45
"encoding/json"
56
"errors"
67
"fmt"
@@ -42,7 +43,7 @@ func findLLamaCPPBackend(galleries string, systemState *system.SystemState) (str
4243
log.Error().Err(err).Msg("failed loading galleries")
4344
return "", err
4445
}
45-
err := gallery.InstallBackendFromGallery(gals, systemState, ml, llamaCPPGalleryName, nil, true)
46+
err := gallery.InstallBackendFromGallery(context.Background(), gals, systemState, ml, llamaCPPGalleryName, nil, true)
4647
if err != nil {
4748
log.Error().Err(err).Msg("llama-cpp backend not found, failed to install it")
4849
return "", err

‎core/gallery/backends.go‎

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
package gallery
44

55
import (
6+
"context"
67
"encoding/json"
78
"errors"
89
"fmt"
@@ -69,7 +70,7 @@ func writeBackendMetadata(backendPath string, metadata *BackendMetadata) error {
6970
}
7071

7172
// InstallBackendFromGallery installs a backend from the gallery.
72-
func InstallBackendFromGallery(galleries []config.Gallery, systemState *system.SystemState, modelLoader *model.ModelLoader, name string, downloadStatus func(string, string, string, float64), force bool) error {
73+
func InstallBackendFromGallery(ctx context.Context, galleries []config.Gallery, systemState *system.SystemState, modelLoader *model.ModelLoader, name string, downloadStatus func(string, string, string, float64), force bool) error {
7374
if !force {
7475
// check if we already have the backend installed
7576
backends, err := ListSystemBackends(systemState)
@@ -109,7 +110,7 @@ func InstallBackendFromGallery(galleries []config.Gallery, systemState *system.S
109110
log.Debug().Str("name", name).Str("bestBackend", bestBackend.Name).Msg("Installing backend from meta backend")
110111

111112
// Then, let's install the best backend
112-
if err := InstallBackend(systemState, modelLoader, bestBackend, downloadStatus); err != nil {
113+
if err := InstallBackend(ctx, systemState, modelLoader, bestBackend, downloadStatus); err != nil {
113114
return err
114115
}
115116

@@ -134,10 +135,10 @@ func InstallBackendFromGallery(galleries []config.Gallery, systemState *system.S
134135
return nil
135136
}
136137

137-
return InstallBackend(systemState, modelLoader, backend, downloadStatus)
138+
return InstallBackend(ctx, systemState, modelLoader, backend, downloadStatus)
138139
}
139140

140-
func InstallBackend(systemState *system.SystemState, modelLoader *model.ModelLoader, config *GalleryBackend, downloadStatus func(string, string, string, float64)) error {
141+
func InstallBackend(ctx context.Context, systemState *system.SystemState, modelLoader *model.ModelLoader, config *GalleryBackend, downloadStatus func(string, string, string, float64)) error {
141142
// Create base path if it doesn't exist
142143
err := os.MkdirAll(systemState.Backend.BackendsPath, 0750)
143144
if err != nil {
@@ -164,11 +165,17 @@ func InstallBackend(systemState *system.SystemState, modelLoader *model.ModelLoa
164165
}
165166
} else {
166167
uri := downloader.URI(config.URI)
167-
if err := uri.DownloadFile(backendPath, "", 1, 1, downloadStatus); err != nil {
168+
if err := uri.DownloadFileWithContext(ctx, backendPath, "", 1, 1, downloadStatus); err != nil {
168169
success := false
169170
// Try to download from mirrors
170171
for _, mirror := range config.Mirrors {
171-
if err := downloader.URI(mirror).DownloadFile(backendPath, "", 1, 1, downloadStatus); err == nil {
172+
// Check for cancellation before trying next mirror
173+
select {
174+
case <-ctx.Done():
175+
return ctx.Err()
176+
default:
177+
}
178+
if err := downloader.URI(mirror).DownloadFileWithContext(ctx, backendPath, "", 1, 1, downloadStatus); err == nil {
172179
success = true
173180
break
174181
}

‎core/gallery/backends_test.go‎

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package gallery
22

33
import (
4+
"context"
45
"encoding/json"
56
"os"
67
"path/filepath"
@@ -55,7 +56,7 @@ var _ = Describe("Runtime capability-based backend selection", func() {
5556
)
5657
must(err)
5758
sysDefault.GPUVendor = "" // force default selection
58-
backs, err := ListSystemBackends(sysDefault)
59+
backs, err := ListSystemBackends(sysDefault)
5960
must(err)
6061
aliasBack, ok := backs.Get("llama-cpp")
6162
Expect(ok).To(BeTrue())
@@ -77,7 +78,7 @@ var _ = Describe("Runtime capability-based backend selection", func() {
7778
must(err)
7879
sysNvidia.GPUVendor = "nvidia"
7980
sysNvidia.VRAM = 8 * 1024 * 1024 * 1024
80-
backs, err = ListSystemBackends(sysNvidia)
81+
backs, err = ListSystemBackends(sysNvidia)
8182
must(err)
8283
aliasBack, ok = backs.Get("llama-cpp")
8384
Expect(ok).To(BeTrue())
@@ -116,13 +117,13 @@ var _ = Describe("Gallery Backends", func() {
116117

117118
Describe("InstallBackendFromGallery", func() {
118119
It("should return error when backend is not found", func() {
119-
err := InstallBackendFromGallery(galleries, systemState, ml, "non-existent", nil, true)
120+
err := InstallBackendFromGallery(context.TODO(), galleries, systemState, ml, "non-existent", nil, true)
120121
Expect(err).To(HaveOccurred())
121122
Expect(err.Error()).To(ContainSubstring("no backend found with name \"non-existent\""))
122123
})
123124

124125
It("should install backend from gallery", func() {
125-
err := InstallBackendFromGallery(galleries, systemState, ml, "test-backend", nil, true)
126+
err := InstallBackendFromGallery(context.TODO(), galleries, systemState, ml, "test-backend", nil, true)
126127
Expect(err).ToNot(HaveOccurred())
127128
Expect(filepath.Join(tempDir, "test-backend", "run.sh")).To(BeARegularFile())
128129
})
@@ -298,7 +299,7 @@ var _ = Describe("Gallery Backends", func() {
298299
VRAM: 1000000000000,
299300
Backend: system.Backend{BackendsPath: tempDir},
300301
}
301-
err = InstallBackendFromGallery([]config.Gallery{gallery}, nvidiaSystemState, ml, "meta-backend", nil, true)
302+
err = InstallBackendFromGallery(context.TODO(), []config.Gallery{gallery}, nvidiaSystemState, ml, "meta-backend", nil, true)
302303
Expect(err).NotTo(HaveOccurred())
303304

304305
metaBackendPath := filepath.Join(tempDir, "meta-backend")
@@ -378,7 +379,7 @@ var _ = Describe("Gallery Backends", func() {
378379
VRAM: 1000000000000,
379380
Backend: system.Backend{BackendsPath: tempDir},
380381
}
381-
err = InstallBackendFromGallery([]config.Gallery{gallery}, nvidiaSystemState, ml, "meta-backend", nil, true)
382+
err = InstallBackendFromGallery(context.TODO(), []config.Gallery{gallery}, nvidiaSystemState, ml, "meta-backend", nil, true)
382383
Expect(err).NotTo(HaveOccurred())
383384

384385
metaBackendPath := filepath.Join(tempDir, "meta-backend")
@@ -462,7 +463,7 @@ var _ = Describe("Gallery Backends", func() {
462463
VRAM: 1000000000000,
463464
Backend: system.Backend{BackendsPath: tempDir},
464465
}
465-
err = InstallBackendFromGallery([]config.Gallery{gallery}, nvidiaSystemState, ml, "meta-backend", nil, true)
466+
err = InstallBackendFromGallery(context.TODO(), []config.Gallery{gallery}, nvidiaSystemState, ml, "meta-backend", nil, true)
466467
Expect(err).NotTo(HaveOccurred())
467468

468469
metaBackendPath := filepath.Join(tempDir, "meta-backend")
@@ -561,7 +562,7 @@ var _ = Describe("Gallery Backends", func() {
561562
system.WithBackendPath(newPath),
562563
)
563564
Expect(err).NotTo(HaveOccurred())
564-
err = InstallBackend(systemState, ml, &backend, nil)
565+
err = InstallBackend(context.TODO(), systemState, ml, &backend, nil)
565566
Expect(err).To(HaveOccurred()) // Will fail due to invalid URI, but path should be created
566567
Expect(newPath).To(BeADirectory())
567568
})
@@ -593,7 +594,7 @@ var _ = Describe("Gallery Backends", func() {
593594
system.WithBackendPath(tempDir),
594595
)
595596
Expect(err).NotTo(HaveOccurred())
596-
err = InstallBackend(systemState, ml, &backend, nil)
597+
err = InstallBackend(context.TODO(), systemState, ml, &backend, nil)
597598
Expect(err).ToNot(HaveOccurred())
598599
Expect(filepath.Join(tempDir, "test-backend", "metadata.json")).To(BeARegularFile())
599600
dat, err := os.ReadFile(filepath.Join(tempDir, "test-backend", "metadata.json"))
@@ -626,7 +627,7 @@ var _ = Describe("Gallery Backends", func() {
626627

627628
Expect(filepath.Join(tempDir, "test-backend", "metadata.json")).ToNot(BeARegularFile())
628629

629-
err = InstallBackend(systemState, ml, &backend, nil)
630+
err = InstallBackend(context.TODO(), systemState, ml, &backend, nil)
630631
Expect(err).ToNot(HaveOccurred())
631632
Expect(filepath.Join(tempDir, "test-backend", "metadata.json")).To(BeARegularFile())
632633
})
@@ -647,7 +648,7 @@ var _ = Describe("Gallery Backends", func() {
647648
system.WithBackendPath(tempDir),
648649
)
649650
Expect(err).NotTo(HaveOccurred())
650-
err = InstallBackend(systemState, ml, &backend, nil)
651+
err = InstallBackend(context.TODO(), systemState, ml, &backend, nil)
651652
Expect(err).ToNot(HaveOccurred())
652653
Expect(filepath.Join(tempDir, "test-backend", "metadata.json")).To(BeARegularFile())
653654

‎core/gallery/gallery.go‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package gallery
22

33
import (
4+
"context"
45
"fmt"
56
"os"
67
"path/filepath"
@@ -28,6 +29,19 @@ func GetGalleryConfigFromURL[T any](url string, basePath string) (T, error) {
2829
return config, nil
2930
}
3031

32+
func GetGalleryConfigFromURLWithContext[T any](ctx context.Context, url string, basePath string) (T, error) {
33+
var config T
34+
uri := downloader.URI(url)
35+
err := uri.DownloadWithAuthorizationAndCallback(ctx, basePath, "", func(url string, d []byte) error {
36+
return yaml.Unmarshal(d, &config)
37+
})
38+
if err != nil {
39+
log.Error().Err(err).Str("url", url).Msg("failed to get gallery config for url")
40+
return config, err
41+
}
42+
return config, nil
43+
}
44+
3145
func ReadConfigFile[T any](filePath string) (*T, error) {
3246
// Read the YAML file
3347
yamlFile, err := os.ReadFile(filePath)

‎core/gallery/models.go‎

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package gallery
22

33
import (
4+
"context"
45
"errors"
56
"fmt"
67
"os"
@@ -72,6 +73,7 @@ type PromptTemplate struct {
7273

7374
// Installs a model from the gallery
7475
func InstallModelFromGallery(
76+
ctx context.Context,
7577
modelGalleries, backendGalleries []config.Gallery,
7678
systemState *system.SystemState,
7779
modelLoader *model.ModelLoader,
@@ -84,7 +86,7 @@ func InstallModelFromGallery(
8486

8587
if len(model.URL) > 0 {
8688
var err error
87-
config, err = GetGalleryConfigFromURL[ModelConfig](model.URL, systemState.Model.ModelsPath)
89+
config, err = GetGalleryConfigFromURLWithContext[ModelConfig](ctx, model.URL, systemState.Model.ModelsPath)
8890
if err != nil {
8991
return err
9092
}
@@ -125,15 +127,15 @@ func InstallModelFromGallery(
125127
return err
126128
}
127129

128-
installedModel, err := InstallModel(systemState, installName, &config, model.Overrides, downloadStatus, enforceScan)
130+
installedModel, err := InstallModel(ctx, systemState, installName, &config, model.Overrides, downloadStatus, enforceScan)
129131
if err != nil {
130132
return err
131133
}
132134
log.Debug().Msgf("Installed model %q", installedModel.Name)
133135
if automaticallyInstallBackend && installedModel.Backend != "" {
134136
log.Debug().Msgf("Installing backend %q", installedModel.Backend)
135137

136-
if err := InstallBackendFromGallery(backendGalleries, systemState, modelLoader, installedModel.Backend, downloadStatus, false); err != nil {
138+
if err := InstallBackendFromGallery(ctx, backendGalleries, systemState, modelLoader, installedModel.Backend, downloadStatus, false); err != nil {
137139
return err
138140
}
139141
}
@@ -154,7 +156,7 @@ func InstallModelFromGallery(
154156
return applyModel(model)
155157
}
156158

157-
func InstallModel(systemState *system.SystemState, nameOverride string, config *ModelConfig, configOverrides map[string]interface{}, downloadStatus func(string, string, string, float64), enforceScan bool) (*lconfig.ModelConfig, error) {
159+
func InstallModel(ctx context.Context, systemState *system.SystemState, nameOverride string, config *ModelConfig, configOverrides map[string]interface{}, downloadStatus func(string, string, string, float64), enforceScan bool) (*lconfig.ModelConfig, error) {
158160
basePath := systemState.Model.ModelsPath
159161
// Create base path if it doesn't exist
160162
err := os.MkdirAll(basePath, 0750)
@@ -168,6 +170,13 @@ func InstallModel(systemState *system.SystemState, nameOverride string, config *
168170

169171
// Download files and verify their SHA
170172
for i, file := range config.Files {
173+
// Check for cancellation before each file
174+
select {
175+
case <-ctx.Done():
176+
return nil, ctx.Err()
177+
default:
178+
}
179+
171180
log.Debug().Msgf("Checking %q exists and matches SHA", file.Filename)
172181

173182
if err := utils.VerifyPath(file.Filename, basePath); err != nil {
@@ -185,7 +194,7 @@ func InstallModel(systemState *system.SystemState, nameOverride string, config *
185194
}
186195
}
187196
uri := downloader.URI(file.URI)
188-
if err := uri.DownloadFile(filePath, file.SHA256, i, len(config.Files), downloadStatus); err != nil {
197+
if err := uri.DownloadFileWithContext(ctx, filePath, file.SHA256, i, len(config.Files), downloadStatus); err != nil {
189198
return nil, err
190199
}
191200
}

‎core/gallery/models_test.go‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package gallery_test
22

33
import (
4+
"context"
45
"errors"
56
"os"
67
"path/filepath"
@@ -34,7 +35,7 @@ var _ = Describe("Model test", func() {
3435
system.WithModelPath(tempdir),
3536
)
3637
Expect(err).ToNot(HaveOccurred())
37-
_, err = InstallModel(systemState, "", c, map[string]interface{}{}, func(string, string, string, float64) {}, true)
38+
_, err = InstallModel(context.TODO(), systemState, "", c, map[string]interface{}{}, func(string, string, string, float64) {}, true)
3839
Expect(err).ToNot(HaveOccurred())
3940

4041
for _, f := range []string{"cerebras", "cerebras-completion.tmpl", "cerebras-chat.tmpl", "cerebras.yaml"} {
@@ -88,7 +89,7 @@ var _ = Describe("Model test", func() {
8889
Expect(models[0].URL).To(Equal(bertEmbeddingsURL))
8990
Expect(models[0].Installed).To(BeFalse())
9091

91-
err = InstallModelFromGallery(galleries, []config.Gallery{}, systemState, nil, "test@bert", GalleryModel{}, func(s1, s2, s3 string, f float64) {}, true, true)
92+
err = InstallModelFromGallery(context.TODO(), galleries, []config.Gallery{}, systemState, nil, "test@bert", GalleryModel{}, func(s1, s2, s3 string, f float64) {}, true, true)
9293
Expect(err).ToNot(HaveOccurred())
9394

9495
dat, err := os.ReadFile(filepath.Join(tempdir, "bert.yaml"))
@@ -129,7 +130,7 @@ var _ = Describe("Model test", func() {
129130
system.WithModelPath(tempdir),
130131
)
131132
Expect(err).ToNot(HaveOccurred())
132-
_, err = InstallModel(systemState, "foo", c, map[string]interface{}{}, func(string, string, string, float64) {}, true)
133+
_, err = InstallModel(context.TODO(), systemState, "foo", c, map[string]interface{}{}, func(string, string, string, float64) {}, true)
133134
Expect(err).ToNot(HaveOccurred())
134135

135136
for _, f := range []string{"cerebras", "cerebras-completion.tmpl", "cerebras-chat.tmpl", "foo.yaml"} {
@@ -149,7 +150,7 @@ var _ = Describe("Model test", func() {
149150
system.WithModelPath(tempdir),
150151
)
151152
Expect(err).ToNot(HaveOccurred())
152-
_, err = InstallModel(systemState, "foo", c, map[string]interface{}{"backend": "foo"}, func(string, string, string, float64) {}, true)
153+
_, err = InstallModel(context.TODO(), systemState, "foo", c, map[string]interface{}{"backend": "foo"}, func(string, string, string, float64) {}, true)
153154
Expect(err).ToNot(HaveOccurred())
154155

155156
for _, f := range []string{"cerebras", "cerebras-completion.tmpl", "cerebras-chat.tmpl", "foo.yaml"} {
@@ -179,7 +180,7 @@ var _ = Describe("Model test", func() {
179180
system.WithModelPath(tempdir),
180181
)
181182
Expect(err).ToNot(HaveOccurred())
182-
_, err = InstallModel(systemState, "../../../foo", c, map[string]interface{}{}, func(string, string, string, float64) {}, true)
183+
_, err = InstallModel(context.TODO(), systemState, "../../../foo", c, map[string]interface{}{}, func(string, string, string, float64) {}, true)
183184
Expect(err).To(HaveOccurred())
184185
})
185186
})

0 commit comments

Comments
 (0)