mirror of
https://github.com/mudler/LocalAI.git
synced 2026-07-30 18:09:05 -04:00
feat(gallery): rank model variants by host backend preference
Variant auto-selection filtered candidates by whether their backend can run on the host, then ranked the survivors by size alone. The backend never influenced the choice beyond that gate, so a Mac offered both an MLX build and a llama.cpp build kept neither filtered and installed whichever was larger, leaving the native accelerated runtime unused. The same held for CUDA against CPU on NVIDIA and ROCm against Vulkan on AMD. Rank by the host's backend preference between the fit tier and size: fit stays a filter, preference decides among the builds the host can equally hold, and size still separates builds on equally preferred runtimes. The preference data stays in one declarative table in pkg/system, now read by a prefix lookup instead of a switch, so adding a capability or reordering one host's runtimes is a one-line edit and the gallery's ranking code carries no per-backend branching. MLX joins the metal rule ahead of metal itself, which is inert for the existing alias-resolution consumer because no alias group holds a candidate named for mlx. An unrecognised backend, an unrecognised capability and an absent preference list all collapse to the previous size-only ordering rather than erroring or dropping candidates. Assisted-by: Claude:claude-opus-4-8 Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
@@ -224,6 +224,28 @@ var _ = Describe("DescribeVariants", func() {
|
||||
Expect(view.AutoSelected).To(Equal("qwen3-8b-vllm-awq"))
|
||||
})
|
||||
|
||||
It("matches what installing would pick when the host prefers a runtime", func() {
|
||||
// The picker and the installer both have to apply backend
|
||||
// preference, or a Mac would be shown the GGUF build and handed the
|
||||
// MLX one. Asserting the agreement AND the name pins both halves.
|
||||
mlx := newModel("qwen3-8b-mlx-4bit", "mlx")
|
||||
models = append(models, mlx)
|
||||
base.Variants = append(base.Variants, gallery.Variant{Model: "qwen3-8b-mlx-4bit"})
|
||||
|
||||
env := probing(gib(64), map[string]uint64{
|
||||
"qwen3-8b-vllm-awq": gib(20),
|
||||
"qwen3-8b-gguf-q8": gib(9),
|
||||
"qwen3-8b-mlx-4bit": gib(5),
|
||||
})
|
||||
env.BackendPreference = []string{"mlx", "metal", "cpu"}
|
||||
|
||||
agreesWithInstall(env)
|
||||
|
||||
view, err := gallery.DescribeVariants(models, base, env)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(view.AutoSelected).To(Equal("qwen3-8b-mlx-4bit"))
|
||||
})
|
||||
|
||||
It("names the entry itself when no variant fits", func() {
|
||||
view, err := gallery.DescribeVariants(models, base, probing(gib(2), map[string]uint64{
|
||||
"qwen3-8b-vllm-awq": gib(20),
|
||||
|
||||
@@ -275,6 +275,11 @@ func HostResolveEnv(ctx context.Context, systemState *system.SystemState) Resolv
|
||||
BackendCompatible: func(backend string) bool {
|
||||
return systemState.IsBackendCompatible(backend, "")
|
||||
},
|
||||
// The ranking half of the hardware story, from the same table that
|
||||
// already picks among installed backend builds. IsBackendCompatible
|
||||
// only rules out what cannot run; on a Mac both an MLX and a llama.cpp
|
||||
// build can, and this is what makes the native runtime win.
|
||||
BackendPreference: systemState.BackendPreferenceTokens(),
|
||||
ProbeMemory: func(target *GalleryModel) uint64 {
|
||||
return probeEntryMemory(ctx, target)
|
||||
},
|
||||
|
||||
@@ -66,6 +66,23 @@ type ResolveEnv struct {
|
||||
// A nil func treats every backend as runnable, the right default for a
|
||||
// caller with no view of the hardware.
|
||||
BackendCompatible func(backend string) bool
|
||||
// BackendPreference lists the runtime tokens this host prefers, best first,
|
||||
// as SystemState.BackendPreferenceTokens reports them (e.g. metal gives
|
||||
// ["mlx", "metal", "cpu"], NVIDIA gives ["cuda", "vulkan", "cpu"]). A token
|
||||
// is matched as a substring of a variant's backend name.
|
||||
//
|
||||
// BackendCompatible answers "can this run here at all" and is a filter.
|
||||
// This answers "which of the things that CAN run here should win" and is
|
||||
// only a ranking. The two are deliberately separate: a Mac can run both an
|
||||
// MLX build and a llama.cpp build, so nothing is filtered, yet the native
|
||||
// accelerated runtime should still be installed even when the GGUF build is
|
||||
// the larger download.
|
||||
//
|
||||
// An empty list ranks every backend equally, which reduces selection to
|
||||
// ordering by size alone. That is the right default for a caller with no
|
||||
// view of the hardware, and it is what every host looked like before
|
||||
// preference existed.
|
||||
BackendPreference []string
|
||||
// ProbeMemory measures how much memory a referenced gallery entry needs,
|
||||
// without downloading it. A zero result means "could not tell", never
|
||||
// "needs nothing".
|
||||
@@ -86,6 +103,30 @@ func (e ResolveEnv) backendRuns(backend string) bool {
|
||||
return e.BackendCompatible(backend)
|
||||
}
|
||||
|
||||
// preferenceRank scores a backend against this host's preferred runtime order,
|
||||
// lower being better.
|
||||
//
|
||||
// It walks BackendPreference generically and knows nothing about any particular
|
||||
// backend or capability: everything a new runtime needs is expressed by the
|
||||
// token table in pkg/system, so adding one never reaches this function.
|
||||
//
|
||||
// A backend matching no token scores just below the least preferred known one.
|
||||
// It is never dropped, and a field where nothing matches scores uniformly, so
|
||||
// an unrecognised backend, an unrecognised capability and an absent preference
|
||||
// list all degrade to the same predictable place: ordering by size alone.
|
||||
func (e ResolveEnv) preferenceRank(backend string) int {
|
||||
if len(e.BackendPreference) == 0 {
|
||||
return 0
|
||||
}
|
||||
name := strings.ToLower(backend)
|
||||
for i, token := range e.BackendPreference {
|
||||
if token != "" && strings.Contains(name, token) {
|
||||
return i
|
||||
}
|
||||
}
|
||||
return len(e.BackendPreference)
|
||||
}
|
||||
|
||||
// VariantSelection is the outcome of a selection pass.
|
||||
type VariantSelection struct {
|
||||
Option VariantOption
|
||||
@@ -118,7 +159,9 @@ type VariantSelection struct {
|
||||
// a candidate like any other, not a last resort: an entry whose own build is
|
||||
// the largest thing that fits must win against a smaller variant, and a base
|
||||
// of known size must win against a variant whose size nothing could measure.
|
||||
// 5. The survivors are ranked and the best one wins, by rankOf below.
|
||||
// 5. The survivors are ranked and the best one wins: fit tier first (rankOf
|
||||
// below), then this host's backend preference, then size. Preference beats
|
||||
// size deliberately, so a Mac takes an MLX build over a larger GGUF one.
|
||||
// 6. With no survivor at all, which can only happen when the caller supplied no
|
||||
// base, there is nothing to install and this reports ErrNoVariantMatch.
|
||||
func SelectVariant(options []VariantOption, env ResolveEnv, pin string) (VariantSelection, error) {
|
||||
@@ -132,9 +175,10 @@ func SelectVariant(options []VariantOption, env ResolveEnv, pin string) (Variant
|
||||
}
|
||||
|
||||
type ranked struct {
|
||||
option VariantOption
|
||||
memory uint64
|
||||
rank int
|
||||
option VariantOption
|
||||
memory uint64
|
||||
rank int
|
||||
preference int
|
||||
}
|
||||
|
||||
survivors := make([]ranked, 0, len(options))
|
||||
@@ -160,7 +204,12 @@ func SelectVariant(options []VariantOption, env ResolveEnv, pin string) (Variant
|
||||
survivingVariants++
|
||||
}
|
||||
|
||||
survivors = append(survivors, ranked{option: o, memory: memory, rank: rankOf(o, env)})
|
||||
survivors = append(survivors, ranked{
|
||||
option: o,
|
||||
memory: memory,
|
||||
rank: rankOf(o, env),
|
||||
preference: env.preferenceRank(o.Backend),
|
||||
})
|
||||
}
|
||||
|
||||
if len(survivors) == 0 {
|
||||
@@ -170,12 +219,29 @@ func SelectVariant(options []VariantOption, env ResolveEnv, pin string) (Variant
|
||||
)
|
||||
}
|
||||
|
||||
// Stable so that options within one rank and of identical size keep their
|
||||
// authored order, which is the only thing order still decides.
|
||||
// Three keys, in this order, and the order is the whole design:
|
||||
//
|
||||
// 1. rank, the fit tier. Fit is a fact about whether the host can hold the
|
||||
// build at all, so nothing outranks it.
|
||||
// 2. preference, the host's runtime order. Among builds the host can
|
||||
// equally hold, the one on the more native runtime wins even when it is
|
||||
// the smaller download. A Mac given an MLX build and a larger llama.cpp
|
||||
// build should run MLX; ranking by size first would install the GGUF and
|
||||
// leave the accelerated build unused. Do NOT move this below size.
|
||||
// 3. memory, largest first, because among builds on equally preferred
|
||||
// runtimes a bigger build is a higher quality quantization of the same
|
||||
// model. Preference must not flatten this: two builds on one backend are
|
||||
// still ordered by size.
|
||||
//
|
||||
// Stable so that options equal on all three keep their authored order, which
|
||||
// is the only thing order still decides.
|
||||
sort.SliceStable(survivors, func(i, j int) bool {
|
||||
if survivors[i].rank != survivors[j].rank {
|
||||
return survivors[i].rank < survivors[j].rank
|
||||
}
|
||||
if survivors[i].preference != survivors[j].preference {
|
||||
return survivors[i].preference < survivors[j].preference
|
||||
}
|
||||
return survivors[i].memory > survivors[j].memory
|
||||
})
|
||||
|
||||
@@ -189,8 +255,8 @@ func SelectVariant(options []VariantOption, env ResolveEnv, pin string) (Variant
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Ranks, best first. Within a rank the larger footprint wins, because a bigger
|
||||
// build is a higher quality quantization of the same model.
|
||||
// Fit tiers, best first. Within a tier the host's preferred runtime wins, and
|
||||
// within one runtime the larger footprint wins; see the sort in SelectVariant.
|
||||
const (
|
||||
// rankProvenFit is a measured size that the host is measured to satisfy.
|
||||
rankProvenFit = iota
|
||||
|
||||
@@ -243,6 +243,159 @@ var _ = Describe("SelectVariant", func() {
|
||||
})
|
||||
})
|
||||
|
||||
Describe("ranking by host backend preference", func() {
|
||||
// The token lists below are exactly what
|
||||
// SystemState.BackendPreferenceTokens reports for these hosts. They are
|
||||
// spelled out rather than read from the live machine so the specs pin
|
||||
// the intended behaviour on every CI runner.
|
||||
darwinMetal := []string{"mlx", "metal", "cpu"}
|
||||
nvidia := []string{"cuda", "vulkan", "cpu"}
|
||||
|
||||
// A Mac runs both engines, so nothing is filtered here and preference is
|
||||
// the only thing that can decide.
|
||||
darwinRunsEverything := func(string) bool { return true }
|
||||
|
||||
It("prefers an MLX build to a larger llama.cpp build on darwin", func() {
|
||||
options := []gallery.VariantOption{
|
||||
option("m-mlx-4bit", "mlx", gib(8)),
|
||||
option("m-gguf-q8", "llama-cpp", gib(24)),
|
||||
base("m-gguf-q4", gib(4)),
|
||||
}
|
||||
|
||||
selection, err := gallery.SelectVariant(options, gallery.ResolveEnv{
|
||||
AvailableMemory: gib(64),
|
||||
BackendCompatible: darwinRunsEverything,
|
||||
BackendPreference: darwinMetal,
|
||||
}, "")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(selection.Option.Variant.Model).To(Equal("m-mlx-4bit"))
|
||||
})
|
||||
|
||||
It("takes the larger llama.cpp build on the same host once preference is unknown", func() {
|
||||
// The mirror of the spec above, proving the MLX win comes from the
|
||||
// preference list and not from anything intrinsic to the option set.
|
||||
options := []gallery.VariantOption{
|
||||
option("m-mlx-4bit", "mlx", gib(8)),
|
||||
option("m-gguf-q8", "llama-cpp", gib(24)),
|
||||
base("m-gguf-q4", gib(4)),
|
||||
}
|
||||
|
||||
selection, err := gallery.SelectVariant(options, gallery.ResolveEnv{
|
||||
AvailableMemory: gib(64),
|
||||
BackendCompatible: darwinRunsEverything,
|
||||
}, "")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(selection.Option.Variant.Model).To(Equal("m-gguf-q8"))
|
||||
})
|
||||
|
||||
It("prefers a CUDA build to a larger CPU-only build on nvidia", func() {
|
||||
options := []gallery.VariantOption{
|
||||
option("m-cuda-q4", "cuda12-llama-cpp", gib(8)),
|
||||
option("m-cpu-q8", "cpu-llama-cpp", gib(24)),
|
||||
base("m-base", gib(4)),
|
||||
}
|
||||
|
||||
selection, err := gallery.SelectVariant(options, gallery.ResolveEnv{
|
||||
AvailableMemory: gib(64),
|
||||
BackendPreference: nvidia,
|
||||
}, "")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(selection.Option.Variant.Model).To(Equal("m-cuda-q4"))
|
||||
})
|
||||
|
||||
It("still picks the largest fitting build among equally preferred backends", func() {
|
||||
// Preference must not flatten size ordering: with one runtime in
|
||||
// play there is nothing left for it to decide.
|
||||
options := []gallery.VariantOption{
|
||||
option("m-cuda-q4", "cuda12-llama-cpp", gib(6)),
|
||||
option("m-cuda-q8", "cuda12-llama-cpp", gib(12)),
|
||||
option("m-cuda-f16", "cuda12-llama-cpp", gib(48)),
|
||||
base("m-base", gib(2)),
|
||||
}
|
||||
|
||||
selection, err := gallery.SelectVariant(options, gallery.ResolveEnv{
|
||||
AvailableMemory: gib(16),
|
||||
BackendPreference: nvidia,
|
||||
}, "")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(selection.Option.Variant.Model).To(Equal("m-cuda-q8"))
|
||||
})
|
||||
|
||||
It("does not let a preferred backend rescue a build that does not fit", func() {
|
||||
// Fit is a filter and preference is only a ranking among survivors.
|
||||
// The CUDA build is both preferred and too large, so the CPU build
|
||||
// the host can actually hold has to win.
|
||||
options := []gallery.VariantOption{
|
||||
option("m-cuda-f16", "cuda12-llama-cpp", gib(48)),
|
||||
option("m-cpu-q4", "cpu-llama-cpp", gib(6)),
|
||||
base("m-base", gib(2)),
|
||||
}
|
||||
|
||||
selection, err := gallery.SelectVariant(options, gallery.ResolveEnv{
|
||||
AvailableMemory: gib(16),
|
||||
BackendPreference: nvidia,
|
||||
}, "")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(selection.Option.Variant.Model).To(Equal("m-cpu-q4"))
|
||||
Expect(selection.Reasons).To(ContainElement(ContainSubstring("m-cuda-f16")))
|
||||
})
|
||||
|
||||
It("orders unrecognised backends by size rather than dropping them", func() {
|
||||
// Neither engine name carries an nvidia token. Nothing is discarded
|
||||
// and nothing is arbitrarily favoured; the host falls back to the
|
||||
// size-only behaviour it had before preference existed.
|
||||
options := []gallery.VariantOption{
|
||||
option("m-vllm-awq", "vllm", gib(12)),
|
||||
option("m-sglang-fp8", "sglang", gib(6)),
|
||||
base("m-base", gib(2)),
|
||||
}
|
||||
|
||||
selection, err := gallery.SelectVariant(options, gallery.ResolveEnv{
|
||||
AvailableMemory: gib(16),
|
||||
BackendPreference: nvidia,
|
||||
}, "")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(selection.Option.Variant.Model).To(Equal("m-vllm-awq"))
|
||||
Expect(selection.Reasons).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("still installs the base when nothing fits, whatever the host prefers", func() {
|
||||
options := []gallery.VariantOption{
|
||||
option("m-mlx-8bit", "mlx", gib(48)),
|
||||
option("m-gguf-q8", "llama-cpp", gib(24)),
|
||||
base("m-base", gib(64)),
|
||||
}
|
||||
|
||||
selection, err := gallery.SelectVariant(options, gallery.ResolveEnv{
|
||||
AvailableMemory: gib(4),
|
||||
BackendCompatible: darwinRunsEverything,
|
||||
BackendPreference: darwinMetal,
|
||||
}, "")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(selection.Option.Variant.Model).To(Equal("m-base"))
|
||||
Expect(selection.Option.IsBase).To(BeTrue())
|
||||
Expect(selection.FellBackToBase).To(BeTrue())
|
||||
})
|
||||
|
||||
It("honors a pin for a less preferred backend", func() {
|
||||
// A pin is an operator override, so preference must not quietly
|
||||
// redirect it any more than the memory filter does.
|
||||
options := []gallery.VariantOption{
|
||||
option("m-mlx-4bit", "mlx", gib(8)),
|
||||
option("m-gguf-q8", "llama-cpp", gib(24)),
|
||||
base("m-base", gib(4)),
|
||||
}
|
||||
|
||||
selection, err := gallery.SelectVariant(options, gallery.ResolveEnv{
|
||||
AvailableMemory: gib(64),
|
||||
BackendCompatible: darwinRunsEverything,
|
||||
BackendPreference: darwinMetal,
|
||||
}, "m-gguf-q8")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(selection.Option.Variant.Model).To(Equal("m-gguf-q8"))
|
||||
})
|
||||
})
|
||||
|
||||
Describe("falling back to the base", func() {
|
||||
It("selects the base when nothing else fits", func() {
|
||||
options := []gallery.VariantOption{
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"slices"
|
||||
"strings"
|
||||
|
||||
"github.com/mudler/xlog"
|
||||
@@ -18,8 +19,8 @@ const (
|
||||
Intel = "intel"
|
||||
|
||||
// Private constants - only used within this package
|
||||
defaultCapability = "default"
|
||||
disableCapability = "disable"
|
||||
defaultCapability = "default"
|
||||
disableCapability = "disable"
|
||||
nvidiaL4T = "nvidia-l4t"
|
||||
darwinX86 = "darwin-x86"
|
||||
metal = "metal"
|
||||
@@ -43,8 +44,53 @@ const (
|
||||
backendTokenROCM = "rocm"
|
||||
backendTokenHIP = "hip"
|
||||
backendTokenSYCL = "sycl"
|
||||
backendTokenCPU = "cpu"
|
||||
)
|
||||
|
||||
// backendPreferenceRule maps a detected capability to the runtime tokens that
|
||||
// capability prefers, best first.
|
||||
//
|
||||
// This table is the single source of truth for "what should this host run
|
||||
// first", consumed both by concrete-backend alias resolution and by gallery
|
||||
// variant auto-selection. Keeping it declarative is deliberate: LocalAI gains
|
||||
// backends and hardware targets continuously, and expressing a new preference
|
||||
// has to stay a one-line edit here rather than a change to any ranking code.
|
||||
//
|
||||
// To add a capability, append a rule. To reorder one host's runtimes, reorder
|
||||
// its tokens. Nothing else needs to change.
|
||||
//
|
||||
// Matching is by capability PREFIX, because a detected capability is refined at
|
||||
// runtime ("nvidia" becomes "nvidia-cuda-12" when the toolkit is present) and a
|
||||
// preference is about the vendor rather than the point release. Rules are tried
|
||||
// in order, so a more specific prefix must precede any rule it shares a prefix
|
||||
// with.
|
||||
//
|
||||
// Tokens are matched as substrings of a backend name, so "cuda" covers
|
||||
// "cuda12-llama-cpp" and "mlx" covers "metal-mlx" alike. A backend matching no
|
||||
// token is not an error and is never discarded; it simply sorts below every
|
||||
// recognised one.
|
||||
type backendPreferenceRule struct {
|
||||
capabilityPrefix string
|
||||
tokens []string
|
||||
}
|
||||
|
||||
var backendPreferenceRules = []backendPreferenceRule{
|
||||
{Nvidia, []string{backendTokenCUDA, vulkan, backendTokenCPU}},
|
||||
{AMD, []string{backendTokenROCM, backendTokenHIP, vulkan, backendTokenCPU}},
|
||||
{Intel, []string{backendTokenSYCL, Intel, backendTokenCPU}},
|
||||
// MLX outranks metal because on Apple silicon it is the native accelerated
|
||||
// runtime, whereas a metal-enabled GGUF build is the portable engine merely
|
||||
// compiled with GPU offload.
|
||||
{metal, []string{backendTokenMLX, backendTokenMetal, backendTokenCPU}},
|
||||
{darwinX86, []string{darwinX86, backendTokenCPU}},
|
||||
{vulkan, []string{vulkan, backendTokenCPU}},
|
||||
}
|
||||
|
||||
// defaultBackendPreferenceTokens is what a host with no matching rule prefers.
|
||||
// A capability nobody has taught this table about degrades to plain CPU rather
|
||||
// than to an error or to an empty list.
|
||||
var defaultBackendPreferenceTokens = []string{backendTokenCPU}
|
||||
|
||||
var (
|
||||
cuda13DirExists bool
|
||||
cuda12DirExists bool
|
||||
@@ -186,24 +232,18 @@ func (s *SystemState) getSystemCapabilities() string {
|
||||
// backend implementation order for the current system capability. Callers can use
|
||||
// these tokens to select the most appropriate concrete backend among multiple
|
||||
// candidates sharing the same alias (e.g., "llama-cpp").
|
||||
// The rules live in backendPreferenceRules above; this function only looks them
|
||||
// up, so teaching LocalAI about a new runtime never means editing logic.
|
||||
func (s *SystemState) BackendPreferenceTokens() []string {
|
||||
capStr := strings.ToLower(s.getSystemCapabilities())
|
||||
switch {
|
||||
case strings.HasPrefix(capStr, Nvidia):
|
||||
return []string{backendTokenCUDA, vulkan, "cpu"}
|
||||
case strings.HasPrefix(capStr, AMD):
|
||||
return []string{backendTokenROCM, backendTokenHIP, vulkan, "cpu"}
|
||||
case strings.HasPrefix(capStr, Intel):
|
||||
return []string{backendTokenSYCL, Intel, "cpu"}
|
||||
case strings.HasPrefix(capStr, metal):
|
||||
return []string{backendTokenMetal, "cpu"}
|
||||
case strings.HasPrefix(capStr, darwinX86):
|
||||
return []string{"darwin-x86", "cpu"}
|
||||
case strings.HasPrefix(capStr, vulkan):
|
||||
return []string{vulkan, "cpu"}
|
||||
default:
|
||||
return []string{"cpu"}
|
||||
for _, rule := range backendPreferenceRules {
|
||||
if strings.HasPrefix(capStr, rule.capabilityPrefix) {
|
||||
// Copied so a caller cannot mutate the shared table out from under
|
||||
// every other host lookup.
|
||||
return slices.Clone(rule.tokens)
|
||||
}
|
||||
}
|
||||
return slices.Clone(defaultBackendPreferenceTokens)
|
||||
}
|
||||
|
||||
// DetectedCapability returns the raw detected capability string (e.g. "metal",
|
||||
|
||||
@@ -128,6 +128,56 @@ var _ = Describe("getSystemCapabilities", func() {
|
||||
)
|
||||
})
|
||||
|
||||
var _ = Describe("BackendPreferenceTokens", func() {
|
||||
var origEnv string
|
||||
|
||||
BeforeEach(func() {
|
||||
origEnv = os.Getenv(capabilityEnv)
|
||||
})
|
||||
|
||||
AfterEach(func() {
|
||||
if origEnv != "" {
|
||||
os.Setenv(capabilityEnv, origEnv)
|
||||
} else {
|
||||
os.Unsetenv(capabilityEnv)
|
||||
}
|
||||
})
|
||||
|
||||
tokensFor := func(capability string) []string {
|
||||
GinkgoHelper()
|
||||
Expect(os.Setenv(capabilityEnv, capability)).To(Succeed())
|
||||
return (&SystemState{}).BackendPreferenceTokens()
|
||||
}
|
||||
|
||||
It("ranks MLX above metal on apple silicon", func() {
|
||||
// MLX is the native accelerated runtime there, so it must outrank a
|
||||
// metal-enabled build of the portable engine.
|
||||
Expect(tokensFor(metal)).To(Equal([]string{"mlx", "metal", "cpu"}))
|
||||
})
|
||||
|
||||
It("keeps the vendor order every other capability already had", func() {
|
||||
Expect(tokensFor("nvidia-cuda-12")).To(Equal([]string{"cuda", "vulkan", "cpu"}))
|
||||
Expect(tokensFor(AMD)).To(Equal([]string{"rocm", "hip", "vulkan", "cpu"}))
|
||||
Expect(tokensFor(Intel)).To(Equal([]string{"sycl", "intel", "cpu"}))
|
||||
Expect(tokensFor(darwinX86)).To(Equal([]string{"darwin-x86", "cpu"}))
|
||||
Expect(tokensFor(vulkan)).To(Equal([]string{"vulkan", "cpu"}))
|
||||
})
|
||||
|
||||
It("degrades an unknown capability to cpu rather than to nothing", func() {
|
||||
Expect(tokensFor("some-future-accelerator")).To(Equal([]string{"cpu"}))
|
||||
})
|
||||
|
||||
It("matches a capability by prefix so a refined one keeps its vendor order", func() {
|
||||
Expect(tokensFor("nvidia-l4t-cuda-13")).To(Equal(tokensFor("nvidia")))
|
||||
})
|
||||
|
||||
It("hands out a copy so one caller cannot corrupt another's lookup", func() {
|
||||
first := tokensFor("nvidia")
|
||||
first[0] = "clobbered"
|
||||
Expect(tokensFor("nvidia")[0]).To(Equal("cuda"))
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("CapabilityFilterDisabled", func() {
|
||||
var origEnv string
|
||||
|
||||
|
||||
Reference in New Issue
Block a user