Skip to content

Commit ebbcba3

Browse files
sozercansozercan
authored
fix: runtime capability detection for backends (#6149)
* runtime capability detection for backends Signed-off-by: Sertac Ozercan <sozercan@gmail.com> * test Signed-off-by: Sertac Ozercan <sozercan@gmail.com> * skip nvidia on darwin Signed-off-by: Sertac Ozercan <sozercan@gmail.com> * address review comments Signed-off-by: Sertac Ozercan <sozercan@gmail.com> * fix apple test Signed-off-by: Sertac Ozercan <sozercan@gmail.com> * remove unused func Signed-off-by: Sertac Ozercan <sozercan@gmail.com> --------- Signed-off-by: Sertac Ozercan <sozercan@gmail.com>
1 parent 0de7551 commit ebbcba3

File tree

3 files changed

+189
-54
lines changed

3 files changed

+189
-54
lines changed

‎core/gallery/backends.go

Lines changed: 98 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,13 @@
1+
// Package gallery provides installation and registration utilities for LocalAI backends,
2+
// including meta-backend resolution based on system capabilities.
13
package gallery
24

35
import (
46
"encoding/json"
57
"fmt"
68
"os"
79
"path/filepath"
10+
"strings"
811
"time"
912

1013
"github.com/mudler/LocalAI/core/config"
@@ -20,6 +23,12 @@ const (
2023
runFile = "run.sh"
2124
)
2225

26+
// backendCandidate represents an installed concrete backend option for a given alias
27+
type backendCandidate struct {
28+
name string
29+
runFile string
30+
}
31+
2332
// readBackendMetadata reads the metadata JSON file for a backend
2433
func readBackendMetadata(backendPath string) (*BackendMetadata, error) {
2534
metadataPath := filepath.Join(backendPath, metadataFile)
@@ -58,7 +67,7 @@ func writeBackendMetadata(backendPath string, metadata *BackendMetadata) error {
5867
return nil
5968
}
6069

61-
// Installs a model from the gallery
70+
// InstallBackendFromGallery installs a backend from the gallery.
6271
func InstallBackendFromGallery(galleries []config.Gallery, systemState *system.SystemState, modelLoader *model.ModelLoader, name string, downloadStatus func(string, string, string, float64), force bool) error {
6372
if !force {
6473
// check if we already have the backend installed
@@ -282,23 +291,18 @@ func (b SystemBackends) GetAll() []SystemBackend {
282291
}
283292

284293
func ListSystemBackends(systemState *system.SystemState) (SystemBackends, error) {
285-
potentialBackends, err := os.ReadDir(systemState.Backend.BackendsPath)
286-
if err != nil {
287-
return nil, err
288-
}
289-
294+
// Gather backends from system and user paths, then resolve alias conflicts by capability.
290295
backends := make(SystemBackends)
291296

292-
systemBackends, err := os.ReadDir(systemState.Backend.BackendsSystemPath)
293-
if err == nil {
294-
// system backends are special, they are provided by the system and not managed by LocalAI
297+
// System-provided backends
298+
if systemBackends, err := os.ReadDir(systemState.Backend.BackendsSystemPath); err == nil {
295299
for _, systemBackend := range systemBackends {
296300
if systemBackend.IsDir() {
297-
systemBackendRunFile := filepath.Join(systemState.Backend.BackendsSystemPath, systemBackend.Name(), runFile)
298-
if _, err := os.Stat(systemBackendRunFile); err == nil {
301+
run := filepath.Join(systemState.Backend.BackendsSystemPath, systemBackend.Name(), runFile)
302+
if _, err := os.Stat(run); err == nil {
299303
backends[systemBackend.Name()] = SystemBackend{
300304
Name: systemBackend.Name(),
301-
RunFile: filepath.Join(systemState.Backend.BackendsSystemPath, systemBackend.Name(), runFile),
305+
RunFile: run,
302306
IsMeta: false,
303307
IsSystem: true,
304308
Metadata: nil,
@@ -307,64 +311,104 @@ func ListSystemBackends(systemState *system.SystemState) (SystemBackends, error)
307311
}
308312
}
309313
} else {
310-
log.Warn().Err(err).Msg("Failed to read system backends, but that's ok, we will just use the backends managed by LocalAI")
314+
log.Warn().Err(err).Msg("Failed to read system backends, proceeding with user-managed backends")
311315
}
312316

313-
for _, potentialBackend := range potentialBackends {
314-
if potentialBackend.IsDir() {
315-
potentialBackendRunFile := filepath.Join(systemState.Backend.BackendsPath, potentialBackend.Name(), runFile)
317+
// User-managed backends and alias collection
318+
entries, err := os.ReadDir(systemState.Backend.BackendsPath)
319+
if err != nil {
320+
return nil, err
321+
}
316322

317-
var metadata *BackendMetadata
323+
aliasGroups := make(map[string][]backendCandidate)
324+
metaMap := make(map[string]*BackendMetadata)
318325

319-
// If metadata file does not exist, we just use the directory name
320-
// and we do not fill the other metadata (such as potential backend Aliases)
321-
metadataFilePath := filepath.Join(systemState.Backend.BackendsPath, potentialBackend.Name(), metadataFile)
322-
if _, err := os.Stat(metadataFilePath); os.IsNotExist(err) {
323-
metadata = &BackendMetadata{
324-
Name: potentialBackend.Name(),
325-
}
326+
for _, e := range entries {
327+
if !e.IsDir() {
328+
continue
329+
}
330+
dir := e.Name()
331+
run := filepath.Join(systemState.Backend.BackendsPath, dir, runFile)
332+
333+
var metadata *BackendMetadata
334+
metadataPath := filepath.Join(systemState.Backend.BackendsPath, dir, metadataFile)
335+
if _, err := os.Stat(metadataPath); os.IsNotExist(err) {
336+
metadata = &BackendMetadata{Name: dir}
337+
} else {
338+
m, rerr := readBackendMetadata(filepath.Join(systemState.Backend.BackendsPath, dir))
339+
if rerr != nil {
340+
return nil, rerr
341+
}
342+
if m == nil {
343+
metadata = &BackendMetadata{Name: dir}
326344
} else {
327-
// Check for alias in metadata
328-
metadata, err = readBackendMetadata(filepath.Join(systemState.Backend.BackendsPath, potentialBackend.Name()))
329-
if err != nil {
330-
return nil, err
331-
}
345+
metadata = m
332346
}
347+
}
333348

334-
if !backends.Exists(potentialBackend.Name()) {
335-
// We don't want to override aliases if already set, and if we are meta backend
336-
if _, err := os.Stat(potentialBackendRunFile); err == nil {
337-
backends[potentialBackend.Name()] = SystemBackend{
338-
Name: potentialBackend.Name(),
339-
RunFile: potentialBackendRunFile,
340-
IsMeta: false,
341-
Metadata: metadata,
342-
}
343-
}
349+
metaMap[dir] = metadata
350+
351+
// Concrete backend entry
352+
if _, err := os.Stat(run); err == nil {
353+
backends[dir] = SystemBackend{
354+
Name: dir,
355+
RunFile: run,
356+
IsMeta: false,
357+
Metadata: metadata,
344358
}
359+
}
345360

346-
if metadata == nil {
347-
continue
361+
// Alias candidates
362+
if metadata.Alias != "" {
363+
aliasGroups[metadata.Alias] = append(aliasGroups[metadata.Alias], backendCandidate{name: dir, runFile: run})
364+
}
365+
366+
// Meta backends indirection
367+
if metadata.MetaBackendFor != "" {
368+
backends[metadata.Name] = SystemBackend{
369+
Name: metadata.Name,
370+
RunFile: filepath.Join(systemState.Backend.BackendsPath, metadata.MetaBackendFor, runFile),
371+
IsMeta: true,
372+
Metadata: metadata,
348373
}
374+
}
375+
}
349376

350-
if metadata.Alias != "" {
351-
backends[metadata.Alias] = SystemBackend{
352-
Name: metadata.Alias,
353-
RunFile: potentialBackendRunFile,
354-
IsMeta: false,
355-
Metadata: metadata,
377+
// Resolve aliases using system capability preferences
378+
tokens := systemState.BackendPreferenceTokens()
379+
for alias, cands := range aliasGroups {
380+
chosen := backendCandidate{}
381+
// Try preference tokens
382+
for _, t := range tokens {
383+
for _, c := range cands {
384+
if strings.Contains(strings.ToLower(c.name), t) && c.runFile != "" {
385+
chosen = c
386+
break
356387
}
357388
}
358-
359-
if metadata.MetaBackendFor != "" {
360-
backends[metadata.Name] = SystemBackend{
361-
Name: metadata.Name,
362-
RunFile: filepath.Join(systemState.Backend.BackendsPath, metadata.MetaBackendFor, runFile),
363-
IsMeta: true,
364-
Metadata: metadata,
389+
if chosen.runFile != "" {
390+
break
391+
}
392+
}
393+
// Fallback: first runnable
394+
if chosen.runFile == "" {
395+
for _, c := range cands {
396+
if c.runFile != "" {
397+
chosen = c
398+
break
365399
}
366400
}
367401
}
402+
if chosen.runFile == "" {
403+
continue
404+
}
405+
md := metaMap[chosen.name]
406+
backends[alias] = SystemBackend{
407+
Name: alias,
408+
RunFile: chosen.runFile,
409+
IsMeta: false,
410+
Metadata: md,
411+
}
368412
}
369413

370414
return backends, nil

‎core/gallery/backends_test.go

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,73 @@ const (
1818
testImage = "quay.io/mudler/tests:localai-backend-test"
1919
)
2020

21+
var _ = Describe("Runtime capability-based backend selection", func() {
22+
var tempDir string
23+
24+
BeforeEach(func() {
25+
var err error
26+
tempDir, err = os.MkdirTemp("", "gallery-caps-*")
27+
Expect(err).NotTo(HaveOccurred())
28+
})
29+
30+
AfterEach(func() {
31+
os.RemoveAll(tempDir)
32+
})
33+
34+
It("ListSystemBackends prefers optimal alias candidate", func() {
35+
// Arrange two installed backends sharing the same alias
36+
must := func(err error) { Expect(err).NotTo(HaveOccurred()) }
37+
38+
cpuDir := filepath.Join(tempDir, "cpu-llama-cpp")
39+
must(os.MkdirAll(cpuDir, 0o750))
40+
cpuMeta := &BackendMetadata{Alias: "llama-cpp", Name: "cpu-llama-cpp"}
41+
b, _ := json.Marshal(cpuMeta)
42+
must(os.WriteFile(filepath.Join(cpuDir, "metadata.json"), b, 0o644))
43+
must(os.WriteFile(filepath.Join(cpuDir, "run.sh"), []byte(""), 0o755))
44+
45+
cudaDir := filepath.Join(tempDir, "cuda12-llama-cpp")
46+
must(os.MkdirAll(cudaDir, 0o750))
47+
cudaMeta := &BackendMetadata{Alias: "llama-cpp", Name: "cuda12-llama-cpp"}
48+
b, _ = json.Marshal(cudaMeta)
49+
must(os.WriteFile(filepath.Join(cudaDir, "metadata.json"), b, 0o644))
50+
must(os.WriteFile(filepath.Join(cudaDir, "run.sh"), []byte(""), 0o755))
51+
52+
// Default system: alias should point to CPU
53+
sysDefault, err := system.GetSystemState(
54+
system.WithBackendPath(tempDir),
55+
)
56+
must(err)
57+
sysDefault.GPUVendor = "" // force default selection
58+
backs, err := ListSystemBackends(sysDefault)
59+
must(err)
60+
aliasBack, ok := backs.Get("llama-cpp")
61+
Expect(ok).To(BeTrue())
62+
Expect(aliasBack.RunFile).To(Equal(filepath.Join(cpuDir, "run.sh")))
63+
// concrete entries remain
64+
_, ok = backs.Get("cpu-llama-cpp")
65+
Expect(ok).To(BeTrue())
66+
_, ok = backs.Get("cuda12-llama-cpp")
67+
Expect(ok).To(BeTrue())
68+
69+
// NVIDIA system: alias should point to CUDA
70+
// Force capability to nvidia to make the test deterministic on platforms like darwin/arm64 (which default to metal)
71+
os.Setenv("LOCALAI_FORCE_META_BACKEND_CAPABILITY", "nvidia")
72+
defer os.Unsetenv("LOCALAI_FORCE_META_BACKEND_CAPABILITY")
73+
74+
sysNvidia, err := system.GetSystemState(
75+
system.WithBackendPath(tempDir),
76+
)
77+
must(err)
78+
sysNvidia.GPUVendor = "nvidia"
79+
sysNvidia.VRAM = 8 * 1024 * 1024 * 1024
80+
backs, err = ListSystemBackends(sysNvidia)
81+
must(err)
82+
aliasBack, ok = backs.Get("llama-cpp")
83+
Expect(ok).To(BeTrue())
84+
Expect(aliasBack.RunFile).To(Equal(filepath.Join(cudaDir, "run.sh")))
85+
})
86+
})
87+
2188
var _ = Describe("Gallery Backends", func() {
2289
var (
2390
tempDir string

‎pkg/system/capabilities.go

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
// Package system provides system detection utilities, including GPU/vendor detection
2+
// and capability classification used to select optimal backends at runtime.
13
package system
24

35
import (
@@ -116,3 +118,25 @@ func detectGPUVendor(gpus []*gpu.GraphicsCard) (string, error) {
116118

117119
return "", nil
118120
}
121+
122+
// BackendPreferenceTokens returns a list of substrings that represent the preferred
123+
// backend implementation order for the current system capability. Callers can use
124+
// these tokens to select the most appropriate concrete backend among multiple
125+
// candidates sharing the same alias (e.g., "llama-cpp").
126+
func (s *SystemState) BackendPreferenceTokens() []string {
127+
capStr := strings.ToLower(s.getSystemCapabilities())
128+
switch {
129+
case strings.HasPrefix(capStr, nvidia):
130+
return []string{"cuda", "vulkan", "cpu"}
131+
case strings.HasPrefix(capStr, amd):
132+
return []string{"rocm", "hip", "vulkan", "cpu"}
133+
case strings.HasPrefix(capStr, intel):
134+
return []string{"sycl", "intel", "cpu"}
135+
case strings.HasPrefix(capStr, metal):
136+
return []string{"metal", "cpu"}
137+
case strings.HasPrefix(capStr, darwinX86):
138+
return []string{"darwin-x86", "cpu"}
139+
default:
140+
return []string{"cpu"}
141+
}
142+
}

0 commit comments

Comments
 (0)