Skip to content

Commit 8d930f4

Browse files
mudlermudler
committed
feat(cli): allow to install backends from OCI tar files
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
1 parent 238c334 commit 8d930f4

4 files changed

Lines changed: 62 additions & 71 deletions

File tree

‎core/cli/backends.go‎

Lines changed: 3 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@ import (
88
"github.com/mudler/LocalAI/core/config"
99

1010
"github.com/mudler/LocalAI/core/gallery"
11-
"github.com/mudler/LocalAI/pkg/downloader"
1211
"github.com/mudler/LocalAI/pkg/startup"
1312
"github.com/rs/zerolog/log"
1413
"github.com/schollz/progressbar/v3"
@@ -23,12 +22,6 @@ type BackendsList struct {
2322
BackendsCMDFlags `embed:""`
2423
}
2524

26-
type BackendsInstallSingle struct {
27-
InstallArgs []string `arg:"" optional:"" name:"backend" help:"Backend images to install"`
28-
29-
BackendsCMDFlags `embed:""`
30-
}
31-
3225
type BackendsInstall struct {
3326
BackendArgs []string `arg:"" optional:"" name:"backends" help:"Backend configuration URLs to load"`
3427

@@ -42,36 +35,9 @@ type BackendsUninstall struct {
4235
}
4336

4437
type BackendsCMD struct {
45-
List BackendsList `cmd:"" help:"List the backends available in your galleries" default:"withargs"`
46-
Install BackendsInstall `cmd:"" help:"Install a backend from the gallery"`
47-
InstallSingle BackendsInstallSingle `cmd:"" help:"Install a single backend from the gallery"`
48-
Uninstall BackendsUninstall `cmd:"" help:"Uninstall a backend"`
49-
}
50-
51-
func (bi *BackendsInstallSingle) Run(ctx *cliContext.Context) error {
52-
for _, backend := range bi.InstallArgs {
53-
progressBar := progressbar.NewOptions(
54-
1000,
55-
progressbar.OptionSetDescription(fmt.Sprintf("downloading backend %s", backend)),
56-
progressbar.OptionShowBytes(false),
57-
progressbar.OptionClearOnFinish(),
58-
)
59-
progressCallback := func(fileName string, current string, total string, percentage float64) {
60-
v := int(percentage * 10)
61-
err := progressBar.Set(v)
62-
if err != nil {
63-
log.Error().Err(err).Str("filename", fileName).Int("value", v).Msg("error while updating progress bar")
64-
}
65-
}
66-
67-
if err := gallery.InstallBackend(bi.BackendsPath, &gallery.GalleryBackend{
68-
URI: backend,
69-
}, progressCallback); err != nil {
70-
return err
71-
}
72-
}
73-
74-
return nil
38+
List BackendsList `cmd:"" help:"List the backends available in your galleries" default:"withargs"`
39+
Install BackendsInstall `cmd:"" help:"Install a backend from the gallery"`
40+
Uninstall BackendsUninstall `cmd:"" help:"Uninstall a backend"`
7541
}
7642

7743
func (bl *BackendsList) Run(ctx *cliContext.Context) error {
@@ -116,23 +82,6 @@ func (bi *BackendsInstall) Run(ctx *cliContext.Context) error {
11682
}
11783
}
11884

119-
backendURI := downloader.URI(backendName)
120-
121-
if !backendURI.LooksLikeOCI() {
122-
backends, err := gallery.AvailableBackends(galleries, bi.BackendsPath)
123-
if err != nil {
124-
return err
125-
}
126-
127-
backend := gallery.FindGalleryElement(backends, backendName, bi.BackendsPath)
128-
if backend == nil {
129-
log.Error().Str("backend", backendName).Msg("backend not found")
130-
return fmt.Errorf("backend not found: %s", backendName)
131-
}
132-
133-
log.Info().Str("backend", backendName).Str("license", backend.License).Msg("installing backend")
134-
}
135-
13685
err := startup.InstallExternalBackends(galleries, bi.BackendsPath, progressCallback, backendName)
13786
if err != nil {
13887
return err

‎core/gallery/backends.go‎

Lines changed: 7 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,8 @@ import (
99

1010
"github.com/mudler/LocalAI/core/config"
1111
"github.com/mudler/LocalAI/core/system"
12+
"github.com/mudler/LocalAI/pkg/downloader"
1213
"github.com/mudler/LocalAI/pkg/model"
13-
"github.com/mudler/LocalAI/pkg/oci"
1414
"github.com/rs/zerolog/log"
1515
)
1616

@@ -151,19 +151,15 @@ func InstallBackend(basePath string, config *GalleryBackend, downloadStatus func
151151
}
152152

153153
name := config.Name
154-
155-
img, err := oci.GetImage(config.URI, "", nil, nil)
156-
if err != nil {
157-
return fmt.Errorf("failed to get image %q: %v", config.URI, err)
158-
}
159-
160154
backendPath := filepath.Join(basePath, name)
161-
if err := os.MkdirAll(backendPath, 0750); err != nil {
162-
return fmt.Errorf("failed to create backend path %q: %v", backendPath, err)
155+
err = os.MkdirAll(backendPath, 0750)
156+
if err != nil {
157+
return fmt.Errorf("failed to create base path: %v", err)
163158
}
164159

165-
if err := oci.ExtractOCIImage(img, config.URI, backendPath, downloadStatus); err != nil {
166-
return fmt.Errorf("failed to extract image %q: %v", config.URI, err)
160+
uri := downloader.URI(config.URI)
161+
if err := uri.DownloadFile(backendPath, "", 1, 1, downloadStatus); err != nil {
162+
return fmt.Errorf("failed to download backend %q: %v", config.URI, err)
167163
}
168164

169165
// Create metadata for the backend

‎pkg/downloader/uri.go‎

Lines changed: 37 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import (
1313
"strconv"
1414
"strings"
1515

16+
"github.com/google/go-containerregistry/pkg/v1/tarball"
1617
ocispec "github.com/opencontainers/image-spec/specs-go/v1"
1718

1819
"github.com/mudler/LocalAI/pkg/oci"
@@ -25,6 +26,7 @@ const (
2526
HuggingFacePrefix1 = "hf://"
2627
HuggingFacePrefix2 = "hf.co/"
2728
OCIPrefix = "oci://"
29+
OCIFilePrefix = "ocifile://"
2830
OllamaPrefix = "ollama://"
2931
HTTPPrefix = "http://"
3032
HTTPSPrefix = "https://"
@@ -137,8 +139,18 @@ func (u URI) LooksLikeURL() bool {
137139
strings.HasPrefix(string(u), GithubURI2)
138140
}
139141

142+
func (u URI) LooksLikeHTTPURL() bool {
143+
return strings.HasPrefix(string(u), HTTPPrefix) ||
144+
strings.HasPrefix(string(u), HTTPSPrefix)
145+
}
146+
140147
func (s URI) LooksLikeOCI() bool {
141-
return strings.HasPrefix(string(s), OCIPrefix) || strings.HasPrefix(string(s), OllamaPrefix)
148+
return strings.HasPrefix(string(s), "quay.io") ||
149+
strings.HasPrefix(string(s), OCIPrefix) ||
150+
strings.HasPrefix(string(s), OllamaPrefix) ||
151+
strings.HasPrefix(string(s), OCIFilePrefix) ||
152+
strings.HasPrefix(string(s), "ghcr.io") ||
153+
strings.HasPrefix(string(s), "docker.io")
142154
}
143155

144156
func (s URI) ResolveURL() string {
@@ -234,6 +246,13 @@ func (uri URI) checkSeverSupportsRangeHeader() (bool, error) {
234246
func (uri URI) DownloadFile(filePath, sha string, fileN, total int, downloadStatus func(string, string, string, float64)) error {
235247
url := uri.ResolveURL()
236248
if uri.LooksLikeOCI() {
249+
250+
// Only Ollama wants to download to the file, for the rest, we want to download to the directory
251+
// so we check if filepath has any extension, otherwise we assume it's a directory
252+
if filepath.Ext(filePath) != "" && !strings.HasPrefix(url, OllamaPrefix) {
253+
filePath = filepath.Dir(filePath)
254+
}
255+
237256
progressStatus := func(desc ocispec.Descriptor) io.Writer {
238257
return &progressWriter{
239258
fileName: filePath,
@@ -245,18 +264,32 @@ func (uri URI) DownloadFile(filePath, sha string, fileN, total int, downloadStat
245264
}
246265
}
247266

248-
if strings.HasPrefix(url, OllamaPrefix) {
249-
url = strings.TrimPrefix(url, OllamaPrefix)
267+
if url, ok := strings.CutPrefix(url, OllamaPrefix); ok {
250268
return oci.OllamaFetchModel(url, filePath, progressStatus)
251269
}
252270

271+
if url, ok := strings.CutPrefix(url, OCIFilePrefix); ok {
272+
// Open the tarball
273+
img, err := tarball.ImageFromPath(url, nil)
274+
if err != nil {
275+
return fmt.Errorf("failed to open tarball: %s", err.Error())
276+
}
277+
278+
return oci.ExtractOCIImage(img, url, filePath, downloadStatus)
279+
}
280+
253281
url = strings.TrimPrefix(url, OCIPrefix)
254282
img, err := oci.GetImage(url, "", nil, nil)
255283
if err != nil {
256284
return fmt.Errorf("failed to get image %q: %v", url, err)
257285
}
258286

259-
return oci.ExtractOCIImage(img, url, filepath.Dir(filePath), downloadStatus)
287+
return oci.ExtractOCIImage(img, url, filePath, downloadStatus)
288+
}
289+
290+
// We need to check if url looks like an URL or bail out
291+
if !uri.LooksLikeHTTPURL() {
292+
return fmt.Errorf("url %q does not look like an HTTP URL", url)
260293
}
261294

262295
// Check if the file already exists

‎pkg/startup/backend_preload.go‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,14 @@ package startup
33
import (
44
"errors"
55
"fmt"
6+
"path/filepath"
67
"strings"
78

89
"github.com/mudler/LocalAI/core/config"
910
"github.com/mudler/LocalAI/core/gallery"
1011
"github.com/mudler/LocalAI/core/system"
12+
"github.com/mudler/LocalAI/pkg/downloader"
13+
"github.com/rs/zerolog/log"
1114
)
1215

1316
func InstallExternalBackends(galleries []config.Gallery, backendPath string, downloadStatus func(string, string, string, float64), backends ...string) error {
@@ -17,11 +20,21 @@ func InstallExternalBackends(galleries []config.Gallery, backendPath string, dow
1720
return fmt.Errorf("failed to get system state: %w", err)
1821
}
1922
for _, backend := range backends {
23+
uri := downloader.URI(backend)
2024
switch {
21-
case strings.HasPrefix(backend, "oci://"):
22-
backend = strings.TrimPrefix(backend, "oci://")
25+
case uri.LooksLikeOCI():
26+
name, err := uri.FilenameFromUrl()
27+
if err != nil {
28+
return fmt.Errorf("failed to get filename from URL: %w", err)
29+
}
30+
// strip extension if any
31+
name = strings.TrimSuffix(name, filepath.Ext(name))
2332

33+
log.Info().Str("backend", backend).Str("name", name).Msg("Installing backend from OCI image")
2434
if err := gallery.InstallBackend(backendPath, &gallery.GalleryBackend{
35+
Metadata: gallery.Metadata{
36+
Name: name,
37+
},
2538
URI: backend,
2639
}, downloadStatus); err != nil {
2740
errs = errors.Join(err, fmt.Errorf("error installing backend %s", backend))

0 commit comments

Comments
 (0)