From 8cec22c3b7926c737d36afe80e8bbf4ab3e75d97 Mon Sep 17 00:00:00 2001 From: "LocalAI [bot]" <139863280+localai-bot@users.noreply.github.com> Date: Wed, 15 Jul 2026 09:58:45 +0200 Subject: [PATCH] feat(vram): per-node VRAM allocation budget (LOCALAI_VRAM_BUDGET) (#10833) * feat(vram): add vrambudget primitive for per-node VRAM caps Signed-off-by: Ettore Di Giacinto * feat(vram): apply default VRAM budget in xsysinfo aggregate getters Signed-off-by: Ettore Di Giacinto * feat(vram): wire LOCALAI_VRAM_BUDGET flag to xsysinfo default budget Signed-off-by: Ettore Di Giacinto * feat(vram): persist VRAM budget via runtime settings with live apply Signed-off-by: Ettore Di Giacinto * test(vram): reset process-global VRAM budget after runtime-settings spec Signed-off-by: Ettore Di Giacinto * feat(vram): add VRAM budget field to Settings page Signed-off-by: Ettore Di Giacinto * feat(vram): store and enforce per-node VRAM budget in the node registry Signed-off-by: Ettore Di Giacinto * feat(vram): apply per-node VRAM budget in router hardware defaults Signed-off-by: Ettore Di Giacinto * feat(vram): report worker VRAM budget in node registration The distributed worker now reports its operator-set VRAM budget string (LOCALAI_VRAM_BUDGET) to the server on registration. The worker keeps reporting RAW total/available VRAM and never sets the xsysinfo process-global budget (that stays standalone-only); the server resolves and enforces the budget uniformly (Task 6). Also closes a Task 6 gap: on re-registration, a struct Updates zero-skips an empty budget, so a worker that dropped LOCALAI_VRAM_BUDGET left the stale cap in place. For non-admin-override nodes the budget columns are now force-written (map Updates) even when empty, so removing the env var clears the cap; admin overrides are preserved unchanged. Signed-off-by: Ettore Di Giacinto * style(vram): drop em dash from worker-clear comment Signed-off-by: Ettore Di Giacinto * feat(vram): add node VRAM budget admin endpoints Signed-off-by: Ettore Di Giacinto * feat(vram): add node VRAM budget control to the node UI Signed-off-by: Ettore Di Giacinto * feat(vram): expose set_node_vram_budget MCP admin tool Signed-off-by: Ettore Di Giacinto * docs(vram): document LOCALAI_VRAM_BUDGET and node VRAM budget UI Signed-off-by: Ettore Di Giacinto * fix(vram): avoid double-applying VRAM budget in GetResourceAggregateInfo The GPU-branch aggregate returned by GetResourceInfo is sourced from GetGPUAggregateInfo, which already caps total/free/used against the process-wide VRAM budget. GetResourceAggregateInfo then applied the budget a second time. For an absolute budget this is idempotent, but for a percentage budget b.Apply resolves the ceiling as a fraction of its input total, so a second pass yields P*(P*T) instead of P*T and distorts UsagePercent (read by the memory reclaimer in pkg/model/watchdog.go). Remove the redundant second application so the budget is applied exactly once, against the raw physical totals, upstream in GetGPUAggregateInfo. Signed-off-by: Ettore Di Giacinto * fix(vram): implement SetNodeVRAMBudget on mcp assistant test stub The LocalAIClient interface gained SetNodeVRAMBudget; the stubClient in core/http/endpoints/mcp used by the assistant tests is a separate implementer and needs the method too (broke golangci-lint typecheck and both test jobs). Signed-off-by: Ettore Di Giacinto --------- Signed-off-by: Ettore Di Giacinto Co-authored-by: Ettore Di Giacinto --- core/cli/run.go | 17 ++ core/config/application_config.go | 23 +++ core/config/application_config_test.go | 7 + core/config/runtime_settings.go | 19 +- core/config/runtime_settings_persist_test.go | 34 ++++ core/http/endpoints/localai/nodes.go | 73 ++++++++ .../localai/nodes_vrambudget_test.go | 99 ++++++++++ core/http/endpoints/localai/settings.go | 13 ++ .../endpoints/mcp/localai_assistant_test.go | 4 + .../src/components/nodes/CapacityEditor.jsx | 168 +++++++++++++++++ core/http/react-ui/src/pages/Settings.jsx | 3 + core/http/react-ui/src/utils/api.js | 11 ++ core/http/react-ui/src/utils/config.js | 1 + core/http/routes/nodes.go | 6 + core/services/nodes/registry.go | 148 ++++++++++++++- .../nodes/registry_vrambudget_test.go | 169 ++++++++++++++++++ core/services/nodes/router.go | 12 +- core/services/nodes/router_vrambudget_test.go | 40 +++++ core/services/worker/config.go | 5 + core/services/worker/registration.go | 8 + core/services/worker/registration_test.go | 33 ++++ docs/content/advanced/vram-management.md | 43 +++++ docs/content/features/distributed-mode.md | 33 ++++ pkg/mcp/localaitools/client.go | 4 + pkg/mcp/localaitools/coverage_test.go | 1 + pkg/mcp/localaitools/dto.go | 9 + pkg/mcp/localaitools/fakes_test.go | 9 + pkg/mcp/localaitools/httpapi/client.go | 7 + pkg/mcp/localaitools/httpapi/routes.go | 4 + pkg/mcp/localaitools/inproc/client.go | 8 + pkg/mcp/localaitools/server_test.go | 1 + pkg/mcp/localaitools/tools.go | 1 + pkg/mcp/localaitools/tools_system.go | 19 +- pkg/vrambudget/budget.go | 158 ++++++++++++++++ pkg/vrambudget/budget_test.go | 102 +++++++++++ pkg/vrambudget/vrambudget_suite_test.go | 13 ++ pkg/xsysinfo/gpu.go | 58 +++++- pkg/xsysinfo/vrambudget_boundary_test.go | 26 +++ 38 files changed, 1365 insertions(+), 24 deletions(-) create mode 100644 core/http/endpoints/localai/nodes_vrambudget_test.go create mode 100644 core/services/nodes/registry_vrambudget_test.go create mode 100644 core/services/nodes/router_vrambudget_test.go create mode 100644 core/services/worker/registration_test.go create mode 100644 pkg/vrambudget/budget.go create mode 100644 pkg/vrambudget/budget_test.go create mode 100644 pkg/vrambudget/vrambudget_suite_test.go create mode 100644 pkg/xsysinfo/vrambudget_boundary_test.go diff --git a/core/cli/run.go b/core/cli/run.go index b067098b1..2b42c9efa 100644 --- a/core/cli/run.go +++ b/core/cli/run.go @@ -19,6 +19,8 @@ import ( "github.com/mudler/LocalAI/pkg/modelartifacts" "github.com/mudler/LocalAI/pkg/signals" "github.com/mudler/LocalAI/pkg/system" + "github.com/mudler/LocalAI/pkg/vrambudget" + "github.com/mudler/LocalAI/pkg/xsysinfo" "github.com/mudler/xlog" ) @@ -96,6 +98,7 @@ type RunCMD struct { WatchdogInterval string `env:"LOCALAI_WATCHDOG_INTERVAL,WATCHDOG_INTERVAL" default:"500ms" help:"Interval between watchdog checks (e.g., 500ms, 5s, 1m) (default: 500ms)" group:"backends"` EnableMemoryReclaimer bool `env:"LOCALAI_MEMORY_RECLAIMER,MEMORY_RECLAIMER,LOCALAI_GPU_RECLAIMER,GPU_RECLAIMER" default:"false" help:"Enable memory threshold monitoring to auto-evict backends when memory usage exceeds threshold (uses GPU VRAM if available, otherwise RAM)" group:"backends"` MemoryReclaimerThreshold float64 `env:"LOCALAI_MEMORY_RECLAIMER_THRESHOLD,MEMORY_RECLAIMER_THRESHOLD,LOCALAI_GPU_RECLAIMER_THRESHOLD,GPU_RECLAIMER_THRESHOLD" default:"0.95" help:"Memory usage threshold (0.0-1.0) that triggers backend eviction (default 0.95 = 95%%)" group:"backends"` + VRAMBudget string `env:"LOCALAI_VRAM_BUDGET" help:"Cap VRAM used for model allocation on this node, as a percentage (e.g. 80%) or absolute amount (e.g. 12GB). Empty uses all detected VRAM." group:"backends"` ForceEvictionWhenBusy bool `env:"LOCALAI_FORCE_EVICTION_WHEN_BUSY,FORCE_EVICTION_WHEN_BUSY" default:"false" help:"Force eviction even when models have active API calls (default: false for safety)" group:"backends"` SizeAwareEviction bool `env:"LOCALAI_SIZE_AWARE_EVICTION,SIZE_AWARE_EVICTION" default:"false" help:"Evict the largest loaded model first rather than the least-recently-used one, keeping small utility models resident and maximizing freed memory per eviction" group:"backends"` LRUEvictionMaxRetries int `env:"LOCALAI_LRU_EVICTION_MAX_RETRIES,LRU_EVICTION_MAX_RETRIES" default:"30" help:"Maximum number of retries when waiting for busy models to become idle before eviction (default: 30)" group:"backends"` @@ -667,6 +670,20 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error { opts = append(opts, config.WithPreferDevelopmentBackends(r.PreferDevelopmentBackends)) } + // Per-node VRAM allocation budget. Record it on the ApplicationConfig and, + // fail-open, install it as the process-global xsysinfo default so allocation + // heuristics cap at it. A malformed value must never block startup: it is + // logged and treated as unset (full detected VRAM). + if r.VRAMBudget != "" { + opts = append(opts, config.SetVRAMBudget(r.VRAMBudget)) + if b, err := vrambudget.Parse(r.VRAMBudget); err != nil { + xlog.Warn("Ignoring invalid LOCALAI_VRAM_BUDGET", "value", r.VRAMBudget, "error", err) + } else { + xsysinfo.SetDefaultVRAMBudget(b) + xlog.Info("VRAM allocation budget set", "budget", b.String()) + } + } + if r.PreloadBackendOnly { _, err := application.New(opts...) return err diff --git a/core/config/application_config.go b/core/config/application_config.go index 047a16703..5e6bb074f 100644 --- a/core/config/application_config.go +++ b/core/config/application_config.go @@ -9,6 +9,7 @@ import ( "github.com/mudler/LocalAI/pkg/model" "github.com/mudler/LocalAI/pkg/modelartifacts" "github.com/mudler/LocalAI/pkg/system" + "github.com/mudler/LocalAI/pkg/vrambudget" "github.com/mudler/LocalAI/pkg/xsysinfo" "github.com/mudler/xlog" ) @@ -133,6 +134,10 @@ type ApplicationConfig struct { MemoryReclaimerEnabled bool // Enable memory threshold monitoring MemoryReclaimerThreshold float64 // Threshold 0.0-1.0 (e.g., 0.95 = 95%) + // VRAMBudget optionally caps how much VRAM this instance uses for model + // allocation, as "80%" or "12GB". Empty = use full detected VRAM. + VRAMBudget string + // Eviction settings ForceEvictionWhenBusy bool // Force eviction even when models have active API calls (default: false for safety) SizeAwareEviction bool // Evict largest models first rather than least-recently-used (default: false) @@ -464,6 +469,13 @@ func SetMemoryReclaimerThreshold(threshold float64) AppOption { } } +// SetVRAMBudget sets the VRAM allocation cap ("80%" or "12GB", "" = no cap). +func SetVRAMBudget(v string) AppOption { + return func(o *ApplicationConfig) { + o.VRAMBudget = v + } +} + // WithMemoryReclaimer configures the memory reclaimer with the given settings func WithMemoryReclaimer(enabled bool, threshold float64) AppOption { return func(o *ApplicationConfig) { @@ -1124,6 +1136,7 @@ func (o *ApplicationConfig) ToRuntimeSettings() RuntimeSettings { lruEvictionMaxRetries := o.LRUEvictionMaxRetries threads := o.Threads contextSize := o.ContextSize + vramBudget := o.VRAMBudget f16 := o.F16 debug := o.Debug tracingMaxItems := o.TracingMaxItems @@ -1218,6 +1231,7 @@ func (o *ApplicationConfig) ToRuntimeSettings() RuntimeSettings { LRUEvictionRetryInterval: &lruEvictionRetryInterval, Threads: &threads, ContextSize: &contextSize, + VRAMBudget: &vramBudget, F16: &f16, Debug: &debug, TracingMaxItems: &tracingMaxItems, @@ -1358,6 +1372,15 @@ func (o *ApplicationConfig) ApplyRuntimeSettings(settings *RuntimeSettings) (req if settings.ContextSize != nil { o.ContextSize = *settings.ContextSize } + if settings.VRAMBudget != nil { + o.VRAMBudget = *settings.VRAMBudget + // Live-apply so the cap takes effect without a restart. An empty string + // clears the cap; a malformed value is rejected by the settings endpoint, + // but stay fail-open here too so a bad persisted value cannot wedge apply. + if b, err := vrambudget.Parse(o.VRAMBudget); err == nil { + xsysinfo.SetDefaultVRAMBudget(b) + } + } if settings.F16 != nil { o.F16 = *settings.F16 } diff --git a/core/config/application_config_test.go b/core/config/application_config_test.go index bdfa1a271..6860388e7 100644 --- a/core/config/application_config_test.go +++ b/core/config/application_config_test.go @@ -648,4 +648,11 @@ var _ = Describe("ApplicationConfig RuntimeSettings Conversion", func() { Expect(appConfig.SingleBackend).To(BeFalse()) // 3 != 1, so single backend is false }) }) + + Describe("SetVRAMBudget", func() { + It("stores the VRAM budget via SetVRAMBudget", func() { + o := NewApplicationConfig(SetVRAMBudget("80%")) + Expect(o.VRAMBudget).To(Equal("80%")) + }) + }) }) diff --git a/core/config/runtime_settings.go b/core/config/runtime_settings.go index a7f15e658..73c78ccda 100644 --- a/core/config/runtime_settings.go +++ b/core/config/runtime_settings.go @@ -28,19 +28,20 @@ type RuntimeSettings struct { // Eviction settings ForceEvictionWhenBusy *bool `json:"force_eviction_when_busy,omitempty"` // Force eviction even when models have active API calls (default: false for safety) - SizeAwareEviction *bool `json:"size_aware_eviction,omitempty"` // Evict largest models first rather than least-recently-used (default: false) + SizeAwareEviction *bool `json:"size_aware_eviction,omitempty"` // Evict largest models first rather than least-recently-used (default: false) LRUEvictionMaxRetries *int `json:"lru_eviction_max_retries,omitempty"` // Maximum number of retries when waiting for busy models to become idle (default: 30) LRUEvictionRetryInterval *string `json:"lru_eviction_retry_interval,omitempty"` // Interval between retries when waiting for busy models (e.g., 1s, 2s) (default: 1s) // Performance settings - Threads *int `json:"threads,omitempty"` - ContextSize *int `json:"context_size,omitempty"` - F16 *bool `json:"f16,omitempty"` - Debug *bool `json:"debug,omitempty"` - EnableTracing *bool `json:"enable_tracing,omitempty"` - TracingMaxItems *int `json:"tracing_max_items,omitempty"` - TracingMaxBodyBytes *int `json:"tracing_max_body_bytes,omitempty"` // Per-body cap in bytes; 0 disables the cap - EnableBackendLogging *bool `json:"enable_backend_logging,omitempty"` + Threads *int `json:"threads,omitempty"` + ContextSize *int `json:"context_size,omitempty"` + VRAMBudget *string `json:"vram_budget,omitempty"` // Cap VRAM for allocation ("80%" or "12GB"; "" = no cap) + F16 *bool `json:"f16,omitempty"` + Debug *bool `json:"debug,omitempty"` + EnableTracing *bool `json:"enable_tracing,omitempty"` + TracingMaxItems *int `json:"tracing_max_items,omitempty"` + TracingMaxBodyBytes *int `json:"tracing_max_body_bytes,omitempty"` // Per-body cap in bytes; 0 disables the cap + EnableBackendLogging *bool `json:"enable_backend_logging,omitempty"` // Security/CORS settings CORS *bool `json:"cors,omitempty"` diff --git a/core/config/runtime_settings_persist_test.go b/core/config/runtime_settings_persist_test.go index 7f5eb07a8..f5d9644d3 100644 --- a/core/config/runtime_settings_persist_test.go +++ b/core/config/runtime_settings_persist_test.go @@ -9,6 +9,8 @@ import ( . "github.com/onsi/gomega" "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/pkg/vrambudget" + "github.com/mudler/LocalAI/pkg/xsysinfo" ) func strPtr(s string) *string { return &s } @@ -93,6 +95,38 @@ var _ = Describe("RuntimeSettings persistence helpers", func() { }) }) + // VRAMBudget round trip pins the Settings-page persistence contract: the + // operator-set cap must survive ToRuntimeSettings (GET /api/settings) -> + // ApplyRuntimeSettings (POST /api/settings) so it lives past a save, and an + // empty value must clear the cap rather than being dropped. + Describe("VRAMBudget round trip", func() { + // ApplyRuntimeSettings live-applies the cap through the process-global + // xsysinfo.SetDefaultVRAMBudget. Reset it after each spec so a cap set + // here cannot bleed into other core/config specs under Ginkgo's + // randomized ordering. + AfterEach(func() { + xsysinfo.SetDefaultVRAMBudget(vrambudget.Budget{}) + }) + + It("round-trips the VRAM budget", func() { + o := config.NewApplicationConfig(config.SetVRAMBudget("80%")) + rs := o.ToRuntimeSettings() + Expect(rs.VRAMBudget).NotTo(BeNil()) + Expect(*rs.VRAMBudget).To(Equal("80%")) + + o2 := config.NewApplicationConfig() + o2.ApplyRuntimeSettings(&rs) + Expect(o2.VRAMBudget).To(Equal("80%")) + }) + + It("applies an empty VRAM budget as clearing the cap", func() { + o := config.NewApplicationConfig(config.SetVRAMBudget("80%")) + empty := "" + o.ApplyRuntimeSettings(&config.RuntimeSettings{VRAMBudget: &empty}) + Expect(o.VRAMBudget).To(Equal("")) + }) + }) + // MITM round trip pins the contract that loadRuntimeSettingsFromFile // MITM listener address must survive a write/read round trip so the // next process restart can bring the listener back up. (Intercept diff --git a/core/http/endpoints/localai/nodes.go b/core/http/endpoints/localai/nodes.go index 2fb4ac530..f5a658501 100644 --- a/core/http/endpoints/localai/nodes.go +++ b/core/http/endpoints/localai/nodes.go @@ -30,6 +30,7 @@ import ( "github.com/mudler/LocalAI/core/services/nodes/prefixcache" "github.com/mudler/LocalAI/pkg/httpclient" "github.com/mudler/LocalAI/pkg/natsauth" + "github.com/mudler/LocalAI/pkg/vrambudget" ) // nodeError builds a schema.ErrorResponse for node endpoints. @@ -92,6 +93,9 @@ type RegisterNodeRequest struct { // Workers older than this field omit it; we coerce 0 → 1 below to preserve // historical single-replica behavior. MaxReplicasPerModel int `json:"max_replicas_per_model,omitempty"` + // VRAMBudget is the worker's operator-set VRAM cap ("80%" or "12GB"). The + // registry resolves and enforces it against the raw reported VRAM. + VRAMBudget string `json:"vram_budget,omitempty"` } // RegisterNodeEndpoint registers a new backend node. @@ -171,6 +175,7 @@ func RegisterNodeEndpoint(registry *nodes.NodeRegistry, expectedToken string, au GPUVendor: req.GPUVendor, GPUComputeCapability: req.GPUComputeCapability, MaxReplicasPerModel: maxReplicasPerModel, + VRAMBudget: req.VRAMBudget, } ctx := c.Request().Context() @@ -942,6 +947,74 @@ func ResetMaxReplicasPerModelEndpoint(registry *nodes.NodeRegistry) echo.Handler } } +// UpdateVRAMBudgetRequest is the body for the per-node VRAM budget endpoint. +type UpdateVRAMBudgetRequest struct { + // Value is the VRAM cap ("80%" or "12GB"). Empty string clears the cap. + Value string `json:"value"` +} + +// UpdateVRAMBudgetEndpoint sets a node's VRAM allocation cap as a sticky admin +// override. The value is validated, resolved against the node's total VRAM, and +// applied to available_vram immediately so scheduling reflects it before the +// next heartbeat. +// +// @Summary Update a node's VRAM allocation budget +// @Tags Nodes +// @Param id path string true "Node ID" +// @Param request body UpdateVRAMBudgetRequest true "New value (\"80%\" or \"12GB\"; empty clears)" +// @Success 200 {object} map[string]string +// @Failure 400 {object} map[string]any "invalid budget" +// @Failure 404 {object} map[string]any "node not found" +// @Router /api/nodes/{id}/vram-budget [put] +func UpdateVRAMBudgetEndpoint(registry *nodes.NodeRegistry) echo.HandlerFunc { + return func(c echo.Context) error { + ctx := c.Request().Context() + nodeID := c.Param("id") + if _, err := registry.Get(ctx, nodeID); err != nil { + return c.JSON(http.StatusNotFound, nodeError(http.StatusNotFound, "node not found")) + } + var req UpdateVRAMBudgetRequest + if err := c.Bind(&req); err != nil { + return c.JSON(http.StatusBadRequest, nodeError(http.StatusBadRequest, "invalid request body")) + } + // An empty value clears the cap; only a non-empty value must parse. + if req.Value != "" { + if _, err := vrambudget.Parse(req.Value); err != nil { + return c.JSON(http.StatusBadRequest, nodeError(http.StatusBadRequest, fmt.Sprintf("invalid vram budget: %v", err))) + } + } + if err := registry.UpdateVRAMBudget(ctx, nodeID, req.Value); err != nil { + xlog.Error("Failed to update vram_budget", "node", nodeID, "value", req.Value, "error", err) + return c.JSON(http.StatusInternalServerError, nodeError(http.StatusInternalServerError, "failed to update vram budget")) + } + return c.JSON(http.StatusOK, map[string]string{"vram_budget": req.Value}) + } +} + +// ResetVRAMBudgetEndpoint clears the admin override so the worker's +// LOCALAI_VRAM_BUDGET takes over again on next registration. +// +// @Summary Reset a node's VRAM budget to the worker default +// @Tags Nodes +// @Param id path string true "Node ID" +// @Success 200 {object} map[string]bool +// @Failure 404 {object} map[string]any "node not found" +// @Router /api/nodes/{id}/vram-budget [delete] +func ResetVRAMBudgetEndpoint(registry *nodes.NodeRegistry) echo.HandlerFunc { + return func(c echo.Context) error { + ctx := c.Request().Context() + nodeID := c.Param("id") + if _, err := registry.Get(ctx, nodeID); err != nil { + return c.JSON(http.StatusNotFound, nodeError(http.StatusNotFound, "node not found")) + } + if err := registry.ResetVRAMBudget(ctx, nodeID); err != nil { + xlog.Error("Failed to reset vram_budget override", "node", nodeID, "error", err) + return c.JSON(http.StatusInternalServerError, nodeError(http.StatusInternalServerError, "failed to reset vram budget")) + } + return c.JSON(http.StatusOK, map[string]bool{"reset": true}) + } +} + // SetNodeLabelsEndpoint replaces all labels for a node. func SetNodeLabelsEndpoint(registry *nodes.NodeRegistry) echo.HandlerFunc { return func(c echo.Context) error { diff --git a/core/http/endpoints/localai/nodes_vrambudget_test.go b/core/http/endpoints/localai/nodes_vrambudget_test.go new file mode 100644 index 000000000..51eb60013 --- /dev/null +++ b/core/http/endpoints/localai/nodes_vrambudget_test.go @@ -0,0 +1,99 @@ +package localai + +import ( + "context" + "net/http" + "net/http/httptest" + "strings" + + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/testutil" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Node VRAM budget HTTP handlers", func() { + var registry *nodes.NodeRegistry + + BeforeEach(func() { + db := testutil.SetupTestDB() + var err error + registry, err = nodes.NewNodeRegistry(db) + Expect(err).ToNot(HaveOccurred()) + }) + + // seedHealthyNode registers a healthy node with the given total/available + // VRAM and returns its ID. + seedHealthyNode := func(ctx context.Context, total, available uint64) string { + node := &nodes.BackendNode{ + ID: "vram-node", + Name: "vram-node", + Address: "10.0.0.9:50051", + Status: nodes.StatusHealthy, + TotalVRAM: total, + AvailableVRAM: available, + } + Expect(registry.Register(ctx, node, true)).To(Succeed()) + return node.ID + } + + // doPUT drives UpdateVRAMBudgetEndpoint for the given node ID and body. + doPUT := func(id, body string) *httptest.ResponseRecorder { + e := echo.New() + req := httptest.NewRequest(http.MethodPut, "/", strings.NewReader(body)) + req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON) + rec := httptest.NewRecorder() + c := e.NewContext(req, rec) + c.SetParamNames("id") + c.SetParamValues(id) + Expect(UpdateVRAMBudgetEndpoint(registry)(c)).To(Succeed()) + return rec + } + + // doDELETE drives ResetVRAMBudgetEndpoint for the given node ID. + doDELETE := func(id string) *httptest.ResponseRecorder { + e := echo.New() + req := httptest.NewRequest(http.MethodDelete, "/", nil) + rec := httptest.NewRecorder() + c := e.NewContext(req, rec) + c.SetParamNames("id") + c.SetParamValues(id) + Expect(ResetVRAMBudgetEndpoint(registry)(c)).To(Succeed()) + return rec + } + + It("sets and clears a node's VRAM budget", func(ctx SpecContext) { + id := seedHealthyNode(ctx, 16_000_000_000, 16_000_000_000) + + rec := doPUT(id, `{"value":"50%"}`) + Expect(rec.Code).To(Equal(http.StatusOK)) + node, err := registry.Get(ctx, id) + Expect(err).ToNot(HaveOccurred()) + Expect(node.VRAMBudget).To(Equal("50%")) + Expect(node.AvailableVRAM).To(Equal(uint64(8_000_000_000))) + + rec = doDELETE(id) + Expect(rec.Code).To(Equal(http.StatusOK)) + node, err = registry.Get(ctx, id) + Expect(err).ToNot(HaveOccurred()) + Expect(node.VRAMBudget).To(Equal("")) + }) + + It("rejects a malformed budget with 400", func(ctx SpecContext) { + id := seedHealthyNode(ctx, 16_000_000_000, 16_000_000_000) + rec := doPUT(id, `{"value":"120%"}`) + Expect(rec.Code).To(Equal(http.StatusBadRequest)) + }) + + It("returns 404 for an unknown node on update", func(ctx SpecContext) { + rec := doPUT("does-not-exist", `{"value":"50%"}`) + Expect(rec.Code).To(Equal(http.StatusNotFound)) + }) + + It("returns 404 for an unknown node on reset", func(ctx SpecContext) { + rec := doDELETE("does-not-exist") + Expect(rec.Code).To(Equal(http.StatusNotFound)) + }) +}) diff --git a/core/http/endpoints/localai/settings.go b/core/http/endpoints/localai/settings.go index 8033f07d5..dad6bbb44 100644 --- a/core/http/endpoints/localai/settings.go +++ b/core/http/endpoints/localai/settings.go @@ -12,6 +12,7 @@ import ( "github.com/mudler/LocalAI/core/http/endpoints/openresponses" "github.com/mudler/LocalAI/core/p2p" "github.com/mudler/LocalAI/core/schema" + "github.com/mudler/LocalAI/pkg/vrambudget" "github.com/mudler/xlog" ) @@ -94,6 +95,18 @@ func UpdateSettingsEndpoint(app *application.Application) echo.HandlerFunc { } } + // Reject a malformed VRAM budget up front so a bad value never reaches + // the config (ApplyRuntimeSettings is fail-open and would silently skip + // installing the cap). An empty string is allowed: it clears the cap. + if settings.VRAMBudget != nil && *settings.VRAMBudget != "" { + if _, err := vrambudget.Parse(*settings.VRAMBudget); err != nil { + return c.JSON(http.StatusBadRequest, schema.SettingsResponse{ + Success: false, + Error: "Invalid vram_budget format: " + err.Error(), + }) + } + } + // Generate P2P token before saving so the real token is persisted (not "0") if settings.P2PToken != nil && *settings.P2PToken == "0" { token := p2p.GenerateToken(60, 60) diff --git a/core/http/endpoints/mcp/localai_assistant_test.go b/core/http/endpoints/mcp/localai_assistant_test.go index d2f2481ff..2231350d0 100644 --- a/core/http/endpoints/mcp/localai_assistant_test.go +++ b/core/http/endpoints/mcp/localai_assistant_test.go @@ -84,6 +84,10 @@ func (stubClient) ListNodes(_ context.Context) ([]localaitools.Node, error) { return []localaitools.Node{}, nil } +func (stubClient) SetNodeVRAMBudget(_ context.Context, _, _ string) error { + return nil +} + func (stubClient) VRAMEstimate(_ context.Context, _ localaitools.VRAMEstimateRequest) (*vram.EstimateResult, error) { return &vram.EstimateResult{SizeDisplay: "stub"}, nil } diff --git a/core/http/react-ui/src/components/nodes/CapacityEditor.jsx b/core/http/react-ui/src/components/nodes/CapacityEditor.jsx index 9bdfe356a..ffe79b862 100644 --- a/core/http/react-ui/src/components/nodes/CapacityEditor.jsx +++ b/core/http/react-ui/src/components/nodes/CapacityEditor.jsx @@ -1,6 +1,7 @@ import { useState, useEffect, useCallback } from 'react' import { nodesApi } from '../../utils/api' import LoadingSpinner from '../LoadingSpinner' +import { formatVRAM } from './nodeStatus' /** * Inline editor for a node's per-model replica capacity. @@ -81,7 +82,69 @@ export default function CapacityEditor({ node, loadedModelCounts, onUpdate, conf } }, [node, onUpdate, addToast]) + // VRAM allocation budget: a free-form string ("80%" or "12GB") resolved to a + // byte ceiling server-side. Presence of vram_budget signals an active budget. + const currentBudget = node.vram_budget || '' + const hasBudget = !!currentBudget + const budgetBytes = node.vram_budget_bytes || 0 + const [budgetEditing, setBudgetEditing] = useState(false) + const [budgetDraft, setBudgetDraft] = useState(currentBudget) + const [budgetSaving, setBudgetSaving] = useState(false) + const [budgetClearing, setBudgetClearing] = useState(false) + + useEffect(() => { + if (!budgetEditing) setBudgetDraft(currentBudget) + }, [currentBudget, budgetEditing]) + + const cancelBudget = useCallback(() => { + setBudgetEditing(false) + setBudgetDraft(currentBudget) + }, [currentBudget]) + + const saveBudget = useCallback(async () => { + const value = budgetDraft.trim() + if (value === currentBudget) { + setBudgetEditing(false) + return + } + setBudgetSaving(true) + try { + await nodesApi.updateVramBudget(node.id, value) + addToast( + value + ? `VRAM budget set to ${value} on ${node.name}` + : `VRAM budget cleared on ${node.name}`, + 'success', + ) + setBudgetEditing(false) + onUpdate?.(value) + } catch (err) { + addToast(`Could not change VRAM budget: ${err.message || err}`, 'error') + } finally { + setBudgetSaving(false) + } + }, [budgetDraft, currentBudget, node, onUpdate, addToast]) + + const onBudgetKeyDown = (e) => { + if (e.key === 'Enter') { e.preventDefault(); saveBudget() } + else if (e.key === 'Escape') { e.preventDefault(); cancelBudget() } + } + + const clearBudget = useCallback(async () => { + setBudgetClearing(true) + try { + await nodesApi.resetVramBudget(node.id) + addToast(`VRAM budget cleared on ${node.name}`, 'success') + onUpdate?.(null) + } catch (err) { + addToast(`Could not clear VRAM budget: ${err.message || err}`, 'error') + } finally { + setBudgetClearing(false) + } + }, [node, onUpdate, addToast]) + return ( +
@@ -192,5 +255,110 @@ export default function CapacityEditor({ node, loadedModelCounts, onUpdate, conf
+ +
+