mirror of
https://github.com/mudler/LocalAI.git
synced 2026-07-30 09:57:57 -04:00
fix(distributed): backend discovery hid GPU-only backends behind the controller's capability (#10947)
* fix(backends): list backends runnable on worker nodes in distributed mode GET /backends/available filtered the gallery against the system state of the host serving the request. In a distributed deployment that host is the controller, which typically has no GPU, while the GPUs live on worker nodes. Any meta backend whose capabilities map lacks a "default" (or "cpu") key was therefore dropped from the listing entirely — longcat-video, vllm-omni, ltx-video, parakeet, edgetam and qwentts were invisible in the UI even though installing them by name on a GPU worker worked fine. Workers now report their own meta-backend capability at registration and the controller persists it on the node row. The controller cannot derive it: OS-dependent capabilities (metal, darwin-x86, nvidia-l4t) and the CUDA runtime refinements are only observable on the worker. Nodes registered before this field existed fall back to a coarse capability derived from their GPU vendor and VRAM. Backend discovery then evaluates compatibility as the union over healthy backend nodes, so a backend runnable on any node is offered while one no node can run stays hidden. Each remote capability is evaluated through a capability-pinned system state, otherwise a forced capability on the controller image (LOCALAI_FORCE_META_BACKEND_CAPABILITY or /run/localai/capability) would silently override every worker's verdict. With no registered nodes the listing is byte-for-byte what it was, so single-node deployments are unaffected. Also fixes the same-root-cause misclassification in /api/operations, which used the capability-filtered listing to decide whether an operation was a backend or a model install. A GPU-only backend installing on a worker is still a backend operation on the controller, so that lookup is now unfiltered. Assisted-by: Claude:claude-opus-4-8 golangci-lint Signed-off-by: Ettore Di Giacinto <mudler@localai.io> * fix(backends): union worker capabilities in backend discovery Implementation for the specs added in the previous commit, plus the two remaining discovery endpoints. Capability-filtered backend discovery evaluated compatibility against the system state of the host serving the request. In a distributed deployment that host is the controller, which typically has no GPU, while the GPUs live on worker nodes. Any meta backend whose capabilities map lacks a "default" (or "cpu") key was dropped entirely — longcat-video, vllm-omni, ltx-video, parakeet, edgetam and qwentts were invisible in the UI even though installing them by name on a GPU worker worked fine. Workers now report their own meta-backend capability at registration and the controller persists it on the node row. The controller cannot derive it: OS-dependent capabilities (metal, darwin-x86, nvidia-l4t) and the CUDA runtime refinements are only observable on the worker. Nodes registered before this field existed fall back to a coarse capability derived from their GPU vendor and VRAM. Discovery then evaluates compatibility as the union over healthy backend nodes, so a backend runnable on any node is offered while one no node can run stays hidden. Each remote capability is evaluated through a capability-pinned system state, otherwise a forced capability on the controller image (LOCALAI_FORCE_META_BACKEND_CAPABILITY or /run/localai/capability) would silently override every worker's verdict. With no registered nodes the listing is byte-for-byte what it was, so single-node deployments are unaffected. Four surfaces shared this root cause and are all routed through the same helper now: - GET /backends/available - GET /api/fine-tuning/backends - GET /api/quantization/backends - /api/operations backend-vs-model classification, which additionally had no reason to filter by capability at all: a GPU-only backend installing on a worker is still a backend operation on the controller, so that lookup is now unfiltered. Assisted-by: Claude:claude-opus-4-8 golangci-lint Signed-off-by: Ettore Di Giacinto <mudler@localai.io> --------- Signed-off-by: Ettore Di Giacinto <mudler@localai.io> Co-authored-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
committed by
GitHub
parent
963c637130
commit
b19afb192a
131
core/gallery/backends_nodecapability_test.go
Normal file
131
core/gallery/backends_nodecapability_test.go
Normal file
@@ -0,0 +1,131 @@
|
||||
package gallery_test
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
. "github.com/mudler/LocalAI/core/gallery"
|
||||
"github.com/mudler/LocalAI/pkg/system"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
// These specs cover backend discovery in distributed mode, where the machine
|
||||
// answering GET /backends/available (the controller) is not the machine that
|
||||
// will run the backend (a worker). Filtering the listing against the
|
||||
// controller's own hardware hid every GPU-only meta backend from admins.
|
||||
var _ = Describe("AvailableBackendsForCapabilities", func() {
|
||||
var (
|
||||
tempDir string
|
||||
galleryPath string
|
||||
galleries []config.Gallery
|
||||
// controller stands in for a GPU-less frontend pod: no vendor, so its
|
||||
// reported capability is "default".
|
||||
controller *system.SystemState
|
||||
)
|
||||
|
||||
writeGalleryYAML := func(backends []GalleryBackend) {
|
||||
data, err := yaml.Marshal(backends)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(os.WriteFile(galleryPath, data, 0644)).To(Succeed())
|
||||
}
|
||||
|
||||
names := func(backends GalleryElements[*GalleryBackend]) []string {
|
||||
out := []string{}
|
||||
for _, b := range backends {
|
||||
out = append(out, b.GetName())
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// longcatVideo mirrors the real gallery entry that triggered this bug: a
|
||||
// meta backend enumerating only NVIDIA variants, with neither a "default"
|
||||
// nor a "cpu" key for Capability() to fall back to.
|
||||
longcatVideo := GalleryBackend{
|
||||
Metadata: Metadata{Name: "longcat-video"},
|
||||
CapabilitiesMap: map[string]string{
|
||||
"nvidia": "longcat-video-nvidia",
|
||||
"nvidia-cuda-12": "longcat-video-cuda-12",
|
||||
"nvidia-cuda-13": "longcat-video-cuda-13",
|
||||
"nvidia-l4t-cuda-13": "longcat-video-l4t-cuda-13",
|
||||
},
|
||||
}
|
||||
|
||||
// cpuMeta is compatible with every host and must keep showing up in all
|
||||
// scenarios, proving the union widens the listing without replacing it.
|
||||
cpuMeta := GalleryBackend{
|
||||
Metadata: Metadata{Name: "whisper"},
|
||||
CapabilitiesMap: map[string]string{"default": "whisper-cpu", "nvidia": "whisper-cuda"},
|
||||
}
|
||||
|
||||
BeforeEach(func() {
|
||||
var err error
|
||||
tempDir, err = os.MkdirTemp("", "node-capability-test-*")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
DeferCleanup(func() {
|
||||
Expect(os.RemoveAll(tempDir)).To(Succeed())
|
||||
})
|
||||
|
||||
galleryPath = filepath.Join(tempDir, "gallery.yaml")
|
||||
writeGalleryYAML([]GalleryBackend{longcatVideo, cpuMeta})
|
||||
|
||||
galleries = []config.Gallery{{Name: "test-gallery", URL: "file://" + galleryPath}}
|
||||
controller = system.NewCapabilityState("default", system.WithBackendPath(tempDir))
|
||||
})
|
||||
|
||||
It("hides a GPU-only meta when no node capabilities are supplied", func() {
|
||||
backends, err := AvailableBackendsForCapabilities(galleries, controller, nil)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(names(backends)).To(ContainElement("whisper"))
|
||||
Expect(names(backends)).NotTo(ContainElement("longcat-video"))
|
||||
})
|
||||
|
||||
It("lists a GPU-only meta runnable on a registered worker node", func() {
|
||||
backends, err := AvailableBackendsForCapabilities(galleries, controller, []string{"nvidia-l4t-cuda-13"})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(names(backends)).To(ContainElement("longcat-video"))
|
||||
Expect(names(backends)).To(ContainElement("whisper"))
|
||||
})
|
||||
|
||||
It("unions across heterogeneous nodes rather than intersecting them", func() {
|
||||
writeGalleryYAML([]GalleryBackend{
|
||||
longcatVideo,
|
||||
{
|
||||
Metadata: Metadata{Name: "amd-only"},
|
||||
CapabilitiesMap: map[string]string{"amd": "amd-only-rocm"},
|
||||
},
|
||||
})
|
||||
|
||||
backends, err := AvailableBackendsForCapabilities(galleries, controller, []string{"nvidia", "amd"})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(names(backends)).To(ContainElements("longcat-video", "amd-only"))
|
||||
})
|
||||
|
||||
It("still excludes a meta no node can satisfy", func() {
|
||||
backends, err := AvailableBackendsForCapabilities(galleries, controller, []string{"amd"})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(names(backends)).NotTo(ContainElement("longcat-video"))
|
||||
})
|
||||
|
||||
It("filters concrete backends by node capability too", func() {
|
||||
writeGalleryYAML([]GalleryBackend{
|
||||
{Metadata: Metadata{Name: "some-backend-cuda"}, URI: "quay.io/test/cuda"},
|
||||
{Metadata: Metadata{Name: "some-backend-rocm"}, URI: "quay.io/test/rocm"},
|
||||
})
|
||||
|
||||
backends, err := AvailableBackendsForCapabilities(galleries, controller, []string{"nvidia"})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(names(backends)).To(ContainElement("some-backend-cuda"))
|
||||
Expect(names(backends)).NotTo(ContainElement("some-backend-rocm"))
|
||||
})
|
||||
|
||||
It("returns exactly the single-node listing when the node list is empty", func() {
|
||||
withNodes, err := AvailableBackendsForCapabilities(galleries, controller, []string{})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
baseline, err := AvailableBackends(galleries, controller)
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(names(withNodes)).To(Equal(names(baseline)))
|
||||
})
|
||||
})
|
||||
@@ -395,16 +395,57 @@ func triggerGalleryRefresh(galleries []config.Gallery, systemState *system.Syste
|
||||
|
||||
// List available backends
|
||||
func AvailableBackends(galleries []config.Gallery, systemState *system.SystemState) (GalleryElements[*GalleryBackend], error) {
|
||||
return availableBackendsWithFilter(galleries, systemState, true)
|
||||
return availableBackendsWithFilter(galleries, systemState, func(backend *GalleryBackend) bool {
|
||||
return backend.IsCompatibleWith(systemState)
|
||||
})
|
||||
}
|
||||
|
||||
// AvailableBackendsUnfiltered returns all available backends without filtering by system capability.
|
||||
func AvailableBackendsUnfiltered(galleries []config.Gallery, systemState *system.SystemState) (GalleryElements[*GalleryBackend], error) {
|
||||
return availableBackendsWithFilter(galleries, systemState, false)
|
||||
return availableBackendsWithFilter(galleries, systemState, nil)
|
||||
}
|
||||
|
||||
// availableBackendsWithFilter is a helper function that lists available backends with optional filtering.
|
||||
func availableBackendsWithFilter(galleries []config.Gallery, systemState *system.SystemState, filterByCapability bool) (GalleryElements[*GalleryBackend], error) {
|
||||
// AvailableBackendsForCapabilities lists backends runnable on the local system
|
||||
// OR on any remote host reporting one of the supplied capabilities.
|
||||
//
|
||||
// In a distributed deployment the host serving this listing (the controller)
|
||||
// is usually a GPU-less pod while the GPUs live on worker nodes. Filtering
|
||||
// only against the controller hid every GPU-only meta backend from admins even
|
||||
// though installing it by name on a worker worked fine, so compatibility is
|
||||
// evaluated as a union over the cluster. An empty capabilities slice reproduces
|
||||
// AvailableBackends exactly, keeping single-node behavior untouched.
|
||||
func AvailableBackendsForCapabilities(galleries []config.Gallery, systemState *system.SystemState, capabilities []string) (GalleryElements[*GalleryBackend], error) {
|
||||
if len(capabilities) == 0 {
|
||||
return AvailableBackends(galleries, systemState)
|
||||
}
|
||||
|
||||
// Each remote capability is evaluated through a state pinned to that exact
|
||||
// capability, so the controller's own detection (and any forced capability
|
||||
// on the controller image) cannot leak into the worker's verdict. Backend
|
||||
// paths still come from the controller's state because that is where the
|
||||
// gallery metadata is read from.
|
||||
nodeStates := make([]*system.SystemState, 0, len(capabilities))
|
||||
for _, capability := range capabilities {
|
||||
nodeStates = append(nodeStates, system.NewCapabilityState(capability,
|
||||
system.WithBackendPath(systemState.Backend.BackendsPath)))
|
||||
}
|
||||
|
||||
return availableBackendsWithFilter(galleries, systemState, func(backend *GalleryBackend) bool {
|
||||
if backend.IsCompatibleWith(systemState) {
|
||||
return true
|
||||
}
|
||||
for _, nodeState := range nodeStates {
|
||||
if backend.IsCompatibleWith(nodeState) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
})
|
||||
}
|
||||
|
||||
// availableBackendsWithFilter lists available backends, keeping only those
|
||||
// accepted by compatible. A nil compatible keeps everything.
|
||||
func availableBackendsWithFilter(galleries []config.Gallery, systemState *system.SystemState, compatible func(*GalleryBackend) bool) (GalleryElements[*GalleryBackend], error) {
|
||||
var backends []*GalleryBackend
|
||||
|
||||
systemBackends, err := ListSystemBackends(systemState)
|
||||
@@ -421,15 +462,15 @@ func availableBackendsWithFilter(galleries []config.Gallery, systemState *system
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Filter backends by system capability if requested
|
||||
if filterByCapability {
|
||||
for _, backend := range galleryBackends {
|
||||
if backend.IsCompatibleWith(systemState) {
|
||||
backends = append(backends, backend)
|
||||
}
|
||||
}
|
||||
} else {
|
||||
if compatible == nil {
|
||||
backends = append(backends, galleryBackends...)
|
||||
continue
|
||||
}
|
||||
|
||||
for _, backend := range galleryBackends {
|
||||
if compatible(backend) {
|
||||
backends = append(backends, backend)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -443,7 +443,7 @@ func API(application *application.Application) (*echo.Echo, error) {
|
||||
ftNats,
|
||||
ftStore,
|
||||
)
|
||||
routes.RegisterFineTuningRoutes(e, ftService, application.ApplicationConfig(), fineTuningMw)
|
||||
routes.RegisterFineTuningRoutes(e, ftService, application.ApplicationConfig(), application, fineTuningMw)
|
||||
|
||||
// Quantization routes
|
||||
quantizationMw := auth.RequireFeature(application.AuthDB(), auth.FeatureQuantization)
|
||||
@@ -465,7 +465,7 @@ func API(application *application.Application) (*echo.Echo, error) {
|
||||
quantNats,
|
||||
quantStore,
|
||||
)
|
||||
routes.RegisterQuantizationRoutes(e, qService, application.ApplicationConfig(), quantizationMw)
|
||||
routes.RegisterQuantizationRoutes(e, qService, application.ApplicationConfig(), application, quantizationMw)
|
||||
|
||||
// Node management routes (distributed mode)
|
||||
distCfg := application.ApplicationConfig().Distributed
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package localai
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"sort"
|
||||
@@ -323,14 +324,41 @@ func (mgs *BackendEndpointService) UpgradeBackendEndpoint() echo.HandlerFunc {
|
||||
}
|
||||
}
|
||||
|
||||
// ClusterCapabilityProvider reports the meta-backend capabilities available
|
||||
// somewhere in the cluster. It is nil in single-node deployments, where the
|
||||
// local system state is the only thing worth filtering against.
|
||||
type ClusterCapabilityProvider func(ctx context.Context) ([]string, error)
|
||||
|
||||
// resolveClusterCapabilities reads the capabilities present in the cluster,
|
||||
// degrading to the local-only listing on error.
|
||||
//
|
||||
// Every capability-filtered discovery endpoint shares this: on a distributed
|
||||
// controller the GPUs live on the workers, so filtering against the local
|
||||
// (usually GPU-less) host hides GPU-only backends the cluster can actually
|
||||
// run. A registry hiccup must never blank the catalog, so a failure falls back
|
||||
// to the pre-existing local-only behavior rather than erroring the request.
|
||||
func resolveClusterCapabilities(ctx context.Context, provider ClusterCapabilityProvider) []string {
|
||||
if provider == nil {
|
||||
return nil
|
||||
}
|
||||
capabilities, err := provider(ctx)
|
||||
if err != nil {
|
||||
xlog.Warn("Could not read cluster capabilities, listing backends for the local system only", "error", err)
|
||||
return nil
|
||||
}
|
||||
return capabilities
|
||||
}
|
||||
|
||||
// ListAvailableBackendsEndpoint list the available backends in the galleries configured in LocalAI
|
||||
// @Summary List all available Backends
|
||||
// @Tags backends
|
||||
// @Success 200 {object} []gallery.GalleryBackend "Response"
|
||||
// @Router /backends/available [get]
|
||||
func (mgs *BackendEndpointService) ListAvailableBackendsEndpoint(systemState *system.SystemState) echo.HandlerFunc {
|
||||
func (mgs *BackendEndpointService) ListAvailableBackendsEndpoint(systemState *system.SystemState, clusterCapabilities ClusterCapabilityProvider) echo.HandlerFunc {
|
||||
return func(c echo.Context) error {
|
||||
backends, err := gallery.AvailableBackends(mgs.galleries, systemState)
|
||||
capabilities := resolveClusterCapabilities(c.Request().Context(), clusterCapabilities)
|
||||
|
||||
backends, err := gallery.AvailableBackendsForCapabilities(mgs.galleries, systemState, capabilities)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
109
core/http/endpoints/localai/backend_available_cluster_test.go
Normal file
109
core/http/endpoints/localai/backend_available_cluster_test.go
Normal file
@@ -0,0 +1,109 @@
|
||||
package localai_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
"github.com/labstack/echo/v4"
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/gallery"
|
||||
. "github.com/mudler/LocalAI/core/http/endpoints/localai"
|
||||
"github.com/mudler/LocalAI/pkg/system"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
// GET /backends/available on a distributed controller must advertise what the
|
||||
// cluster can run, not what the controller itself can run.
|
||||
var _ = Describe("ListAvailableBackendsEndpoint cluster capabilities", func() {
|
||||
var (
|
||||
app *echo.Echo
|
||||
systemState *system.SystemState
|
||||
tmpDir string
|
||||
galleries []config.Gallery
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
app = echo.New()
|
||||
|
||||
var err error
|
||||
tmpDir, err = os.MkdirTemp("", "backends-available-cluster-*")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
DeferCleanup(func() {
|
||||
Expect(os.RemoveAll(tmpDir)).To(Succeed())
|
||||
})
|
||||
|
||||
galleryPath := filepath.Join(tmpDir, "gallery.yaml")
|
||||
data, err := yaml.Marshal([]gallery.GalleryBackend{
|
||||
{
|
||||
Metadata: gallery.Metadata{Name: "longcat-video"},
|
||||
CapabilitiesMap: map[string]string{
|
||||
"nvidia": "longcat-video-nvidia",
|
||||
"nvidia-cuda-13": "longcat-video-cuda-13",
|
||||
"nvidia-l4t-cuda-13": "longcat-video-l4t-cuda-13",
|
||||
},
|
||||
},
|
||||
{
|
||||
Metadata: gallery.Metadata{Name: "whisper"},
|
||||
CapabilitiesMap: map[string]string{"default": "whisper-cpu"},
|
||||
},
|
||||
})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(os.WriteFile(galleryPath, data, 0o644)).To(Succeed())
|
||||
|
||||
galleries = []config.Gallery{{Name: "test-gallery", URL: "file://" + galleryPath}}
|
||||
|
||||
// A GPU-less controller pod: capability "default".
|
||||
systemState = system.NewCapabilityState("default",
|
||||
system.WithBackendPath(tmpDir), system.WithBackendSystemPath(tmpDir))
|
||||
})
|
||||
|
||||
listNames := func(provider ClusterCapabilityProvider) []string {
|
||||
svc := CreateBackendEndpointService(galleries, systemState, nil, nil)
|
||||
app.GET("/backends/available", svc.ListAvailableBackendsEndpoint(systemState, provider))
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/backends/available", nil)
|
||||
rec := httptest.NewRecorder()
|
||||
app.ServeHTTP(rec, req)
|
||||
Expect(rec.Code).To(Equal(http.StatusOK))
|
||||
|
||||
var backends []gallery.GalleryBackend
|
||||
Expect(json.Unmarshal(rec.Body.Bytes(), &backends)).To(Succeed())
|
||||
|
||||
names := []string{}
|
||||
for _, b := range backends {
|
||||
names = append(names, b.Name)
|
||||
}
|
||||
return names
|
||||
}
|
||||
|
||||
It("hides GPU-only backends with no provider (single-node)", func() {
|
||||
names := listNames(nil)
|
||||
Expect(names).To(ContainElement("whisper"))
|
||||
Expect(names).NotTo(ContainElement("longcat-video"))
|
||||
})
|
||||
|
||||
It("lists GPU-only backends a worker node can run", func() {
|
||||
provider := func(ctx context.Context) ([]string, error) {
|
||||
return []string{"nvidia-l4t-cuda-13"}, nil
|
||||
}
|
||||
Expect(listNames(provider)).To(ContainElements("whisper", "longcat-video"))
|
||||
})
|
||||
|
||||
It("falls back to the local listing when the registry errors", func() {
|
||||
// A registry hiccup must degrade to the old behavior, never blank the
|
||||
// backend catalog with a 500.
|
||||
provider := func(ctx context.Context) ([]string, error) {
|
||||
return nil, errors.New("registry unavailable")
|
||||
}
|
||||
names := listNames(provider)
|
||||
Expect(names).To(ContainElement("whisper"))
|
||||
Expect(names).NotTo(ContainElement("longcat-video"))
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,126 @@
|
||||
package localai_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
"github.com/labstack/echo/v4"
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/gallery"
|
||||
. "github.com/mudler/LocalAI/core/http/endpoints/localai"
|
||||
"github.com/mudler/LocalAI/pkg/system"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
// The fine-tune and quantization backend lists are discovery endpoints with
|
||||
// the same semantics as /backends/available: they populate UI dropdowns. On a
|
||||
// GPU-less controller they hid GPU-only backends the cluster could run, so an
|
||||
// admin with a fine-tuning-capable GPU worker saw an empty dropdown.
|
||||
var _ = Describe("Tagged backend discovery with cluster capabilities", func() {
|
||||
var (
|
||||
appCfg *config.ApplicationConfig
|
||||
tmpDir string
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
var err error
|
||||
tmpDir, err = os.MkdirTemp("", "tagged-discovery-cluster-*")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
DeferCleanup(func() {
|
||||
Expect(os.RemoveAll(tmpDir)).To(Succeed())
|
||||
})
|
||||
|
||||
// Both endpoints only surface *installed* backends, so lay the
|
||||
// backend down on disk; the capability filter is then the only thing
|
||||
// that can remove it from the response.
|
||||
writeFakeSystemBackend(tmpDir, "gpu-trainer")
|
||||
writeFakeSystemBackend(tmpDir, "cpu-trainer")
|
||||
|
||||
galleryPath := filepath.Join(tmpDir, "gallery.yaml")
|
||||
data, err := yaml.Marshal([]gallery.GalleryBackend{
|
||||
{
|
||||
Metadata: gallery.Metadata{
|
||||
Name: "gpu-trainer",
|
||||
Tags: []string{"fine-tuning", "quantization"},
|
||||
},
|
||||
CapabilitiesMap: map[string]string{
|
||||
"nvidia": "gpu-trainer-nvidia",
|
||||
"nvidia-cuda-13": "gpu-trainer-cuda-13",
|
||||
},
|
||||
},
|
||||
{
|
||||
Metadata: gallery.Metadata{
|
||||
Name: "cpu-trainer",
|
||||
Tags: []string{"fine-tuning", "quantization"},
|
||||
},
|
||||
CapabilitiesMap: map[string]string{"default": "cpu-trainer-cpu"},
|
||||
},
|
||||
})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(os.WriteFile(galleryPath, data, 0o644)).To(Succeed())
|
||||
|
||||
appCfg = &config.ApplicationConfig{
|
||||
BackendGalleries: []config.Gallery{{Name: "test-gallery", URL: "file://" + galleryPath}},
|
||||
SystemState: system.NewCapabilityState("default",
|
||||
system.WithBackendPath(tmpDir), system.WithBackendSystemPath(tmpDir)),
|
||||
}
|
||||
})
|
||||
|
||||
// listNames drives handler over its route and returns the backend names.
|
||||
listNames := func(path string, handler echo.HandlerFunc) []string {
|
||||
e := echo.New()
|
||||
e.GET(path, handler)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, path, nil)
|
||||
rec := httptest.NewRecorder()
|
||||
e.ServeHTTP(rec, req)
|
||||
Expect(rec.Code).To(Equal(http.StatusOK))
|
||||
|
||||
var backends []struct {
|
||||
Name string `json:"name"`
|
||||
}
|
||||
Expect(json.Unmarshal(rec.Body.Bytes(), &backends)).To(Succeed())
|
||||
|
||||
names := []string{}
|
||||
for _, b := range backends {
|
||||
names = append(names, b.Name)
|
||||
}
|
||||
return names
|
||||
}
|
||||
|
||||
nvidiaWorker := func(ctx context.Context) ([]string, error) {
|
||||
return []string{"nvidia-cuda-13"}, nil
|
||||
}
|
||||
|
||||
Describe("fine-tune backends", func() {
|
||||
It("hides GPU-only backends with no cluster provider", func() {
|
||||
names := listNames("/backends", ListFineTuneBackendsEndpoint(appCfg, nil))
|
||||
Expect(names).To(ContainElement("cpu-trainer"))
|
||||
Expect(names).NotTo(ContainElement("gpu-trainer"))
|
||||
})
|
||||
|
||||
It("lists GPU-only backends a worker node can run", func() {
|
||||
names := listNames("/backends", ListFineTuneBackendsEndpoint(appCfg, nvidiaWorker))
|
||||
Expect(names).To(ContainElements("cpu-trainer", "gpu-trainer"))
|
||||
})
|
||||
})
|
||||
|
||||
Describe("quantization backends", func() {
|
||||
It("hides GPU-only backends with no cluster provider", func() {
|
||||
names := listNames("/backends", ListQuantizationBackendsEndpoint(appCfg, nil))
|
||||
Expect(names).To(ContainElement("cpu-trainer"))
|
||||
Expect(names).NotTo(ContainElement("gpu-trainer"))
|
||||
})
|
||||
|
||||
It("lists GPU-only backends a worker node can run", func() {
|
||||
names := listNames("/backends", ListQuantizationBackendsEndpoint(appCfg, nvidiaWorker))
|
||||
Expect(names).To(ContainElements("cpu-trainer", "gpu-trainer"))
|
||||
})
|
||||
})
|
||||
})
|
||||
@@ -274,9 +274,10 @@ func DownloadExportedModelEndpoint(ftService *finetune.FineTuneService) echo.Han
|
||||
}
|
||||
|
||||
// ListFineTuneBackendsEndpoint returns installed backends tagged with "fine-tuning".
|
||||
func ListFineTuneBackendsEndpoint(appConfig *config.ApplicationConfig) echo.HandlerFunc {
|
||||
func ListFineTuneBackendsEndpoint(appConfig *config.ApplicationConfig, clusterCapabilities ClusterCapabilityProvider) echo.HandlerFunc {
|
||||
return func(c echo.Context) error {
|
||||
backends, err := gallery.AvailableBackends(appConfig.BackendGalleries, appConfig.SystemState)
|
||||
capabilities := resolveClusterCapabilities(c.Request().Context(), clusterCapabilities)
|
||||
backends, err := gallery.AvailableBackendsForCapabilities(appConfig.BackendGalleries, appConfig.SystemState, capabilities)
|
||||
if err != nil {
|
||||
return c.JSON(http.StatusInternalServerError, map[string]string{
|
||||
"error": "failed to list backends: " + err.Error(),
|
||||
|
||||
@@ -87,8 +87,12 @@ type RegisterNodeRequest struct {
|
||||
GPUVendor string `json:"gpu_vendor,omitempty"`
|
||||
// GPUComputeCapability is the worker GPU's compute capability ("major.minor",
|
||||
// e.g. "12.1" for GB10). Used by the router for per-arch option tuning.
|
||||
GPUComputeCapability string `json:"gpu_compute_capability,omitempty"`
|
||||
Labels map[string]string `json:"labels,omitempty"`
|
||||
GPUComputeCapability string `json:"gpu_compute_capability,omitempty"`
|
||||
// Capability is the worker's own meta-backend capability string. The
|
||||
// controller stores it so backend discovery can reflect what the cluster
|
||||
// can run rather than what the (typically GPU-less) controller can run.
|
||||
Capability string `json:"capability,omitempty"`
|
||||
Labels map[string]string `json:"labels,omitempty"`
|
||||
// MaxReplicasPerModel is the per-node cap on replicas of any single model.
|
||||
// Workers older than this field omit it; we coerce 0 → 1 below to preserve
|
||||
// historical single-replica behavior.
|
||||
@@ -174,6 +178,7 @@ func RegisterNodeEndpoint(registry *nodes.NodeRegistry, expectedToken string, au
|
||||
AvailableRAM: req.AvailableRAM,
|
||||
GPUVendor: req.GPUVendor,
|
||||
GPUComputeCapability: req.GPUComputeCapability,
|
||||
Capability: req.Capability,
|
||||
MaxReplicasPerModel: maxReplicasPerModel,
|
||||
VRAMBudget: req.VRAMBudget,
|
||||
}
|
||||
|
||||
@@ -193,9 +193,10 @@ func DownloadQuantizedModelEndpoint(qService *quantization.QuantizationService)
|
||||
}
|
||||
|
||||
// ListQuantizationBackendsEndpoint returns installed backends tagged with "quantization".
|
||||
func ListQuantizationBackendsEndpoint(appConfig *config.ApplicationConfig) echo.HandlerFunc {
|
||||
func ListQuantizationBackendsEndpoint(appConfig *config.ApplicationConfig, clusterCapabilities ClusterCapabilityProvider) echo.HandlerFunc {
|
||||
return func(c echo.Context) error {
|
||||
backends, err := gallery.AvailableBackends(appConfig.BackendGalleries, appConfig.SystemState)
|
||||
capabilities := resolveClusterCapabilities(c.Request().Context(), clusterCapabilities)
|
||||
backends, err := gallery.AvailableBackendsForCapabilities(appConfig.BackendGalleries, appConfig.SystemState, capabilities)
|
||||
if err != nil {
|
||||
return c.JSON(http.StatusInternalServerError, map[string]string{
|
||||
"error": "failed to list backends: " + err.Error(),
|
||||
|
||||
20
core/http/routes/cluster_capabilities.go
Normal file
20
core/http/routes/cluster_capabilities.go
Normal file
@@ -0,0 +1,20 @@
|
||||
package routes
|
||||
|
||||
import (
|
||||
"github.com/mudler/LocalAI/core/application"
|
||||
"github.com/mudler/LocalAI/core/http/endpoints/localai"
|
||||
)
|
||||
|
||||
// ClusterCapabilityProviderFor returns the capability source backing every
|
||||
// capability-filtered backend discovery endpoint, or nil in single-node mode.
|
||||
//
|
||||
// A nil provider makes those endpoints filter against the local system exactly
|
||||
// as they always have. In distributed mode the controller is typically a
|
||||
// GPU-less pod, so discovery instead unions the capabilities of the healthy
|
||||
// worker nodes that would actually run the backend.
|
||||
func ClusterCapabilityProviderFor(app *application.Application) localai.ClusterCapabilityProvider {
|
||||
if app == nil || !app.IsDistributed() || app.Distributed().Registry == nil {
|
||||
return nil
|
||||
}
|
||||
return app.Distributed().Registry.HealthyBackendCapabilities
|
||||
}
|
||||
@@ -4,13 +4,14 @@ import (
|
||||
"net/http"
|
||||
|
||||
"github.com/labstack/echo/v4"
|
||||
"github.com/mudler/LocalAI/core/application"
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/http/endpoints/localai"
|
||||
"github.com/mudler/LocalAI/core/services/finetune"
|
||||
)
|
||||
|
||||
// RegisterFineTuningRoutes registers fine-tuning API routes.
|
||||
func RegisterFineTuningRoutes(e *echo.Echo, ftService *finetune.FineTuneService, appConfig *config.ApplicationConfig, fineTuningMw echo.MiddlewareFunc) {
|
||||
func RegisterFineTuningRoutes(e *echo.Echo, ftService *finetune.FineTuneService, appConfig *config.ApplicationConfig, app *application.Application, fineTuningMw echo.MiddlewareFunc) {
|
||||
if ftService == nil {
|
||||
return
|
||||
}
|
||||
@@ -28,7 +29,7 @@ func RegisterFineTuningRoutes(e *echo.Echo, ftService *finetune.FineTuneService,
|
||||
}
|
||||
|
||||
ft := e.Group("/api/fine-tuning", readyMw, fineTuningMw)
|
||||
ft.GET("/backends", localai.ListFineTuneBackendsEndpoint(appConfig))
|
||||
ft.GET("/backends", localai.ListFineTuneBackendsEndpoint(appConfig, ClusterCapabilityProviderFor(app)))
|
||||
ft.POST("/jobs", localai.StartFineTuneJobEndpoint(ftService))
|
||||
ft.GET("/jobs", localai.ListFineTuneJobsEndpoint(ftService))
|
||||
ft.GET("/jobs/:id", localai.GetFineTuneJobEndpoint(ftService))
|
||||
|
||||
@@ -64,7 +64,7 @@ func RegisterLocalAIRoutes(router *echo.Echo,
|
||||
router.POST("/backends/apply", backendGalleryEndpointService.ApplyBackendEndpoint(appConfig.SystemState), adminMiddleware)
|
||||
router.POST("/backends/delete/:name", backendGalleryEndpointService.DeleteBackendEndpoint(), adminMiddleware)
|
||||
router.GET("/backends", backendGalleryEndpointService.ListBackendsEndpoint(), adminMiddleware)
|
||||
router.GET("/backends/available", backendGalleryEndpointService.ListAvailableBackendsEndpoint(appConfig.SystemState), adminMiddleware)
|
||||
router.GET("/backends/available", backendGalleryEndpointService.ListAvailableBackendsEndpoint(appConfig.SystemState, ClusterCapabilityProviderFor(app)), adminMiddleware)
|
||||
router.GET("/backends/known", backendGalleryEndpointService.ListKnownBackendsEndpoint(appConfig.SystemState), adminMiddleware)
|
||||
router.GET("/backends/galleries", backendGalleryEndpointService.ListBackendGalleriesEndpoint(), adminMiddleware)
|
||||
router.GET("/backends/jobs/:uuid", backendGalleryEndpointService.GetOpStatusEndpoint(), adminMiddleware)
|
||||
|
||||
@@ -4,13 +4,14 @@ import (
|
||||
"net/http"
|
||||
|
||||
"github.com/labstack/echo/v4"
|
||||
"github.com/mudler/LocalAI/core/application"
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/http/endpoints/localai"
|
||||
"github.com/mudler/LocalAI/core/services/quantization"
|
||||
)
|
||||
|
||||
// RegisterQuantizationRoutes registers quantization API routes.
|
||||
func RegisterQuantizationRoutes(e *echo.Echo, qService *quantization.QuantizationService, appConfig *config.ApplicationConfig, quantizationMw echo.MiddlewareFunc) {
|
||||
func RegisterQuantizationRoutes(e *echo.Echo, qService *quantization.QuantizationService, appConfig *config.ApplicationConfig, app *application.Application, quantizationMw echo.MiddlewareFunc) {
|
||||
if qService == nil {
|
||||
return
|
||||
}
|
||||
@@ -28,7 +29,7 @@ func RegisterQuantizationRoutes(e *echo.Echo, qService *quantization.Quantizatio
|
||||
}
|
||||
|
||||
q := e.Group("/api/quantization", readyMw, quantizationMw)
|
||||
q.GET("/backends", localai.ListQuantizationBackendsEndpoint(appConfig))
|
||||
q.GET("/backends", localai.ListQuantizationBackendsEndpoint(appConfig, ClusterCapabilityProviderFor(app)))
|
||||
q.POST("/jobs", localai.StartQuantizationJobEndpoint(qService))
|
||||
q.GET("/jobs", localai.ListQuantizationJobsEndpoint(qService))
|
||||
q.GET("/jobs/:id", localai.GetQuantizationJobEndpoint(qService))
|
||||
|
||||
@@ -211,9 +211,13 @@ func RegisterUIAPIRoutes(app *echo.Echo, cl *config.ModelConfigLoader, ml *model
|
||||
// Determine if it's a model or backend
|
||||
// First check if it was explicitly marked as a backend operation
|
||||
isBackend := opcache.IsBackendOp(galleryID)
|
||||
// If not explicitly marked, check if it matches a known backend from the gallery
|
||||
// If not explicitly marked, check if it matches a known backend from the gallery.
|
||||
// Unfiltered on purpose: this only classifies an operation as
|
||||
// backend-vs-model, and a GPU-only backend installed on a worker is
|
||||
// still a backend op even when this host cannot run it. The
|
||||
// capability-filtered listing misclassified those as model ops.
|
||||
if !isBackend {
|
||||
backends, _ := gallery.AvailableBackends(appConfig.BackendGalleries, appConfig.SystemState)
|
||||
backends, _ := gallery.AvailableBackendsUnfiltered(appConfig.BackendGalleries, appConfig.SystemState)
|
||||
for _, b := range backends {
|
||||
backendID := fmt.Sprintf("%s@%s", b.Gallery.Name, b.Name)
|
||||
if backendID == galleryID || b.Name == galleryID {
|
||||
|
||||
@@ -0,0 +1,87 @@
|
||||
package routes_test
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
"github.com/labstack/echo/v4"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
"gopkg.in/yaml.v3"
|
||||
|
||||
"github.com/mudler/LocalAI/core/application"
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/gallery"
|
||||
"github.com/mudler/LocalAI/core/http/routes"
|
||||
"github.com/mudler/LocalAI/core/services/galleryop"
|
||||
"github.com/mudler/LocalAI/pkg/system"
|
||||
)
|
||||
|
||||
// /api/operations classifies an op as backend-vs-model by looking the name up
|
||||
// in the backend gallery. That lookup must not be capability-filtered: on a
|
||||
// distributed controller the GPU lives on a worker, so a GPU-only backend
|
||||
// install was misfiled as a model operation in the UI panel.
|
||||
var _ = Describe("/api/operations backend classification", func() {
|
||||
noopMw := func(next echo.HandlerFunc) echo.HandlerFunc { return next }
|
||||
|
||||
It("classifies a GPU-only backend op as a backend on a GPU-less host", func() {
|
||||
tmpDir, err := os.MkdirTemp("", "ops-classification-*")
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
DeferCleanup(func() {
|
||||
Expect(os.RemoveAll(tmpDir)).To(Succeed())
|
||||
})
|
||||
|
||||
galleryPath := filepath.Join(tmpDir, "gallery.yaml")
|
||||
data, err := yaml.Marshal([]gallery.GalleryBackend{
|
||||
{
|
||||
Metadata: gallery.Metadata{Name: "longcat-video"},
|
||||
CapabilitiesMap: map[string]string{
|
||||
"nvidia": "longcat-video-nvidia",
|
||||
"nvidia-cuda-13": "longcat-video-cuda-13",
|
||||
"nvidia-l4t-cuda-13": "longcat-video-l4t-cuda-13",
|
||||
},
|
||||
},
|
||||
})
|
||||
Expect(err).NotTo(HaveOccurred())
|
||||
Expect(os.WriteFile(galleryPath, data, 0o644)).To(Succeed())
|
||||
|
||||
appCfg := &config.ApplicationConfig{
|
||||
BackendGalleries: []config.Gallery{{Name: "test-gallery", URL: "file://" + galleryPath}},
|
||||
SystemState: system.NewCapabilityState("default",
|
||||
system.WithBackendPath(tmpDir), system.WithBackendSystemPath(tmpDir)),
|
||||
}
|
||||
galleryService := galleryop.NewGalleryService(appCfg, nil)
|
||||
opcache := galleryop.NewOpCache(galleryService)
|
||||
|
||||
jobID := "job-longcat-1"
|
||||
// Set (not SetBackend) so the handler has to fall back to the gallery
|
||||
// lookup — exactly the path that was misclassifying.
|
||||
opcache.Set("longcat-video", jobID)
|
||||
|
||||
e := echo.New()
|
||||
routes.RegisterUIAPIRoutes(e, nil, nil, appCfg, galleryService, opcache, &application.Application{}, noopMw)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/api/operations", nil)
|
||||
rec := httptest.NewRecorder()
|
||||
e.ServeHTTP(rec, req)
|
||||
Expect(rec.Code).To(Equal(http.StatusOK))
|
||||
|
||||
var envelope struct {
|
||||
Operations []map[string]any `json:"operations"`
|
||||
}
|
||||
Expect(json.Unmarshal(rec.Body.Bytes(), &envelope)).To(Succeed())
|
||||
|
||||
var found map[string]any
|
||||
for _, op := range envelope.Operations {
|
||||
if op["jobID"] == jobID {
|
||||
found = op
|
||||
break
|
||||
}
|
||||
}
|
||||
Expect(found).ToNot(BeNil(), "operation should appear in /api/operations")
|
||||
Expect(found["isBackend"]).To(Equal(true))
|
||||
})
|
||||
})
|
||||
@@ -9,6 +9,7 @@ import (
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/mudler/LocalAI/core/services/advisorylock"
|
||||
"github.com/mudler/LocalAI/pkg/system"
|
||||
"github.com/mudler/LocalAI/pkg/vrambudget"
|
||||
"github.com/mudler/xlog"
|
||||
"gorm.io/gorm"
|
||||
@@ -42,6 +43,13 @@ type BackendNode struct {
|
||||
// on registration; used by the router to pick per-arch options (e.g. a
|
||||
// larger physical batch on Blackwell). Empty when unknown / non-NVIDIA.
|
||||
GPUComputeCapability string `gorm:"column:gpu_compute_capability;size:16" json:"gpu_compute_capability"`
|
||||
// Capability is the worker's own meta-backend capability string (e.g.
|
||||
// "nvidia-cuda-13", "metal", "default"), reported at registration. The
|
||||
// controller cannot derive it: OS-dependent capabilities (metal,
|
||||
// darwin-x86, nvidia-l4t) and the CUDA-runtime refinements are only
|
||||
// visible on the worker itself. Empty for workers registered before this
|
||||
// field existed; readers fall back to system.CapabilityFromGPU.
|
||||
Capability string `gorm:"column:capability;size:64" json:"capability,omitempty"`
|
||||
// MaxReplicasPerModel caps how many replicas of any one model can run on
|
||||
// this node concurrently. Default 1 preserves the historical "one
|
||||
// (node, model)" assumption; set higher (via worker --max-replicas-per-model)
|
||||
@@ -747,6 +755,38 @@ func (r *NodeRegistry) List(ctx context.Context) ([]BackendNode, error) {
|
||||
return nodes, nil
|
||||
}
|
||||
|
||||
// HealthyBackendCapabilities returns the deduplicated meta-backend capability
|
||||
// strings of every healthy backend node.
|
||||
//
|
||||
// Backend discovery on the controller unions these so a GPU-only backend that
|
||||
// no node can run today stays hidden, while one that any worker can run is
|
||||
// offered. Only healthy backend nodes count — the same predicate the scheduler
|
||||
// places against — so an offline worker does not keep advertising hardware the
|
||||
// cluster cannot currently use.
|
||||
func (r *NodeRegistry) HealthyBackendCapabilities(ctx context.Context) ([]string, error) {
|
||||
var nodes []BackendNode
|
||||
if err := r.db.WithContext(ctx).
|
||||
Where("status = ? AND node_type = ?", StatusHealthy, NodeTypeBackend).
|
||||
Find(&nodes).Error; err != nil {
|
||||
return nil, fmt.Errorf("listing healthy backend node capabilities: %w", err)
|
||||
}
|
||||
|
||||
seen := make(map[string]struct{}, len(nodes))
|
||||
capabilities := make([]string, 0, len(nodes))
|
||||
for _, node := range nodes {
|
||||
capability := node.Capability
|
||||
if capability == "" {
|
||||
capability = system.CapabilityFromGPU(node.GPUVendor, node.TotalVRAM)
|
||||
}
|
||||
if _, dup := seen[capability]; dup {
|
||||
continue
|
||||
}
|
||||
seen[capability] = struct{}{}
|
||||
capabilities = append(capabilities, capability)
|
||||
}
|
||||
return capabilities, nil
|
||||
}
|
||||
|
||||
// Get returns a single node by ID.
|
||||
func (r *NodeRegistry) Get(ctx context.Context, nodeID string) (*BackendNode, error) {
|
||||
var node BackendNode
|
||||
|
||||
123
core/services/nodes/registry_capability_test.go
Normal file
123
core/services/nodes/registry_capability_test.go
Normal file
@@ -0,0 +1,123 @@
|
||||
package nodes
|
||||
|
||||
import (
|
||||
"context"
|
||||
"runtime"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/testutil"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// Backend discovery on a GPU-less controller unions these capabilities, so the
|
||||
// set this returns decides which GPU-only backends admins can see at all.
|
||||
var _ = Describe("NodeRegistry HealthyBackendCapabilities", func() {
|
||||
var (
|
||||
db *gorm.DB
|
||||
registry *NodeRegistry
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
if runtime.GOOS == "darwin" {
|
||||
Skip("testcontainers requires Docker, not available on macOS CI")
|
||||
}
|
||||
db = testutil.SetupTestDB()
|
||||
var err error
|
||||
registry, err = NewNodeRegistry(db)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
})
|
||||
|
||||
register := func(node *BackendNode) {
|
||||
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
||||
}
|
||||
|
||||
It("returns the worker-reported capability", func() {
|
||||
register(&BackendNode{
|
||||
Name: "gpu-worker", NodeType: NodeTypeBackend, Address: "10.0.0.1:50051",
|
||||
TotalVRAM: 24_000_000_000, GPUVendor: "nvidia", Capability: "nvidia-cuda-13",
|
||||
})
|
||||
|
||||
caps, err := registry.HealthyBackendCapabilities(context.Background())
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(caps).To(ConsistOf("nvidia-cuda-13"))
|
||||
})
|
||||
|
||||
It("falls back to the GPU vendor for workers that report no capability", func() {
|
||||
register(&BackendNode{
|
||||
Name: "legacy-worker", NodeType: NodeTypeBackend, Address: "10.0.0.2:50051",
|
||||
TotalVRAM: 24_000_000_000, GPUVendor: "nvidia",
|
||||
})
|
||||
|
||||
caps, err := registry.HealthyBackendCapabilities(context.Background())
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(caps).To(ConsistOf("nvidia"))
|
||||
})
|
||||
|
||||
It("reports default for a legacy worker whose VRAM is below the GPU threshold", func() {
|
||||
register(&BackendNode{
|
||||
Name: "tiny-worker", NodeType: NodeTypeBackend, Address: "10.0.0.3:50051",
|
||||
TotalVRAM: 2_000_000_000, GPUVendor: "nvidia",
|
||||
})
|
||||
|
||||
caps, err := registry.HealthyBackendCapabilities(context.Background())
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(caps).To(ConsistOf("default"))
|
||||
})
|
||||
|
||||
It("deduplicates capabilities across a homogeneous fleet", func() {
|
||||
register(&BackendNode{
|
||||
Name: "gpu-a", NodeType: NodeTypeBackend, Address: "10.0.0.4:50051",
|
||||
TotalVRAM: 24_000_000_000, GPUVendor: "nvidia", Capability: "nvidia-cuda-12",
|
||||
})
|
||||
register(&BackendNode{
|
||||
Name: "gpu-b", NodeType: NodeTypeBackend, Address: "10.0.0.5:50051",
|
||||
TotalVRAM: 24_000_000_000, GPUVendor: "nvidia", Capability: "nvidia-cuda-12",
|
||||
})
|
||||
|
||||
caps, err := registry.HealthyBackendCapabilities(context.Background())
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(caps).To(ConsistOf("nvidia-cuda-12"))
|
||||
})
|
||||
|
||||
It("collects every distinct capability in a heterogeneous fleet", func() {
|
||||
register(&BackendNode{
|
||||
Name: "nvidia-worker", NodeType: NodeTypeBackend, Address: "10.0.0.6:50051",
|
||||
TotalVRAM: 24_000_000_000, GPUVendor: "nvidia", Capability: "nvidia-cuda-13",
|
||||
})
|
||||
register(&BackendNode{
|
||||
Name: "mac-worker", NodeType: NodeTypeBackend, Address: "10.0.0.7:50051",
|
||||
TotalVRAM: 32_000_000_000, Capability: "metal",
|
||||
})
|
||||
|
||||
caps, err := registry.HealthyBackendCapabilities(context.Background())
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(caps).To(ConsistOf("nvidia-cuda-13", "metal"))
|
||||
})
|
||||
|
||||
It("ignores nodes that are not healthy backend nodes", func() {
|
||||
// A pending (unapproved) GPU worker must not advertise hardware the
|
||||
// cluster cannot schedule onto yet.
|
||||
pending := &BackendNode{
|
||||
Name: "pending-gpu", NodeType: NodeTypeBackend, Address: "10.0.0.8:50051",
|
||||
TotalVRAM: 24_000_000_000, GPUVendor: "nvidia", Capability: "nvidia-cuda-13",
|
||||
}
|
||||
Expect(registry.Register(context.Background(), pending, false)).To(Succeed())
|
||||
|
||||
register(&BackendNode{
|
||||
Name: "agent-node", NodeType: NodeTypeAgent, Address: "10.0.0.9:50051",
|
||||
Capability: "metal",
|
||||
})
|
||||
|
||||
caps, err := registry.HealthyBackendCapabilities(context.Background())
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(caps).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("returns nothing when no nodes are registered", func() {
|
||||
caps, err := registry.HealthyBackendCapabilities(context.Background())
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(caps).To(BeEmpty())
|
||||
})
|
||||
})
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/mudler/LocalAI/pkg/system"
|
||||
"github.com/mudler/LocalAI/pkg/xsysinfo"
|
||||
"github.com/mudler/xlog"
|
||||
)
|
||||
@@ -103,6 +104,16 @@ func (cfg *Config) registrationBody() map[string]any {
|
||||
// options (e.g. larger physical batch on Blackwell). Detected on the worker
|
||||
// because only the worker sees the GPU in distributed mode.
|
||||
gpuComputeCap := xsysinfo.NVIDIAComputeCapability()
|
||||
// Report our own meta-backend capability so the controller can list the
|
||||
// backends this cluster can actually run. The controller cannot infer it
|
||||
// from the GPU vendor alone: OS-dependent capabilities (metal, darwin-x86,
|
||||
// nvidia-l4t) and the CUDA runtime refinements are only observable here.
|
||||
capability := ""
|
||||
if systemState, err := system.GetSystemState(); err != nil {
|
||||
xlog.Warn("Could not detect system capability for node registration", "error", err)
|
||||
} else {
|
||||
capability = systemState.DetectedCapability()
|
||||
}
|
||||
|
||||
maxReplicas := cfg.MaxReplicasPerModel
|
||||
if maxReplicas < 1 {
|
||||
@@ -116,6 +127,7 @@ func (cfg *Config) registrationBody() map[string]any {
|
||||
"available_vram": totalVRAM, // initially all VRAM is available
|
||||
"gpu_vendor": gpuVendor,
|
||||
"gpu_compute_capability": gpuComputeCap,
|
||||
"capability": capability,
|
||||
"max_replicas_per_model": maxReplicas,
|
||||
}
|
||||
|
||||
|
||||
@@ -57,6 +57,39 @@ func init() {
|
||||
cuda12DirExists = err == nil
|
||||
}
|
||||
|
||||
// NewCapabilityState builds a SystemState that reports exactly the supplied
|
||||
// capability, bypassing all host detection.
|
||||
//
|
||||
// This exists so a controller can evaluate backend compatibility on behalf of
|
||||
// a *remote* worker. Synthesizing a SystemState from the worker's GPU vendor
|
||||
// would not work: getSystemCapabilities consults
|
||||
// LOCALAI_FORCE_META_BACKEND_CAPABILITY and /run/localai/capability first, and
|
||||
// container images routinely set the latter, so the controller's own forced
|
||||
// capability would silently override every worker's.
|
||||
func NewCapabilityState(capability string, opts ...SystemStateOptions) *SystemState {
|
||||
state := &SystemState{}
|
||||
for _, opt := range opts {
|
||||
opt(state)
|
||||
}
|
||||
state.systemCapabilities = capability
|
||||
return state
|
||||
}
|
||||
|
||||
// CapabilityFromGPU derives a coarse capability from a GPU vendor and VRAM
|
||||
// size, mirroring the tail of getSystemCapabilities.
|
||||
//
|
||||
// It is the fallback for worker nodes registered before workers began
|
||||
// reporting their own capability string: the registry only persists the
|
||||
// vendor and VRAM for those, so OS-dependent capabilities (metal, darwin-x86,
|
||||
// nvidia-l4t) and the CUDA-runtime refinements are not recoverable. Prefer the
|
||||
// worker-reported capability whenever it is present.
|
||||
func CapabilityFromGPU(gpuVendor string, vram uint64) string {
|
||||
if gpuVendor == "" || vram <= 4*1024*1024*1024 {
|
||||
return defaultCapability
|
||||
}
|
||||
return gpuVendor
|
||||
}
|
||||
|
||||
// CapabilityFilterDisabled returns true when capability-based backend filtering
|
||||
// is disabled via LOCALAI_FORCE_META_BACKEND_CAPABILITY=disable.
|
||||
func (s *SystemState) CapabilityFilterDisabled() bool {
|
||||
|
||||
58
pkg/system/node_capability_test.go
Normal file
58
pkg/system/node_capability_test.go
Normal file
@@ -0,0 +1,58 @@
|
||||
package system
|
||||
|
||||
import (
|
||||
"os"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
var _ = Describe("NewCapabilityState", func() {
|
||||
It("reports exactly the supplied capability", func() {
|
||||
Expect(NewCapabilityState("nvidia-l4t-cuda-13").DetectedCapability()).To(Equal("nvidia-l4t-cuda-13"))
|
||||
})
|
||||
|
||||
It("ignores a forced capability meant for the local host", func() {
|
||||
// A controller container image commonly pins its own capability via
|
||||
// this env var (or /run/localai/capability). That must never override
|
||||
// the capability we are evaluating on a remote worker's behalf.
|
||||
orig, had := os.LookupEnv(capabilityEnv)
|
||||
Expect(os.Setenv(capabilityEnv, "metal")).To(Succeed())
|
||||
DeferCleanup(func() {
|
||||
if had {
|
||||
Expect(os.Setenv(capabilityEnv, orig)).To(Succeed())
|
||||
return
|
||||
}
|
||||
Expect(os.Unsetenv(capabilityEnv)).To(Succeed())
|
||||
})
|
||||
|
||||
Expect(NewCapabilityState("nvidia-cuda-13").DetectedCapability()).To(Equal("nvidia-cuda-13"))
|
||||
})
|
||||
|
||||
It("applies the supplied state options", func() {
|
||||
state := NewCapabilityState("nvidia", WithBackendPath("/some/backends"))
|
||||
Expect(state.Backend.BackendsPath).To(Equal("/some/backends"))
|
||||
})
|
||||
|
||||
It("honors the disable escape hatch when that is the pinned capability", func() {
|
||||
Expect(NewCapabilityState(disableCapability).CapabilityFilterDisabled()).To(BeTrue())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("CapabilityFromGPU", func() {
|
||||
const eightGB = 8 * 1024 * 1024 * 1024
|
||||
const twoGB = 2 * 1024 * 1024 * 1024
|
||||
|
||||
It("returns the vendor for a GPU with enough VRAM", func() {
|
||||
Expect(CapabilityFromGPU(Nvidia, eightGB)).To(Equal(Nvidia))
|
||||
Expect(CapabilityFromGPU(AMD, eightGB)).To(Equal(AMD))
|
||||
})
|
||||
|
||||
It("returns default when no GPU is reported", func() {
|
||||
Expect(CapabilityFromGPU("", eightGB)).To(Equal(defaultCapability))
|
||||
})
|
||||
|
||||
It("returns default when VRAM is below the usable threshold", func() {
|
||||
Expect(CapabilityFromGPU(Nvidia, twoGB)).To(Equal(defaultCapability))
|
||||
})
|
||||
})
|
||||
Reference in New Issue
Block a user