feat(system): detect and allow to override capabilities (#5785)

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

Dockerfile

@@ -24,6 +24,8 @@ ARG TARGETARCH
 ARG TARGETVARIANT
 ENV BUILD_TYPE=${BUILD_TYPE}
 
+RUN mkdir -p /run/localai
+
 # Vulkan requirements
 RUN <<EOT bash
     if [ "${BUILD_TYPE}" = "vulkan" ] && [ "${SKIP_DRIVERS}" = "false" ]; then
@@ -36,7 +38,8 @@ RUN <<EOT bash
         apt-get install -y \
             vulkan-sdk && \
         apt-get clean && \
-        rm -rf /var/lib/apt/lists/*
+        rm -rf /var/lib/apt/lists/* && \
+        echo "vulkan" > /run/localai/capability
     fi
 EOT
 
@@ -63,7 +66,8 @@ RUN <<EOT bash
             libcusparse-dev-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} \
             libcusolver-dev-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} && \
         apt-get clean && \
-        rm -rf /var/lib/apt/lists/*
+        rm -rf /var/lib/apt/lists/* && \
+        echo "nvidia" > /run/localai/capability
     fi
 EOT
 
@@ -83,6 +87,7 @@ RUN if [ "${BUILD_TYPE}" = "hipblas" ] && [ "${SKIP_DRIVERS}" = "false" ]; then
             rocblas-dev && \
         apt-get clean && \
         rm -rf /var/lib/apt/lists/* && \
+        echo "amd" > /run/localai/capability && \
         # I have no idea why, but the ROCM lib packages don't trigger ldconfig after they install, which results in local-ai and others not being able
         # to locate the libraries. We run ldconfig ourselves to work around this packaging deficiency
         ldconfig \

core/gallery/backends.go

@@ -62,7 +62,7 @@ func findBestBackendFromMeta(backend *GalleryBackend, systemState *system.System
 		return nil
 	}
 
-	realBackend := backend.CapabilitiesMap[systemState.GPUVendor]
+	realBackend := backend.CapabilitiesMap[systemState.Capability()]
 	if realBackend == "" {
 		return nil
 	}

core/system/capabilities.go

@@ -1,6 +1,7 @@
 package system
 
 import (
+	"os"
 	"strings"
 
 	"github.com/mudler/LocalAI/pkg/xsysinfo"
@@ -11,6 +12,29 @@ type SystemState struct {
 	GPUVendor string
 }
 
+func (s *SystemState) Capability() string {
+	if os.Getenv("LOCALAI_FORCE_META_BACKEND_CAPABILITY") != "" {
+		return os.Getenv("LOCALAI_FORCE_META_BACKEND_CAPABILITY")
+	}
+
+	capabilityRunFile := "/run/localai/capability"
+	if os.Getenv("LOCALAI_FORCE_META_BACKEND_CAPABILITY_RUN_FILE") != "" {
+		capabilityRunFile = os.Getenv("LOCALAI_FORCE_META_BACKEND_CAPABILITY_RUN_FILE")
+	}
+
+	// Check if /run/localai/capability exists and use it
+	// This might be used by e.g. container images to specify which
+	// backends to pull in automatically when installing meta backends.
+	if _, err := os.Stat(capabilityRunFile); err == nil {
+		capability, err := os.ReadFile(capabilityRunFile)
+		if err == nil {
+			return string(capability)
+		}
+	}
+
+	return s.GPUVendor
+}
+
 func GetSystemState() (*SystemState, error) {
 	gpuVendor, _ := detectGPUVendor()
 	log.Debug().Str("gpuVendor", gpuVendor).Msg("GPU vendor")