+ Cap how much of this node's VRAM the scheduler may allocate. Accepts a percentage or a size (e.g. 80% or 12GB).
+
+
+
+
)
}
diff --git a/core/http/react-ui/src/pages/Settings.jsx b/core/http/react-ui/src/pages/Settings.jsx
index d455a1bde..2cf6abda2 100644
--- a/core/http/react-ui/src/pages/Settings.jsx
+++ b/core/http/react-ui/src/pages/Settings.jsx
@@ -419,6 +419,9 @@ export default function Settings() {
update('context_size', parseInt(e.target.value) || 0)} placeholder="2048" />
+
+ update('vram_budget', e.target.value)} placeholder="e.g. 80% or 12GB" />
+ update('f16', v)} />
diff --git a/core/http/react-ui/src/utils/api.js b/core/http/react-ui/src/utils/api.js
index 12bc9366b..33e200506 100644
--- a/core/http/react-ui/src/utils/api.js
+++ b/core/http/react-ui/src/utils/api.js
@@ -590,6 +590,17 @@ export const nodesApi = {
resetMaxReplicasPerModel: (id) => fetchJSON(API_CONFIG.endpoints.nodeMaxReplicasPerModel(id), {
method: 'DELETE',
}),
+ // Set a sticky admin override for the per-node VRAM allocation budget. The
+ // value is a string ("80%" or "12GB"); resolution to a byte ceiling happens
+ // server-side. Call resetVramBudget to clear the override entirely.
+ updateVramBudget: (id, value) => fetchJSON(API_CONFIG.endpoints.nodeVramBudget(id), {
+ method: 'PUT',
+ headers: { 'Content-Type': 'application/json' },
+ body: JSON.stringify({ value }),
+ }),
+ resetVramBudget: (id) => fetchJSON(API_CONFIG.endpoints.nodeVramBudget(id), {
+ method: 'DELETE',
+ }),
listScheduling: () => fetchJSON(API_CONFIG.endpoints.nodesScheduling),
allModels: () => fetchJSON(API_CONFIG.endpoints.nodesModels),
setScheduling: (config) => postJSON(API_CONFIG.endpoints.nodesScheduling, config),
diff --git a/core/http/react-ui/src/utils/config.js b/core/http/react-ui/src/utils/config.js
index a1c3c9fa8..d1bd4d335 100644
--- a/core/http/react-ui/src/utils/config.js
+++ b/core/http/react-ui/src/utils/config.js
@@ -145,6 +145,7 @@ export const API_CONFIG = {
nodeLabels: (id) => `/api/nodes/${id}/labels`,
nodeLabelKey: (id, key) => `/api/nodes/${id}/labels/${key}`,
nodeMaxReplicasPerModel: (id) => `/api/nodes/${id}/max-replicas-per-model`,
+ nodeVramBudget: (id) => `/api/nodes/${id}/vram-budget`,
nodesScheduling: '/api/nodes/scheduling',
nodesModels: '/api/nodes/models',
nodesSchedulingModel: (model) => `/api/nodes/scheduling/${encodeURIComponent(model)}`,
diff --git a/core/http/routes/nodes.go b/core/http/routes/nodes.go
index fbbcce081..053d6c19c 100644
--- a/core/http/routes/nodes.go
+++ b/core/http/routes/nodes.go
@@ -116,6 +116,12 @@ func RegisterNodeAdminRoutes(e *echo.Echo, registry *nodes.NodeRegistry, unloade
admin.PUT("/:id/max-replicas-per-model", localai.UpdateMaxReplicasPerModelEndpoint(registry))
admin.DELETE("/:id/max-replicas-per-model", localai.ResetMaxReplicasPerModelEndpoint(registry))
+ // Per-node VRAM allocation budget. PUT sets a sticky admin override that
+ // survives worker restarts; DELETE clears it so the worker's reported
+ // budget takes over again at the next re-registration.
+ admin.PUT("/:id/vram-budget", localai.UpdateVRAMBudgetEndpoint(registry))
+ admin.DELETE("/:id/vram-budget", localai.ResetVRAMBudgetEndpoint(registry))
+
// WebSocket proxy for real-time log streaming from workers
e.GET("/ws/nodes/:id/backend-logs/:modelId", localai.NodeBackendLogsWSEndpoint(registry, registrationToken), readyMw, adminMw)
}
diff --git a/core/services/nodes/registry.go b/core/services/nodes/registry.go
index 8470c67ed..d74fefa6d 100644
--- a/core/services/nodes/registry.go
+++ b/core/services/nodes/registry.go
@@ -9,6 +9,7 @@ import (
"github.com/google/uuid"
"github.com/mudler/LocalAI/core/services/advisorylock"
+ "github.com/mudler/LocalAI/pkg/vrambudget"
"github.com/mudler/xlog"
"gorm.io/gorm"
"gorm.io/gorm/clause"
@@ -50,12 +51,24 @@ type BackendNode struct {
// admin override. When true, the worker's CLI value is ignored on
// re-registration so the override survives worker restarts. Cleared
// by an explicit "reset to worker default" action.
- MaxReplicasPerModelManuallySet bool `gorm:"column:max_replicas_per_model_manually_set;default:false" json:"max_replicas_per_model_manually_set"`
- APIKeyID string `gorm:"size:36" json:"-"` // auto-provisioned API key ID (for cleanup)
- AuthUserID string `gorm:"size:36" json:"-"` // auto-provisioned user ID (for cleanup)
- LastHeartbeat time.Time `gorm:"column:last_heartbeat" json:"last_heartbeat"`
- CreatedAt time.Time `json:"created_at"`
- UpdatedAt time.Time `json:"updated_at"`
+ MaxReplicasPerModelManuallySet bool `gorm:"column:max_replicas_per_model_manually_set;default:false" json:"max_replicas_per_model_manually_set"`
+ // VRAMBudget is the operator-set allocation cap for this node ("80%" or
+ // "12GB"; empty = no cap). VRAMBudgetBytes is that budget resolved to an
+ // absolute ceiling against the node's raw TotalVRAM (0 = none) and is the
+ // value actually enforced: available_vram is written as min(reported,
+ // ceiling) so the SQL scheduler places against budgeted capacity. TotalVRAM
+ // stays raw so a percentage budget can be recomputed if capacity changes.
+ VRAMBudget string `gorm:"column:vram_budget;size:32" json:"vram_budget,omitempty"`
+ VRAMBudgetBytes uint64 `gorm:"column:vram_budget_bytes;default:0" json:"vram_budget_bytes,omitempty"`
+ // VRAMBudgetManuallySet marks the budget as a UI-set admin override so the
+ // worker's re-registration value does not clobber it (mirrors
+ // MaxReplicasPerModelManuallySet).
+ VRAMBudgetManuallySet bool `gorm:"column:vram_budget_manually_set;default:false" json:"vram_budget_manually_set"`
+ APIKeyID string `gorm:"size:36" json:"-"` // auto-provisioned API key ID (for cleanup)
+ AuthUserID string `gorm:"size:36" json:"-"` // auto-provisioned user ID (for cleanup)
+ LastHeartbeat time.Time `gorm:"column:last_heartbeat" json:"last_heartbeat"`
+ CreatedAt time.Time `json:"created_at"`
+ UpdatedAt time.Time `json:"updated_at"`
}
const (
@@ -77,6 +90,8 @@ const (
ColGPUComputeCap = "gpu_compute_capability"
ColLastHeartbeat = "last_heartbeat"
ColMaxReplicasPerModel = "max_replicas_per_model"
+ ColVRAMBudget = "vram_budget"
+ ColVRAMBudgetBytes = "vram_budget_bytes"
)
// NodeModel tracks which models are loaded on which nodes.
@@ -305,6 +320,36 @@ func NewNodeRegistry(db *gorm.DB) (*NodeRegistry, error) {
return &NodeRegistry{db: db}, nil
}
+// resolveVRAMBudgetBytes turns a budget string into an absolute byte ceiling
+// against the node's raw total VRAM, clamped to that total. 0 means no cap or
+// unparseable (fail-open at this layer; the endpoint validates input).
+func (r *NodeRegistry) resolveVRAMBudgetBytes(budget string, totalVRAM uint64) uint64 {
+ if budget == "" {
+ return 0
+ }
+ b, err := vrambudget.Parse(budget)
+ if err != nil {
+ xlog.Warn("Ignoring invalid node VRAM budget", "budget", budget, "error", err)
+ return 0
+ }
+ return b.Ceiling(totalVRAM)
+}
+
+// ResolveVRAMBudgetBytesForTest exposes resolveVRAMBudgetBytes for tests.
+func (r *NodeRegistry) ResolveVRAMBudgetBytesForTest(budget string, totalVRAM uint64) uint64 {
+ return r.resolveVRAMBudgetBytes(budget, totalVRAM)
+}
+
+// capAvailable applies a node's resolved budget ceiling to a reported available
+// figure. ceilingBytes == 0 means no cap; otherwise the result is the smaller
+// of the two so the SQL scheduler never sees more free VRAM than the budget.
+func capAvailable(reported, ceilingBytes uint64) uint64 {
+ if ceilingBytes == 0 || reported <= ceilingBytes {
+ return reported
+ }
+ return ceilingBytes
+}
+
// Register adds or updates a backend node.
// If autoApprove is true, the node goes directly to "healthy" status.
// If false, new nodes start in "pending" status and must be approved by an admin.
@@ -338,9 +383,37 @@ func (r *NodeRegistry) Register(ctx context.Context, node *BackendNode, autoAppr
node.MaxReplicasPerModel = existing.MaxReplicasPerModel
node.MaxReplicasPerModelManuallySet = true
}
+ // Preserve an admin-set VRAM budget across worker re-registration;
+ // otherwise the worker-reported budget (wired in a later task) wins.
+ if existing.VRAMBudgetManuallySet {
+ node.VRAMBudget = existing.VRAMBudget
+ node.VRAMBudgetManuallySet = true
+ }
+ // Resolve the effective budget against the (possibly updated) raw total
+ // VRAM and cap the reported available so the SQL scheduler places against
+ // budgeted capacity. TotalVRAM is written raw.
+ node.VRAMBudgetBytes = r.resolveVRAMBudgetBytes(node.VRAMBudget, node.TotalVRAM)
+ node.AvailableVRAM = capAvailable(node.AvailableVRAM, node.VRAMBudgetBytes)
if err := updateDB.Updates(node).Error; err != nil {
return fmt.Errorf("updating node %s: %w", node.Name, err)
}
+ // The struct Updates above zero-skips fields, so a worker that dropped
+ // LOCALAI_VRAM_BUDGET (now reporting an empty budget / 0 ceiling) would
+ // leave the previously-stored cap in place, so the operator's env removal
+ // would never take effect. For a node whose budget is NOT an admin
+ // override, force-write the worker-authoritative budget columns even
+ // when empty/zero so removing the budget actually clears the cap.
+ // Manual overrides are handled above and keep the admin value untouched.
+ if !existing.VRAMBudgetManuallySet {
+ if err := r.db.WithContext(ctx).Model(&BackendNode{}).Where("id = ?", node.ID).
+ Updates(map[string]any{
+ ColVRAMBudget: node.VRAMBudget,
+ ColVRAMBudgetBytes: node.VRAMBudgetBytes,
+ ColAvailableVRAM: node.AvailableVRAM,
+ }).Error; err != nil {
+ return fmt.Errorf("clearing worker VRAM budget for node %s: %w", node.Name, err)
+ }
+ }
// Preserve auth references from existing record.
// GORM Updates(struct) skips zero-value fields, so the DB retains
// the old auth_user_id/api_key_id but the caller's struct is empty.
@@ -380,6 +453,10 @@ func (r *NodeRegistry) Register(ctx context.Context, node *BackendNode, autoAppr
} else {
node.Status = StatusPending
}
+ // Resolve a worker-reported budget (wired in a later task) against the
+ // node's raw total VRAM and cap the reported available accordingly.
+ node.VRAMBudgetBytes = r.resolveVRAMBudgetBytes(node.VRAMBudget, node.TotalVRAM)
+ node.AvailableVRAM = capAvailable(node.AvailableVRAM, node.VRAMBudgetBytes)
if err := r.db.WithContext(ctx).Create(node).Error; err != nil {
return fmt.Errorf("creating node %s: %w", node.Name, err)
}
@@ -617,7 +694,13 @@ func (r *NodeRegistry) Heartbeat(ctx context.Context, nodeID string, update *Hea
if update != nil {
if update.AvailableVRAM != nil {
- updates[ColAvailableVRAM] = *update.AvailableVRAM
+ // Cap the reported available against the node's resolved budget
+ // ceiling (0 = none) so the SQL scheduler only ever sees budgeted
+ // capacity. TotalVRAM stays raw (written below).
+ var ceiling uint64
+ db.Model(&BackendNode{}).
+ Select(ColVRAMBudgetBytes).Where("id = ?", nodeID).Scan(&ceiling)
+ updates[ColAvailableVRAM] = capAvailable(*update.AvailableVRAM, ceiling)
// The worker is the source of truth for actual free VRAM.
// Whenever it sends us a fresh reading, the in-tick soft
// reservation is no longer needed — clear it. (See ReserveVRAM.)
@@ -1670,6 +1753,57 @@ func (r *NodeRegistry) ResetMaxReplicasPerModel(ctx context.Context, nodeID stri
return nil
}
+// UpdateVRAMBudget sets a node's VRAM allocation cap as a sticky admin override.
+// It resolves the budget against the node's raw TotalVRAM, stores the ceiling,
+// and immediately re-caps available_vram so the scheduler reflects the change
+// before the next heartbeat. Empty budget clears the cap. The override survives
+// worker re-registration (see Register); to hand control back to the worker,
+// call ResetVRAMBudget.
+func (r *NodeRegistry) UpdateVRAMBudget(ctx context.Context, nodeID, budget string) error {
+ node, err := r.Get(ctx, nodeID)
+ if err != nil {
+ return err
+ }
+ ceiling := r.resolveVRAMBudgetBytes(budget, node.TotalVRAM)
+ res := r.db.WithContext(ctx).Model(&BackendNode{}).Where("id = ?", nodeID).
+ Updates(map[string]any{
+ ColVRAMBudget: budget,
+ ColVRAMBudgetBytes: ceiling,
+ "vram_budget_manually_set": true,
+ ColAvailableVRAM: capAvailable(node.AvailableVRAM, ceiling),
+ })
+ if res.Error != nil {
+ return fmt.Errorf("updating vram_budget on %s: %w", nodeID, res.Error)
+ }
+ if res.RowsAffected == 0 {
+ return fmt.Errorf("node %s not found", nodeID)
+ }
+ // Capping available_vram may have freed or constrained capacity; wake any
+ // configs the reconciler put in cooldown so the next tick re-evaluates.
+ if err := r.ClearAllUnsatisfiable(ctx); err != nil {
+ xlog.Warn("Failed to clear unsatisfiable flags after vram budget change", "error", err)
+ }
+ return nil
+}
+
+// ResetVRAMBudget clears the admin override and the resolved ceiling, handing
+// budget control back to the worker's reported budget on next register.
+func (r *NodeRegistry) ResetVRAMBudget(ctx context.Context, nodeID string) error {
+ res := r.db.WithContext(ctx).Model(&BackendNode{}).Where("id = ?", nodeID).
+ Updates(map[string]any{
+ ColVRAMBudget: "",
+ ColVRAMBudgetBytes: 0,
+ "vram_budget_manually_set": false,
+ })
+ if res.Error != nil {
+ return fmt.Errorf("clearing vram_budget on %s: %w", nodeID, res.Error)
+ }
+ if res.RowsAffected == 0 {
+ return fmt.Errorf("node %s not found", nodeID)
+ }
+ return nil
+}
+
// ClearAllUnsatisfiable clears the cooldown flag on every scheduling config.
// Called from cluster-events that could plausibly increase capacity (new
// node registers, node approves pending→healthy, node labels change,
diff --git a/core/services/nodes/registry_vrambudget_test.go b/core/services/nodes/registry_vrambudget_test.go
new file mode 100644
index 000000000..0360c345e
--- /dev/null
+++ b/core/services/nodes/registry_vrambudget_test.go
@@ -0,0 +1,169 @@
+package nodes
+
+import (
+ "context"
+ "runtime"
+
+ . "github.com/onsi/ginkgo/v2"
+ . "github.com/onsi/gomega"
+
+ "github.com/mudler/LocalAI/core/services/testutil"
+ "gorm.io/gorm"
+)
+
+var _ = Describe("Node VRAM budget", func() {
+ var (
+ db *gorm.DB
+ registry *NodeRegistry
+ )
+
+ // gb is 1000-based to match vrambudget's decimal "GB" suffix.
+ const gb = uint64(1000 * 1000 * 1000)
+
+ 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())
+ })
+
+ // seedHealthyNode registers an auto-approved backend node with raw total and
+ // available VRAM and returns its ID.
+ seedHealthyNode := func(ctx context.Context, name string, total, avail uint64) string {
+ node := &BackendNode{
+ Name: name,
+ NodeType: NodeTypeBackend,
+ Address: "10.0.0.1:50051",
+ TotalVRAM: total,
+ AvailableVRAM: avail,
+ }
+ Expect(registry.Register(ctx, node, true)).To(Succeed())
+ return node.ID
+ }
+
+ It("resolves a percentage budget against the node's raw total VRAM", func() {
+ Expect(registry.ResolveVRAMBudgetBytesForTest("80%", 10*gb)).To(Equal(8 * gb))
+ Expect(registry.ResolveVRAMBudgetBytesForTest("12GB", 10*gb)).To(Equal(10 * gb)) // clamped to physical
+ Expect(registry.ResolveVRAMBudgetBytesForTest("", 10*gb)).To(Equal(uint64(0)))
+ })
+
+ It("caps stored available_vram when an admin sets a budget", func(ctx SpecContext) {
+ id := seedHealthyNode(ctx, "worker-budget-1", 16*gb, 16*gb)
+ Expect(registry.UpdateVRAMBudget(ctx, id, "50%")).To(Succeed())
+
+ node, err := registry.Get(ctx, id)
+ Expect(err).NotTo(HaveOccurred())
+ Expect(node.TotalVRAM).To(Equal(16 * gb)) // raw preserved
+ Expect(node.VRAMBudgetBytes).To(Equal(8 * gb))
+ Expect(node.AvailableVRAM).To(Equal(8 * gb)) // capped
+ Expect(node.VRAMBudgetManuallySet).To(BeTrue())
+ })
+
+ It("re-caps available_vram on heartbeat against the stored ceiling", func(ctx SpecContext) {
+ id := seedHealthyNode(ctx, "worker-budget-2", 16*gb, 16*gb)
+ Expect(registry.UpdateVRAMBudget(ctx, id, "8GB")).To(Succeed())
+
+ avail := 15 * gb
+ Expect(registry.Heartbeat(ctx, id, &HeartbeatUpdate{AvailableVRAM: &avail})).To(Succeed())
+ node, err := registry.Get(ctx, id)
+ Expect(err).NotTo(HaveOccurred())
+ Expect(node.AvailableVRAM).To(Equal(8 * gb)) // reported 15GB capped to 8GB budget
+ })
+
+ It("preserves an admin override across worker re-registration", func(ctx SpecContext) {
+ id := seedHealthyNode(ctx, "worker-budget-3", 16*gb, 16*gb)
+ Expect(registry.UpdateVRAMBudget(ctx, id, "50%")).To(Succeed())
+
+ // Worker re-registers reporting full available VRAM and no budget.
+ reReg := &BackendNode{
+ Name: "worker-budget-3",
+ NodeType: NodeTypeBackend,
+ Address: "10.0.0.1:50051",
+ TotalVRAM: 16 * gb,
+ AvailableVRAM: 16 * gb,
+ }
+ Expect(registry.Register(ctx, reReg, true)).To(Succeed())
+
+ node, err := registry.Get(ctx, id)
+ Expect(err).NotTo(HaveOccurred())
+ Expect(node.VRAMBudgetManuallySet).To(BeTrue())
+ Expect(node.VRAMBudget).To(Equal("50%"))
+ Expect(node.VRAMBudgetBytes).To(Equal(8 * gb))
+ Expect(node.AvailableVRAM).To(Equal(8 * gb)) // re-capped despite worker reporting 16GB
+ })
+
+ It("clears the cap when the budget is reset", func(ctx SpecContext) {
+ id := seedHealthyNode(ctx, "worker-budget-4", 16*gb, 16*gb)
+ Expect(registry.UpdateVRAMBudget(ctx, id, "50%")).To(Succeed())
+ Expect(registry.ResetVRAMBudget(ctx, id)).To(Succeed())
+ node, err := registry.Get(ctx, id)
+ Expect(err).NotTo(HaveOccurred())
+ Expect(node.VRAMBudgetBytes).To(Equal(uint64(0)))
+ Expect(node.VRAMBudgetManuallySet).To(BeFalse())
+ })
+
+ // reregister re-registers an existing node by name, mirroring what the worker
+ // does on restart (raw total/available, and whatever budget it currently
+ // reports via LOCALAI_VRAM_BUDGET).
+ reregister := func(ctx context.Context, name, budget string, total, avail uint64) {
+ reReg := &BackendNode{
+ Name: name,
+ NodeType: NodeTypeBackend,
+ Address: "10.0.0.1:50051",
+ TotalVRAM: total,
+ AvailableVRAM: avail,
+ VRAMBudget: budget,
+ }
+ Expect(registry.Register(ctx, reReg, true)).To(Succeed())
+ }
+
+ It("clears a worker-reported budget when the worker re-registers without one", func(ctx SpecContext) {
+ // Worker first reports LOCALAI_VRAM_BUDGET=80%.
+ node := &BackendNode{
+ Name: "worker-clear-1",
+ NodeType: NodeTypeBackend,
+ Address: "10.0.0.1:50051",
+ TotalVRAM: 16 * gb,
+ AvailableVRAM: 16 * gb,
+ VRAMBudget: "80%",
+ }
+ Expect(registry.Register(ctx, node, true)).To(Succeed())
+
+ got, err := registry.Get(ctx, node.ID)
+ Expect(err).NotTo(HaveOccurred())
+ Expect(got.VRAMBudget).To(Equal("80%"))
+ Expect(got.VRAMBudgetBytes).To(BeNumerically(">", uint64(0)))
+ Expect(got.VRAMBudgetBytes).To(Equal(uint64(12.8 * float64(gb)))) // 80% of 16GB
+ Expect(got.AvailableVRAM).To(Equal(got.VRAMBudgetBytes)) // capped
+
+ // Operator removes LOCALAI_VRAM_BUDGET and restarts the worker: it now
+ // re-registers with an empty budget and full raw available. The stale
+ // 80% cap MUST be cleared, not preserved (struct Updates zero-skip bug).
+ reregister(ctx, "worker-clear-1", "", 16*gb, 16*gb)
+
+ got, err = registry.Get(ctx, node.ID)
+ Expect(err).NotTo(HaveOccurred())
+ Expect(got.VRAMBudget).To(Equal(""))
+ Expect(got.VRAMBudgetBytes).To(Equal(uint64(0)))
+ Expect(got.AvailableVRAM).To(Equal(16 * gb)) // back to raw reported
+ })
+
+ It("keeps an admin override when a worker re-registers without a budget", func(ctx SpecContext) {
+ // Regression guard for the force-clear: it must not clobber a sticky
+ // admin-set budget just because the worker stopped reporting one.
+ id := seedHealthyNode(ctx, "worker-clear-2", 16*gb, 16*gb)
+ Expect(registry.UpdateVRAMBudget(ctx, id, "50%")).To(Succeed())
+
+ reregister(ctx, "worker-clear-2", "", 16*gb, 16*gb)
+
+ node, err := registry.Get(ctx, id)
+ Expect(err).NotTo(HaveOccurred())
+ Expect(node.VRAMBudgetManuallySet).To(BeTrue())
+ Expect(node.VRAMBudget).To(Equal("50%"))
+ Expect(node.VRAMBudgetBytes).To(Equal(8 * gb))
+ Expect(node.AvailableVRAM).To(Equal(8 * gb)) // still capped by admin budget
+ })
+})
diff --git a/core/services/nodes/router.go b/core/services/nodes/router.go
index be88a60f4..011001234 100644
--- a/core/services/nodes/router.go
+++ b/core/services/nodes/router.go
@@ -181,10 +181,20 @@ func applyNodeHardwareDefaults(opts *pb.ModelOptions, node *BackendNode, backend
if opts == nil || node == nil || config.HardwareDefaultsDisabled() {
return
}
+ // Gate the throughput heuristics on the node's BUDGETED ceiling, not its
+ // physical card. node.TotalVRAM is stored raw (so a percentage budget can be
+ // recomputed if capacity changes), but the batch/parallel boosts allocate a
+ // per-device compute buffer against real capacity: an operator who budgeted
+ // only a slice of the GPU to LocalAI must not get defaults sized for the full
+ // device, or the raised batch/slots would overflow the allocation (#10485).
+ usableVRAM := node.TotalVRAM
+ if node.VRAMBudgetBytes > 0 && node.VRAMBudgetBytes < usableVRAM {
+ usableVRAM = node.VRAMBudgetBytes
+ }
gpu := config.GPU{
Vendor: node.GPUVendor,
ComputeCapability: node.GPUComputeCapability,
- VRAM: node.TotalVRAM,
+ VRAM: usableVRAM,
}
if config.IsManagedPhysicalBatch(int(opts.NBatch)) {
// Gate the raised batch on the selected node's per-device VRAM at this
diff --git a/core/services/nodes/router_vrambudget_test.go b/core/services/nodes/router_vrambudget_test.go
new file mode 100644
index 000000000..cb16c30f3
--- /dev/null
+++ b/core/services/nodes/router_vrambudget_test.go
@@ -0,0 +1,40 @@
+package nodes
+
+import (
+ . "github.com/onsi/ginkgo/v2"
+ . "github.com/onsi/gomega"
+
+ "github.com/mudler/LocalAI/core/config"
+ pb "github.com/mudler/LocalAI/pkg/grpc/proto"
+)
+
+var _ = Describe("applyNodeHardwareDefaults with VRAM budget", func() {
+ const gb = uint64(1000 * 1000 * 1000)
+
+ It("uses the capped ceiling, not raw TotalVRAM, for batch/parallel gating", func() {
+ // A Blackwell node whose RAW VRAM has ample headroom to keep the raised
+ // physical batch, but whose operator budget is far too small for it. The
+ // batch/parallel heuristics must gate on the budgeted ceiling, not the
+ // physical card, or a budgeted-tiny node would still get OOM-prone
+ // throughput defaults meant for the full device.
+ //
+ // The context is deliberately large: the compute-buffer headroom guard
+ // (PhysicalBatchForContext) scales the extra scratch by n_ctx, so at a
+ // small context even 2GB clears the guard. At 32768 the raised batch's
+ // scratch fits comfortably in 64GB (raw keeps 2048) but overflows a
+ // quarter of a 2GB budget (budget drops it to the conservative default).
+ const largeCtx = int32(32768)
+ node := &BackendNode{
+ GPUVendor: "nvidia",
+ GPUComputeCapability: "12.1",
+ TotalVRAM: 64 * gb,
+ VRAMBudgetBytes: 2 * gb, // tiny operator budget
+ }
+ opts := &pb.ModelOptions{NBatch: int32(config.BlackwellPhysicalBatch), ContextSize: largeCtx}
+
+ applyNodeHardwareDefaults(opts, node, "llama-cpp")
+
+ // With only 2GB budgeted the raised 2048 batch must not survive.
+ Expect(int(opts.NBatch)).To(BeNumerically("<", config.BlackwellPhysicalBatch))
+ })
+})
diff --git a/core/services/worker/config.go b/core/services/worker/config.go
index 049189f0a..817f7f217 100644
--- a/core/services/worker/config.go
+++ b/core/services/worker/config.go
@@ -61,6 +61,11 @@ type Config struct {
// the existing label selector.
MaxReplicasPerModel int `env:"LOCALAI_MAX_REPLICAS_PER_MODEL" default:"1" help:"Max replicas of any single model on this worker. Default 1 preserves single-replica behavior; set higher to allow stacking replicas on a fat node." group:"registration"`
+ // VRAMBudget optionally caps this node's VRAM for model allocation ("80%" or
+ // "12GB"). Reported to the server as a string; the server resolves and
+ // enforces it against the raw VRAM this worker reports. Empty = no cap.
+ VRAMBudget string `env:"LOCALAI_VRAM_BUDGET" help:"Cap VRAM used for model allocation on this worker node, as a percentage (e.g. 80%) or absolute amount (e.g. 12GB)." group:"registration"`
+
// NATS (required)
NatsURL string `env:"LOCALAI_NATS_URL" required:"" help:"NATS server URL" group:"distributed"`
NatsJWT string `env:"LOCALAI_NATS_JWT" help:"NATS user JWT override (normally from registration nats_jwt)" group:"distributed"`
diff --git a/core/services/worker/registration.go b/core/services/worker/registration.go
index 432cc845b..7d3bbf244 100644
--- a/core/services/worker/registration.go
+++ b/core/services/worker/registration.go
@@ -93,6 +93,14 @@ func (cfg *Config) registrationBody() map[string]any {
"max_replicas_per_model": maxReplicas,
}
+ // Report the operator-set budget as a STRING so the server resolves and
+ // enforces it against the raw VRAM above. The worker never caps its own
+ // total_vram/available_vram, and never touches the xsysinfo process-global
+ // budget (that is standalone-only). Omit when unset.
+ if cfg.VRAMBudget != "" {
+ body["vram_budget"] = cfg.VRAMBudget
+ }
+
// If no GPU detected, report system RAM so the scheduler/UI has capacity info
if totalVRAM == 0 {
if ramInfo, err := xsysinfo.GetSystemRAMInfo(); err == nil {
diff --git a/core/services/worker/registration_test.go b/core/services/worker/registration_test.go
new file mode 100644
index 000000000..0fe6d2e51
--- /dev/null
+++ b/core/services/worker/registration_test.go
@@ -0,0 +1,33 @@
+package worker
+
+import (
+ . "github.com/onsi/ginkgo/v2"
+ . "github.com/onsi/gomega"
+
+ "github.com/mudler/LocalAI/pkg/xsysinfo"
+)
+
+var _ = Describe("Worker registration body", func() {
+ It("includes the VRAM budget in the registration body when set", func() {
+ cfg := &Config{VRAMBudget: "80%"}
+ body := cfg.registrationBody()
+ Expect(body["vram_budget"]).To(Equal("80%"))
+ })
+
+ It("omits an empty VRAM budget", func() {
+ cfg := &Config{}
+ body := cfg.registrationBody()
+ _, present := body["vram_budget"]
+ Expect(present).To(BeFalse())
+ })
+
+ It("reports raw total_vram regardless of the budget", func() {
+ // The worker reports RAW VRAM; the server resolves and enforces the
+ // budget. It must never pre-cap total_vram locally, so total_vram equals
+ // the raw xsysinfo reading whether or not a budget is set.
+ rawTotal, _ := xsysinfo.TotalAvailableVRAM()
+ cfg := &Config{VRAMBudget: "50%"}
+ body := cfg.registrationBody()
+ Expect(body["total_vram"]).To(Equal(rawTotal))
+ })
+})
diff --git a/docs/content/advanced/vram-management.md b/docs/content/advanced/vram-management.md
index ffa08b894..bb7583a07 100644
--- a/docs/content/advanced/vram-management.md
+++ b/docs/content/advanced/vram-management.md
@@ -360,6 +360,49 @@ This configuration:
- Waits 2 seconds between retry attempts
- Ensures busy models have time to complete their requests before eviction
+## VRAM Budget (Allocation Ceiling)
+
+By default LocalAI treats **all** detected GPU memory as available for model allocation. On a shared machine (a desktop you also game on, a box that runs other GPU workloads, or a multi-tenant server) you often want to reserve some headroom and let LocalAI use only part of the card. The **VRAM budget** sets a hard ceiling on how much VRAM LocalAI will consider allocatable.
+
+### Configuration
+
+Set the budget with a CLI flag or environment variable:
+
+```bash
+# Percentage of detected VRAM
+./local-ai --vram-budget=80%
+
+# Absolute amount (GB, GiB, MB, or raw bytes)
+./local-ai --vram-budget=12GB
+
+# Using environment variables
+LOCALAI_VRAM_BUDGET=80% ./local-ai
+LOCALAI_VRAM_BUDGET=12GB ./local-ai
+```
+
+Accepted formats:
+
+- A **percentage** such as `80%` (or the decimal `0.8`).
+- An **absolute amount** such as `12GB`, `12GiB`, `12000MB`, or a raw byte count.
+
+Leaving it empty or unset (the default) uses all detected VRAM, which is the historical behavior.
+
+### Semantics
+
+The budget is a **hard ceiling**, never a floor:
+
+- Everywhere LocalAI reads VRAM to decide model allocation (hardware defaults for batch size and parallelism, context auto-fit, GGUF fit warnings, and the watchdog), it uses `min(detected VRAM, budget)`.
+- The budget can only **lower** usable VRAM, never raise it above what the hardware physically has. An absolute value larger than physical VRAM is clamped to physical.
+- A percentage above `100%` is rejected as invalid.
+
+### Editing from the UI
+
+The budget is exposed in the **Settings** page under the **Performance** section as **VRAM Budget**, so you can change it without editing flags or restarting from scratch. It is live-editable and persists to the runtime settings.
+
+### Distributed mode
+
+The same `LOCALAI_VRAM_BUDGET` / `--vram-budget` applies to distributed worker nodes (`local-ai worker ...`), where it caps the VRAM the scheduler will use for per-node placement. Admins can also set a per-node budget live from the node UI (or the admin API); see [Distributed Mode]({{%relref "features/distributed-mode#per-node-vram-budget" %}}).
+
## Limitations and Considerations
### VRAM Usage Estimation
diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md
index adf0e3292..6e2af94bd 100644
--- a/docs/content/features/distributed-mode.md
+++ b/docs/content/features/distributed-mode.md
@@ -220,6 +220,7 @@ local-ai worker \
| `--nats-tls-key` | `LOCALAI_NATS_TLS_KEY` | *(empty)* | Client private key for NATS mTLS |
| `--backends-path` | `LOCALAI_BACKENDS_PATH` | `./backends` | Path to backend binaries |
| `--models-path` | `LOCALAI_MODELS_PATH` | `./models` | Path to model files |
+| `--vram-budget` | `LOCALAI_VRAM_BUDGET` | *(empty)* | Cap the VRAM this node advertises for model placement, as a percentage (e.g. `80%`) or an absolute amount (e.g. `12GB`). Empty uses all detected VRAM. See [Per-node VRAM budget](#per-node-vram-budget). |
{{% notice tip %}}
**Advertise address:** The `--addr` flag is the local bind address for gRPC. The `--advertise-addr` is the address the frontend stores and uses to reach the worker via gRPC. If not set, the worker auto-derives it by replacing `0.0.0.0` with the OS hostname (which in Docker is the container ID, resolvable via Docker DNS). Set `--advertise-addr` explicitly when the auto-detected hostname is not routable from the frontend (e.g., in Kubernetes, use the pod's service DNS name).
@@ -337,9 +338,41 @@ Used by the WebUI and admin API consumers. Requires admin authentication.
| `POST` | `/api/nodes/:id/backends/delete` | Delete a backend from a worker |
| `POST` | `/api/nodes/:id/models/unload` | Unload a model from a worker |
| `POST` | `/api/nodes/:id/models/delete` | Delete model files from a worker |
+| `PUT` | `/api/nodes/:id/vram-budget` | Set a VRAM budget for a worker (`{"value":"80%"}`) |
+| `DELETE` | `/api/nodes/:id/vram-budget` | Clear a worker's VRAM budget (revert to all detected VRAM) |
The **Nodes** page in the React WebUI provides a visual overview of all registered workers, their statuses, and loaded models. The page opens with a one-line **cluster pulse** summarising node health and an **attention callout** that surfaces nodes needing action (for example pending approvals). Below that, a roster of **node panels** lists each worker with its inline model chips (no expand click needed), filtered by an **All / Backend / Agent** segmented control. Selecting a panel opens a dedicated **node detail page** at `/app/nodes/:id` with per-node metrics, models, and backend actions. Model scheduling lives on its own **Scheduling** page (separate nav item), not as a tab on the Nodes page.
+### Per-node VRAM budget
+
+Each worker advertises its detected VRAM, and the SmartRouter uses that number when picking a node with enough free memory. You can cap the VRAM a node offers for placement so it never gets scheduled beyond a chosen limit, leaving headroom for other workloads on that machine.
+
+There are two ways to set the cap:
+
+- **At the worker:** start it with `--vram-budget` / `LOCALAI_VRAM_BUDGET` (see [Worker Configuration](#worker-configuration)).
+- **From the frontend, live:** set it per node in the **node capacity editor** on the node detail page, or via the admin API:
+
+```bash
+# Cap node placement at 80% of its detected VRAM
+curl -X PUT http://frontend:8080/api/nodes//vram-budget \
+ -H "Authorization: Bearer " \
+ -H "Content-Type: application/json" \
+ -d '{"value":"80%"}'
+
+# Or an absolute amount
+curl -X PUT http://frontend:8080/api/nodes//vram-budget \
+ -H "Authorization: Bearer " \
+ -d '{"value":"12GB"}'
+
+# Clear the budget (revert to all detected VRAM)
+curl -X DELETE http://frontend:8080/api/nodes//vram-budget \
+ -H "Authorization: Bearer "
+```
+
+The value accepts the same formats as the standalone budget: a percentage (`80%`) or an absolute amount (`12GB`, `12GiB`, `12000MB`, or raw bytes). It is a **hard ceiling**: the node's advertised VRAM becomes `min(detected, budget)`, so a budget can only lower the number, never raise it above the hardware. An admin-set node budget is **sticky across worker restarts**: it is stored in the node registry and reapplied when the worker re-registers, so it wins over whatever the worker reports on reconnect. For the underlying semantics and the standalone equivalent, see [VRAM Budget]({{%relref "advanced/vram-management#vram-budget-allocation-ceiling" %}}).
+
+The **LocalAI Assistant** can also set a node budget conversationally through the `set_node_vram_budget` MCP tool.
+
## Node Approval
By default, new worker nodes start in **pending** status and must be approved by an admin before they can receive traffic. This prevents unknown machines from joining the cluster.
diff --git a/pkg/mcp/localaitools/client.go b/pkg/mcp/localaitools/client.go
index d52469701..9d76bee58 100644
--- a/pkg/mcp/localaitools/client.go
+++ b/pkg/mcp/localaitools/client.go
@@ -64,6 +64,10 @@ type LocalAIClient interface {
// ---- System ----
SystemInfo(ctx context.Context) (*SystemInfo, error)
ListNodes(ctx context.Context) ([]Node, error)
+ // SetNodeVRAMBudget sets (or, with an empty budget, clears) a federated
+ // node's VRAM allocation cap as a sticky admin override. Only meaningful
+ // in distributed mode; single-process clients report it as unavailable.
+ SetNodeVRAMBudget(ctx context.Context, nodeID, budget string) error
VRAMEstimate(ctx context.Context, req VRAMEstimateRequest) (*vram.EstimateResult, error)
// ---- State ----
diff --git a/pkg/mcp/localaitools/coverage_test.go b/pkg/mcp/localaitools/coverage_test.go
index 76d8b2485..d068180aa 100644
--- a/pkg/mcp/localaitools/coverage_test.go
+++ b/pkg/mcp/localaitools/coverage_test.go
@@ -59,6 +59,7 @@ var toolToHTTPRoute = map[string]string{
ToolSetAlias: "PATCH /api/models/config-json/:name (swap) or POST /models/import (create)",
ToolCreateVoiceProfile: "POST /api/voice-profiles",
ToolDeleteVoiceProfile: "DELETE /api/voice-profiles/:id",
+ ToolSetNodeVRAMBudget: "PUT /api/nodes/:id/vram-budget",
}
// allKnownTools is the union of expectedFullCatalog (defined in
diff --git a/pkg/mcp/localaitools/dto.go b/pkg/mcp/localaitools/dto.go
index 65ecea667..ecc23204a 100644
--- a/pkg/mcp/localaitools/dto.go
+++ b/pkg/mcp/localaitools/dto.go
@@ -105,6 +105,15 @@ type Node struct {
LastSeen string `json:"last_seen,omitempty"`
}
+// SetNodeVRAMBudgetRequest is the input for set_node_vram_budget. It PUTs
+// the value to /api/nodes/{node_id}/vram-budget, where the server validates
+// and resolves the budget against the node's total VRAM. An empty Budget
+// clears the admin override so the worker's own default takes over again.
+type SetNodeVRAMBudgetRequest struct {
+ NodeID string `json:"node_id" jsonschema:"The federated node id (from list_nodes) whose VRAM budget to set."`
+ Budget string `json:"budget,omitempty" jsonschema:"VRAM allocation cap as a percentage (e.g. 80%) or absolute amount (e.g. 12GB). Empty string clears the override."`
+}
+
// ImportModelURIRequest is the input for import_model_uri. It mirrors the
// REST surface (`/models/import-uri`) closely so both clients can produce
// identical responses; the BackendPreference is a flat field rather than the
diff --git a/pkg/mcp/localaitools/fakes_test.go b/pkg/mcp/localaitools/fakes_test.go
index 2f298bb4e..ed328352c 100644
--- a/pkg/mcp/localaitools/fakes_test.go
+++ b/pkg/mcp/localaitools/fakes_test.go
@@ -42,6 +42,7 @@ type fakeClient struct {
upgradeBackend func(string) (string, error)
systemInfo func() (*SystemInfo, error)
listNodes func() ([]Node, error)
+ setNodeVRAMBudget func(string, string) error
vramEstimate func(VRAMEstimateRequest) (*vram.EstimateResult, error)
toggleModelState func(string, modeladmin.Action) error
toggleModelPinned func(string, modeladmin.Action) error
@@ -229,6 +230,14 @@ func (f *fakeClient) ListNodes(_ context.Context) ([]Node, error) {
return nil, nil
}
+func (f *fakeClient) SetNodeVRAMBudget(_ context.Context, nodeID, budget string) error {
+ f.record("SetNodeVRAMBudget", []any{nodeID, budget})
+ if f.setNodeVRAMBudget != nil {
+ return f.setNodeVRAMBudget(nodeID, budget)
+ }
+ return nil
+}
+
func (f *fakeClient) VRAMEstimate(_ context.Context, req VRAMEstimateRequest) (*vram.EstimateResult, error) {
f.record("VRAMEstimate", req)
if f.vramEstimate != nil {
diff --git a/pkg/mcp/localaitools/httpapi/client.go b/pkg/mcp/localaitools/httpapi/client.go
index da94e40a4..771d2908a 100644
--- a/pkg/mcp/localaitools/httpapi/client.go
+++ b/pkg/mcp/localaitools/httpapi/client.go
@@ -480,6 +480,13 @@ func (c *Client) ListNodes(ctx context.Context) ([]localaitools.Node, error) {
return out, nil
}
+func (c *Client) SetNodeVRAMBudget(ctx context.Context, nodeID, budget string) error {
+ // PUT with an empty value clears the override server-side (Task 9), so we
+ // use PUT uniformly rather than switching to DELETE for the clear case.
+ body := map[string]any{"value": budget}
+ return c.do(ctx, http.MethodPut, routeNodeVRAMBudget(nodeID), body, nil)
+}
+
func (c *Client) VRAMEstimate(ctx context.Context, req localaitools.VRAMEstimateRequest) (*vram.EstimateResult, error) {
body := map[string]any{"model": req.ModelName}
if req.ContextSize > 0 {
diff --git a/pkg/mcp/localaitools/httpapi/routes.go b/pkg/mcp/localaitools/httpapi/routes.go
index 81f61e13f..2b8a107dd 100644
--- a/pkg/mcp/localaitools/httpapi/routes.go
+++ b/pkg/mcp/localaitools/httpapi/routes.go
@@ -62,3 +62,7 @@ func routeToggleModelPinned(name, action string) string {
func routeVoiceProfileDelete(id string) string {
return "/api/voice-profiles/" + url.PathEscape(id)
}
+
+func routeNodeVRAMBudget(id string) string {
+ return "/api/nodes/" + url.PathEscape(id) + "/vram-budget"
+}
diff --git a/pkg/mcp/localaitools/inproc/client.go b/pkg/mcp/localaitools/inproc/client.go
index d30a5b7a1..9f6eb8fe2 100644
--- a/pkg/mcp/localaitools/inproc/client.go
+++ b/pkg/mcp/localaitools/inproc/client.go
@@ -500,6 +500,14 @@ func (c *Client) ListNodes(_ context.Context) ([]localaitools.Node, error) {
return []localaitools.Node{}, nil
}
+func (c *Client) SetNodeVRAMBudget(_ context.Context, _, _ string) error {
+ // The node registry is a distributed-mode concern owned by the Application
+ // layer and is not wired into the in-process client (which also returns an
+ // empty node list from ListNodes). Report it as unavailable rather than
+ // silently succeeding so the assistant tells the admin the truth.
+ return errors.New("per-node VRAM budgets are only available in distributed mode")
+}
+
func (c *Client) VRAMEstimate(ctx context.Context, req localaitools.VRAMEstimateRequest) (*vram.EstimateResult, error) {
resp, err := modeladmin.EstimateVRAM(ctx, modeladmin.VRAMRequest{
Model: req.ModelName,
diff --git a/pkg/mcp/localaitools/server_test.go b/pkg/mcp/localaitools/server_test.go
index 0d24f3b5e..a47b84ba3 100644
--- a/pkg/mcp/localaitools/server_test.go
+++ b/pkg/mcp/localaitools/server_test.go
@@ -104,6 +104,7 @@ var expectedFullCatalog = sortedStrings(
ToolVRAMEstimate,
ToolCreateVoiceProfile,
ToolDeleteVoiceProfile,
+ ToolSetNodeVRAMBudget,
)
// expectedReadOnlyCatalog is the tool set when DisableMutating=true. Sorted.
diff --git a/pkg/mcp/localaitools/tools.go b/pkg/mcp/localaitools/tools.go
index a09f47614..0d6890e34 100644
--- a/pkg/mcp/localaitools/tools.go
+++ b/pkg/mcp/localaitools/tools.go
@@ -41,6 +41,7 @@ const (
ToolSetAlias = "set_alias"
ToolCreateVoiceProfile = "create_voice_profile"
ToolDeleteVoiceProfile = "delete_voice_profile"
+ ToolSetNodeVRAMBudget = "set_node_vram_budget"
// ToolListAliases is read-only but lives here so the alias tools stay
// grouped; the catalog tests assert its read-only placement.
diff --git a/pkg/mcp/localaitools/tools_system.go b/pkg/mcp/localaitools/tools_system.go
index cf1c83399..b88ffbe90 100644
--- a/pkg/mcp/localaitools/tools_system.go
+++ b/pkg/mcp/localaitools/tools_system.go
@@ -6,7 +6,7 @@ import (
"github.com/modelcontextprotocol/go-sdk/mcp"
)
-func registerSystemTools(s *mcp.Server, client LocalAIClient, _ Options) {
+func registerSystemTools(s *mcp.Server, client LocalAIClient, opts Options) {
mcp.AddTool(s, &mcp.Tool{
Name: ToolSystemInfo,
Description: "Report LocalAI version, paths, distributed flag, currently loaded models, and installed backends.",
@@ -28,4 +28,21 @@ func registerSystemTools(s *mcp.Server, client LocalAIClient, _ Options) {
}
return jsonResult(nodes), nil, nil
})
+
+ if opts.DisableMutating {
+ return
+ }
+
+ mcp.AddTool(s, &mcp.Tool{
+ Name: ToolSetNodeVRAMBudget,
+ Description: "Set a federated node's VRAM allocation budget (\"80%\" or \"12GB\"; empty clears the override). Only meaningful in distributed mode. Requires user confirmation per safety rule 1.",
+ }, func(ctx context.Context, _ *mcp.CallToolRequest, args SetNodeVRAMBudgetRequest) (*mcp.CallToolResult, any, error) {
+ if args.NodeID == "" {
+ return errorResultf("node_id is required"), nil, nil
+ }
+ if err := client.SetNodeVRAMBudget(ctx, args.NodeID, args.Budget); err != nil {
+ return errorResult(err), nil, nil
+ }
+ return jsonResult(map[string]string{"node_id": args.NodeID, "vram_budget": args.Budget}), nil, nil
+ })
}
diff --git a/pkg/vrambudget/budget.go b/pkg/vrambudget/budget.go
new file mode 100644
index 000000000..672e78294
--- /dev/null
+++ b/pkg/vrambudget/budget.go
@@ -0,0 +1,158 @@
+// Package vrambudget parses an operator-set cap on how much VRAM LocalAI may
+// use for model allocation on a node, and applies it as a hard ceiling.
+//
+// A budget is expressed as either a percentage of detected VRAM ("80%", "0.8")
+// or an absolute amount ("12GB", "12GiB", raw bytes). The zero value is "no
+// cap" so every existing deployment is unaffected until a budget is set.
+package vrambudget
+
+import (
+ "fmt"
+ "math"
+ "strconv"
+ "strings"
+)
+
+// Budget is an optional cap on allocatable VRAM. The zero value means no cap.
+type Budget struct {
+ fraction float64 // (0,1] when percentage form; 0 otherwise
+ absolute uint64 // bytes when absolute form; 0 otherwise
+}
+
+// decimal (1000-based) and binary (1024-based) size suffixes, longest first so
+// "GiB" is matched before "GB"/"B".
+var sizeSuffixes = []struct {
+ suffix string
+ mult uint64
+}{
+ {"KIB", 1 << 10}, {"MIB", 1 << 20}, {"GIB", 1 << 30}, {"TIB", 1 << 40},
+ {"KB", 1000}, {"MB", 1000 * 1000}, {"GB", 1000 * 1000 * 1000}, {"TB", 1000 * 1000 * 1000 * 1000},
+ {"B", 1},
+}
+
+// Parse accepts "", "80%", "0.8", "1.0", "12GB", "12GiB", "12000MB", raw bytes.
+// Empty / zero forms return an unset Budget with no error. A percentage above
+// 100% (a cap that would raise VRAM) is a config error; an absolute value above
+// physical VRAM is harmless and clamped later in Apply.
+func Parse(s string) (Budget, error) {
+ s = strings.TrimSpace(s)
+ if s == "" {
+ return Budget{}, nil
+ }
+ upper := strings.ToUpper(s)
+
+ // Percentage form: trailing '%'.
+ if strings.HasSuffix(upper, "%") {
+ num := strings.TrimSpace(strings.TrimSuffix(upper, "%"))
+ v, err := strconv.ParseFloat(num, 64)
+ if err != nil {
+ return Budget{}, fmt.Errorf("invalid vram budget percentage %q: %w", s, err)
+ }
+ return fromFraction(v/100.0, s)
+ }
+
+ // Absolute form: any known size suffix.
+ for _, sfx := range sizeSuffixes {
+ if strings.HasSuffix(upper, sfx.suffix) {
+ num := strings.TrimSpace(strings.TrimSuffix(upper, sfx.suffix))
+ v, err := strconv.ParseFloat(num, 64)
+ if err != nil || v < 0 {
+ return Budget{}, fmt.Errorf("invalid vram budget %q", s)
+ }
+ return Budget{absolute: uint64(v * float64(sfx.mult))}, nil
+ }
+ }
+
+ // Bare number: (0,1] is a fraction, anything else is absolute bytes.
+ v, err := strconv.ParseFloat(s, 64)
+ if err != nil {
+ return Budget{}, fmt.Errorf("invalid vram budget %q", s)
+ }
+ if v < 0 {
+ return Budget{}, fmt.Errorf("invalid vram budget %q: negative", s)
+ }
+ if v > 0 && v <= 1 {
+ return fromFraction(v, s)
+ }
+ // A bare value above 1 is a byte count and must be whole: a fractional
+ // value >1 is neither a valid fraction (>100%) nor a sensible byte amount.
+ if v != math.Trunc(v) {
+ return Budget{}, fmt.Errorf("invalid vram budget %q: out of range", s)
+ }
+ return Budget{absolute: uint64(v)}, nil
+}
+
+func fromFraction(f float64, orig string) (Budget, error) {
+ if f <= 0 {
+ return Budget{}, nil // 0% == no cap
+ }
+ if f > 1 {
+ return Budget{}, fmt.Errorf("vram budget %q exceeds 100%%", orig)
+ }
+ return Budget{fraction: f}, nil
+}
+
+// IsSet reports whether a cap is configured.
+func (b Budget) IsSet() bool { return b.fraction > 0 || b.absolute > 0 }
+
+// Ceiling resolves the budget to an absolute byte ceiling against detectedTotal,
+// clamped to detectedTotal so an over-large absolute budget can't fabricate
+// VRAM. Returns 0 when unset.
+func (b Budget) Ceiling(detectedTotal uint64) uint64 {
+ if !b.IsSet() {
+ return 0
+ }
+ var ceil uint64
+ if b.fraction > 0 {
+ ceil = uint64(float64(detectedTotal) * b.fraction)
+ } else {
+ ceil = b.absolute
+ }
+ if ceil > detectedTotal {
+ ceil = detectedTotal
+ }
+ return ceil
+}
+
+// Apply caps detected total/free against the budget. Returns the inputs
+// unchanged when the budget is unset.
+func (b Budget) Apply(detectedTotal, detectedFree uint64) (total, free uint64) {
+ if !b.IsSet() {
+ return detectedTotal, detectedFree
+ }
+ ceil := b.Ceiling(detectedTotal)
+ total = min(detectedTotal, ceil)
+ free = min(detectedFree, ceil)
+ return total, free
+}
+
+// String returns the canonical config form ("80%" or "12GB"), or "" when unset.
+func (b Budget) String() string {
+ switch {
+ case b.fraction > 0:
+ return strconv.FormatFloat(b.fraction*100, 'f', -1, 64) + "%"
+ case b.absolute > 0:
+ return canonicalBytes(b.absolute)
+ default:
+ return ""
+ }
+}
+
+// canonicalBytes renders bytes using the largest decimal suffix that divides
+// evenly, falling back to a raw byte count.
+func canonicalBytes(v uint64) string {
+ for _, sfx := range []struct {
+ suffix string
+ mult uint64
+ }{
+ {"TB", 1000 * 1000 * 1000 * 1000},
+ {"GB", 1000 * 1000 * 1000},
+ {"MB", 1000 * 1000},
+ {"KB", 1000},
+ } {
+ if v%sfx.mult == 0 {
+ return strconv.FormatUint(v/sfx.mult, 10) + sfx.suffix
+ }
+ }
+ return strconv.FormatUint(v, 10) + "B"
+}
diff --git a/pkg/vrambudget/budget_test.go b/pkg/vrambudget/budget_test.go
new file mode 100644
index 000000000..a88d15416
--- /dev/null
+++ b/pkg/vrambudget/budget_test.go
@@ -0,0 +1,102 @@
+package vrambudget_test
+
+import (
+ . "github.com/onsi/ginkgo/v2"
+ . "github.com/onsi/gomega"
+
+ "github.com/mudler/LocalAI/pkg/vrambudget"
+)
+
+const gb = uint64(1000 * 1000 * 1000)
+
+var _ = Describe("Budget", func() {
+ Describe("Parse", func() {
+ It("treats empty/zero as unset (no error)", func() {
+ for _, s := range []string{"", " ", "0", "0%", "0b"} {
+ b, err := vrambudget.Parse(s)
+ Expect(err).NotTo(HaveOccurred())
+ Expect(b.IsSet()).To(BeFalse(), "input %q", s)
+ }
+ })
+
+ It("parses percentage forms", func() {
+ for _, s := range []string{"80%", "0.8"} {
+ b, err := vrambudget.Parse(s)
+ Expect(err).NotTo(HaveOccurred())
+ Expect(b.IsSet()).To(BeTrue())
+ total, free := b.Apply(10*gb, 10*gb)
+ Expect(total).To(Equal(8 * gb))
+ Expect(free).To(Equal(8 * gb))
+ }
+ })
+
+ It("treats 100% / 1.0 as a valid no-op ceiling", func() {
+ b, err := vrambudget.Parse("100%")
+ Expect(err).NotTo(HaveOccurred())
+ total, free := b.Apply(10*gb, 6*gb)
+ Expect(total).To(Equal(10 * gb))
+ Expect(free).To(Equal(6 * gb))
+ })
+
+ It("parses absolute forms with decimal and binary suffixes and raw bytes", func() {
+ cases := map[string]uint64{
+ "12GB": 12 * gb,
+ "12000MB": 12 * gb,
+ "12000000000": 12 * gb,
+ "12GiB": 12 * 1024 * 1024 * 1024,
+ }
+ for s, want := range cases {
+ b, err := vrambudget.Parse(s)
+ Expect(err).NotTo(HaveOccurred(), "input %q", s)
+ Expect(b.Ceiling(64 * gb)).To(Equal(want), "input %q", s)
+ }
+ })
+
+ It("rejects malformed and out-of-range input", func() {
+ for _, s := range []string{"120%", "1.5", "-1", "abc", "12ZB", "%", "12 GB x"} {
+ _, err := vrambudget.Parse(s)
+ Expect(err).To(HaveOccurred(), "input %q", s)
+ }
+ })
+ })
+
+ Describe("Apply", func() {
+ It("is a passthrough when unset", func() {
+ b, _ := vrambudget.Parse("")
+ total, free := b.Apply(16*gb, 9*gb)
+ Expect(total).To(Equal(16 * gb))
+ Expect(free).To(Equal(9 * gb))
+ })
+
+ It("caps total and free to the budget ceiling (hard ceiling)", func() {
+ b, _ := vrambudget.Parse("12GB")
+ total, free := b.Apply(16*gb, 15*gb)
+ Expect(total).To(Equal(12 * gb))
+ Expect(free).To(Equal(12 * gb))
+ })
+
+ It("never raises free above detected when free is already below budget", func() {
+ b, _ := vrambudget.Parse("12GB")
+ _, free := b.Apply(16*gb, 5*gb)
+ Expect(free).To(Equal(5 * gb))
+ })
+
+ It("clamps an absolute budget above physical to detected", func() {
+ b, _ := vrambudget.Parse("64GB")
+ total, free := b.Apply(16*gb, 10*gb)
+ Expect(total).To(Equal(16 * gb))
+ Expect(free).To(Equal(10 * gb))
+ })
+ })
+
+ Describe("String", func() {
+ It("round-trips percentage and absolute canonical forms", func() {
+ p, _ := vrambudget.Parse("80%")
+ Expect(p.String()).To(Equal("80%"))
+ a, _ := vrambudget.Parse("12000MB")
+ Expect(a.String()).To(Equal("12GB"))
+ z, _ := vrambudget.Parse("")
+ Expect(z.String()).To(Equal(""))
+ })
+ })
+})
diff --git a/pkg/vrambudget/vrambudget_suite_test.go b/pkg/vrambudget/vrambudget_suite_test.go
new file mode 100644
index 000000000..5c3bdd04b
--- /dev/null
+++ b/pkg/vrambudget/vrambudget_suite_test.go
@@ -0,0 +1,13 @@
+package vrambudget_test
+
+import (
+ "testing"
+
+ . "github.com/onsi/ginkgo/v2"
+ . "github.com/onsi/gomega"
+)
+
+func TestVRAMBudget(t *testing.T) {
+ RegisterFailHandler(Fail)
+ RunSpecs(t, "VRAMBudget test suite")
+}
diff --git a/pkg/xsysinfo/gpu.go b/pkg/xsysinfo/gpu.go
index da183212f..92a85c981 100644
--- a/pkg/xsysinfo/gpu.go
+++ b/pkg/xsysinfo/gpu.go
@@ -10,12 +10,37 @@ import (
"strconv"
"strings"
"sync"
+ "sync/atomic"
"github.com/jaypipes/ghw"
"github.com/jaypipes/ghw/pkg/gpu"
+ "github.com/mudler/LocalAI/pkg/vrambudget"
"github.com/mudler/xlog"
)
+// defaultVRAMBudget is the process-wide allocation cap set by standalone
+// local-ai (LOCALAI_VRAM_BUDGET / the Settings page). nil pointer = no cap.
+// The aggregate getters below apply it before returning so every allocation
+// decision (hardware defaults, context-fit, watchdog) inherits the cap without
+// per-call-site changes. Distributed workers deliberately leave this unset and
+// report raw VRAM; the server applies per-node budgets instead.
+var defaultVRAMBudget atomic.Pointer[vrambudget.Budget]
+
+// SetDefaultVRAMBudget installs the process-wide VRAM allocation cap. Safe to
+// call at startup and again live when the Settings page changes it.
+func SetDefaultVRAMBudget(b vrambudget.Budget) {
+ bb := b
+ defaultVRAMBudget.Store(&bb)
+}
+
+// DefaultVRAMBudget returns the current process-wide cap (unset Budget = none).
+func DefaultVRAMBudget() vrambudget.Budget {
+ if p := defaultVRAMBudget.Load(); p != nil {
+ return *p
+ }
+ return vrambudget.Budget{}
+}
+
// GPU vendor constants
const (
VendorNVIDIA = "nvidia"
@@ -107,7 +132,8 @@ func TotalAvailableVRAM() (uint64, error) {
}
// If we got valid VRAM from ghw, return it
if totalVRAM > 0 {
- return totalVRAM, nil
+ capped, _ := DefaultVRAMBudget().Apply(totalVRAM, totalVRAM)
+ return capped, nil
}
}
@@ -121,7 +147,8 @@ func TotalAvailableVRAM() (uint64, error) {
}
if totalVRAM > 0 {
xlog.Debug("VRAM detected via binary tools", "total_vram", totalVRAM)
- return totalVRAM, nil
+ capped, _ := DefaultVRAMBudget().Apply(totalVRAM, totalVRAM)
+ return capped, nil
}
}
@@ -144,7 +171,8 @@ func MinPerGPUVRAM() (uint64, error) {
// hosts, which is why TotalAvailableVRAM treats it as a sum.
if infos := GetGPUMemoryUsage(); len(infos) > 0 {
if v := minNonZeroVRAM(infos); v > 0 {
- return v, nil
+ capped, _ := DefaultVRAMBudget().Apply(v, v)
+ return capped, nil
}
}
@@ -162,7 +190,8 @@ func MinPerGPUVRAM() (uint64, error) {
}
}
if min > 0 {
- return min, nil
+ capped, _ := DefaultVRAMBudget().Apply(min, min)
+ return capped, nil
}
}
@@ -359,6 +388,16 @@ func GetGPUAggregateInfo() GPUAggregateInfo {
aggregate.UsagePercent = float64(aggregate.UsedVRAM) / float64(aggregate.TotalVRAM) * 100
}
+ // Apply the process-wide allocation cap so scheduling/hardware-default
+ // consumers see budgeted capacity, not physical capacity.
+ if b := DefaultVRAMBudget(); b.IsSet() && aggregate.TotalVRAM > 0 {
+ aggregate.TotalVRAM, aggregate.FreeVRAM = b.Apply(aggregate.TotalVRAM, aggregate.FreeVRAM)
+ aggregate.UsedVRAM = aggregate.TotalVRAM - aggregate.FreeVRAM
+ if aggregate.TotalVRAM > 0 {
+ aggregate.UsagePercent = float64(aggregate.UsedVRAM) / float64(aggregate.TotalVRAM) * 100
+ }
+ }
+
return aggregate
}
@@ -971,8 +1010,15 @@ func GetResourceInfo() ResourceInfo {
// GetResourceAggregateInfo returns aggregate memory info (GPU if available, otherwise RAM)
// This is used by the memory reclaimer to check memory usage
func GetResourceAggregateInfo() AggregateMemoryInfo {
- resourceInfo := GetResourceInfo()
- return resourceInfo.Aggregate
+ // The VRAM budget is applied exactly once, upstream: the GPU-branch
+ // Aggregate is sourced from GetGPUAggregateInfo, which already caps
+ // total/free/used against DefaultVRAMBudget. Re-applying b.Apply here would
+ // double-cap: Apply resolves a percentage budget as a fraction of its input
+ // total, so a second pass shrinks the already-shrunk total again (P*(P*T)
+ // instead of P*T) and corrupts UsagePercent read by the memory reclaimer.
+ // The RAM branch is never budgeted (the budget is VRAM-only). So this
+ // boundary must return the aggregate unchanged.
+ return GetResourceInfo().Aggregate
}
// getVulkanGPUMemory queries GPUs using vulkaninfo as a fallback.
diff --git a/pkg/xsysinfo/vrambudget_boundary_test.go b/pkg/xsysinfo/vrambudget_boundary_test.go
new file mode 100644
index 000000000..8fef3f583
--- /dev/null
+++ b/pkg/xsysinfo/vrambudget_boundary_test.go
@@ -0,0 +1,26 @@
+package xsysinfo_test
+
+import (
+ . "github.com/onsi/ginkgo/v2"
+ . "github.com/onsi/gomega"
+
+ "github.com/mudler/LocalAI/pkg/vrambudget"
+ "github.com/mudler/LocalAI/pkg/xsysinfo"
+)
+
+var _ = Describe("Default VRAM budget", func() {
+ AfterEach(func() {
+ xsysinfo.SetDefaultVRAMBudget(vrambudget.Budget{}) // reset to no cap
+ })
+
+ It("defaults to unset", func() {
+ Expect(xsysinfo.DefaultVRAMBudget().IsSet()).To(BeFalse())
+ })
+
+ It("stores and returns a set budget", func() {
+ b, _ := vrambudget.Parse("80%")
+ xsysinfo.SetDefaultVRAMBudget(b)
+ Expect(xsysinfo.DefaultVRAMBudget().IsSet()).To(BeTrue())
+ Expect(xsysinfo.DefaultVRAMBudget().String()).To(Equal("80%"))
+ })
+})