diff --git a/core/application/application.go b/core/application/application.go index c5853a92569f..7633c8f8169a 100644 --- a/core/application/application.go +++ b/core/application/application.go @@ -553,12 +553,17 @@ func (a *Application) start() error { // once at startup and reused across chat sessions that opt in via metadata. if !a.applicationConfig.DisableLocalAIAssistant { holder := mcpTools.NewLocalAIAssistantHolder() + var nodeRegistry *nodes.NodeRegistry + if a.distributed != nil { + nodeRegistry = a.distributed.Registry + } assistantClient := localaiInproc.New( a.applicationConfig, a.applicationConfig.SystemState, a.backendLoader, a.modelLoader, a.galleryService, + nodeRegistry, ) // Wire usage tracking so the assistant's get_usage_stats tool // returns real data; nil values keep the tool returning a clear diff --git a/pkg/mcp/localaitools/client.go b/pkg/mcp/localaitools/client.go index 9d76bee5885c..e8b7b2ef5848 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) + ListScheduling(ctx context.Context) ([]ModelSchedulingConfig, error) + GetScheduling(ctx context.Context, modelName string) (*ModelSchedulingConfig, error) + SetScheduling(ctx context.Context, req SetSchedulingRequest) (*ModelSchedulingConfig, error) + DeleteScheduling(ctx context.Context, modelName string) 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. diff --git a/pkg/mcp/localaitools/coverage_test.go b/pkg/mcp/localaitools/coverage_test.go index d068180aacb7..dbe132ba193d 100644 --- a/pkg/mcp/localaitools/coverage_test.go +++ b/pkg/mcp/localaitools/coverage_test.go @@ -35,6 +35,8 @@ var toolToHTTPRoute = map[string]string{ ToolListKnownBackends: "GET /backends/known", ToolSystemInfo: "GET / (welcome JSON)", ToolListNodes: "GET /api/nodes", + ToolListScheduling: "GET /api/nodes/scheduling", + ToolGetScheduling: "GET /api/nodes/scheduling/:model", ToolVRAMEstimate: "POST /api/models/vram-estimate", ToolGetBranding: "GET /api/branding", ToolGetUsageStats: "GET /api/usage (or /api/usage/all when all=true)", @@ -60,6 +62,8 @@ var toolToHTTPRoute = map[string]string{ ToolCreateVoiceProfile: "POST /api/voice-profiles", ToolDeleteVoiceProfile: "DELETE /api/voice-profiles/:id", ToolSetNodeVRAMBudget: "PUT /api/nodes/:id/vram-budget", + ToolSetScheduling: "POST /api/nodes/scheduling", + ToolDeleteScheduling: "DELETE /api/nodes/scheduling/:model", } // allKnownTools is the union of expectedFullCatalog (defined in diff --git a/pkg/mcp/localaitools/dto.go b/pkg/mcp/localaitools/dto.go index 693287647665..51938b7867eb 100644 --- a/pkg/mcp/localaitools/dto.go +++ b/pkg/mcp/localaitools/dto.go @@ -115,6 +115,41 @@ type SetNodeVRAMBudgetRequest struct { 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."` } +// ModelSchedulingConfig is the MCP wire shape for one per-model distributed +// scheduling rule. Keep this DTO explicit instead of aliasing the node-registry +// model so the MCP contract only exposes operator-facing scheduling fields. +type ModelSchedulingConfig struct { + ModelName string `json:"model_name"` + NodeSelector string `json:"node_selector,omitempty"` + MinReplicas int `json:"min_replicas"` + MaxReplicas int `json:"max_replicas"` + SpreadAll bool `json:"spread_all,omitempty"` + RoutePolicy string `json:"route_policy,omitempty"` + BalanceAbsThreshold int `json:"balance_abs_threshold,omitempty"` + BalanceRelThreshold float64 `json:"balance_rel_threshold,omitempty"` + MinPrefixMatch float64 `json:"min_prefix_match,omitempty"` +} + +// SetSchedulingRequest is the input for set_scheduling. It mirrors +// /api/nodes/scheduling so standalone MCP and REST callers preserve the same +// PATCH-style semantics for the optional prefix-cache routing fields. +type SetSchedulingRequest struct { + ModelName string `json:"model_name" jsonschema:"Installed model name whose distributed scheduling rule should be created or updated."` + NodeSelector map[string]string `json:"node_selector,omitempty" jsonschema:"Optional node-label selector. Empty means any healthy backend node."` + MinReplicas int `json:"min_replicas" jsonschema:"Minimum desired replicas. Mutually exclusive with spread_all."` + MaxReplicas int `json:"max_replicas" jsonschema:"Maximum desired replicas. Must be >= min_replicas when non-zero. Mutually exclusive with spread_all."` + SpreadAll bool `json:"spread_all,omitempty" jsonschema:"When true, keep one replica on every matching node. Mutually exclusive with min_replicas/max_replicas."` + RoutePolicy *string `json:"route_policy,omitempty" jsonschema:"Optional prefix-cache route policy override. Omit to preserve the existing value on updates."` + BalanceAbsThreshold *int `json:"balance_abs_threshold,omitempty" jsonschema:"Optional absolute imbalance threshold override. Omit to preserve the existing value on updates."` + BalanceRelThreshold *float64 `json:"balance_rel_threshold,omitempty" jsonschema:"Optional relative imbalance threshold override. Omit to preserve the existing value on updates."` + MinPrefixMatch *float64 `json:"min_prefix_match,omitempty" jsonschema:"Optional minimum prefix match threshold override. Omit to preserve the existing value on updates."` +} + +// DeleteSchedulingRequest identifies the model scheduling rule to remove. +type DeleteSchedulingRequest struct { + ModelName string `json:"model_name" jsonschema:"Installed model name whose scheduling config should be removed."` +} + // 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 ed328352c73b..932f429a88e4 100644 --- a/pkg/mcp/localaitools/fakes_test.go +++ b/pkg/mcp/localaitools/fakes_test.go @@ -42,6 +42,10 @@ type fakeClient struct { upgradeBackend func(string) (string, error) systemInfo func() (*SystemInfo, error) listNodes func() ([]Node, error) + listScheduling func() ([]ModelSchedulingConfig, error) + getScheduling func(string) (*ModelSchedulingConfig, error) + setScheduling func(SetSchedulingRequest) (*ModelSchedulingConfig, error) + deleteScheduling func(string) error setNodeVRAMBudget func(string, string) error vramEstimate func(VRAMEstimateRequest) (*vram.EstimateResult, error) toggleModelState func(string, modeladmin.Action) error @@ -230,6 +234,38 @@ func (f *fakeClient) ListNodes(_ context.Context) ([]Node, error) { return nil, nil } +func (f *fakeClient) ListScheduling(_ context.Context) ([]ModelSchedulingConfig, error) { + f.record("ListScheduling", nil) + if f.listScheduling != nil { + return f.listScheduling() + } + return []ModelSchedulingConfig{}, nil +} + +func (f *fakeClient) GetScheduling(_ context.Context, modelName string) (*ModelSchedulingConfig, error) { + f.record("GetScheduling", modelName) + if f.getScheduling != nil { + return f.getScheduling(modelName) + } + return &ModelSchedulingConfig{ModelName: modelName}, nil +} + +func (f *fakeClient) SetScheduling(_ context.Context, req SetSchedulingRequest) (*ModelSchedulingConfig, error) { + f.record("SetScheduling", req) + if f.setScheduling != nil { + return f.setScheduling(req) + } + return &ModelSchedulingConfig{ModelName: req.ModelName, MinReplicas: req.MinReplicas, MaxReplicas: req.MaxReplicas, SpreadAll: req.SpreadAll}, nil +} + +func (f *fakeClient) DeleteScheduling(_ context.Context, modelName string) error { + f.record("DeleteScheduling", modelName) + if f.deleteScheduling != nil { + return f.deleteScheduling(modelName) + } + return nil +} + func (f *fakeClient) SetNodeVRAMBudget(_ context.Context, nodeID, budget string) error { f.record("SetNodeVRAMBudget", []any{nodeID, budget}) if f.setNodeVRAMBudget != nil { diff --git a/pkg/mcp/localaitools/httpapi/client.go b/pkg/mcp/localaitools/httpapi/client.go index 8029655caa29..550c1ef0170d 100644 --- a/pkg/mcp/localaitools/httpapi/client.go +++ b/pkg/mcp/localaitools/httpapi/client.go @@ -483,6 +483,49 @@ func (c *Client) ListNodes(ctx context.Context) ([]localaitools.Node, error) { return out, nil } +func (c *Client) ListScheduling(ctx context.Context) ([]localaitools.ModelSchedulingConfig, error) { + var out []localaitools.ModelSchedulingConfig + if err := c.do(ctx, http.MethodGet, routeScheduling, nil, &out); err != nil { + if errors.Is(err, ErrHTTPNotFound) { + return []localaitools.ModelSchedulingConfig{}, nil + } + return nil, err + } + return out, nil +} + +func (c *Client) GetScheduling(ctx context.Context, modelName string) (*localaitools.ModelSchedulingConfig, error) { + if modelName == "" { + return nil, errors.New("model_name is required") + } + var out localaitools.ModelSchedulingConfig + if err := c.do(ctx, http.MethodGet, routeModelScheduling(modelName), nil, &out); err != nil { + if errors.Is(err, ErrHTTPNotFound) { + return nil, nil + } + return nil, err + } + return &out, nil +} + +func (c *Client) SetScheduling(ctx context.Context, req localaitools.SetSchedulingRequest) (*localaitools.ModelSchedulingConfig, error) { + if req.ModelName == "" { + return nil, errors.New("model_name is required") + } + var out localaitools.ModelSchedulingConfig + if err := c.do(ctx, http.MethodPost, routeScheduling, req, &out); err != nil { + return nil, err + } + return &out, nil +} + +func (c *Client) DeleteScheduling(ctx context.Context, modelName string) error { + if modelName == "" { + return errors.New("model_name is required") + } + return c.do(ctx, http.MethodDelete, routeModelScheduling(modelName), nil, 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. diff --git a/pkg/mcp/localaitools/httpapi/client_test.go b/pkg/mcp/localaitools/httpapi/client_test.go index 38eabf2832b7..c4e596650ed7 100644 --- a/pkg/mcp/localaitools/httpapi/client_test.go +++ b/pkg/mcp/localaitools/httpapi/client_test.go @@ -84,6 +84,43 @@ func fakeLocalAI() *httptest.Server { }) }) + mux.HandleFunc("/api/nodes/scheduling", func(w http.ResponseWriter, r *http.Request) { + switch r.Method { + case http.MethodGet: + _ = json.NewEncoder(w).Encode([]map[string]any{{ + "id": "sched-1", + "model_name": "qwen", + "min_replicas": 1, + "max_replicas": 2, + "unsatisfiable_ticks": 1, + "unsatisfiable_until": "2026-01-01T00:00:00Z", + "created_at": "2026-01-01T00:00:00Z", + "updated_at": "2026-01-01T00:00:00Z", + }}) + case http.MethodPost: + var body map[string]any + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + body["id"] = "sched-1" + _ = json.NewEncoder(w).Encode(body) + default: + http.Error(w, "method", http.StatusMethodNotAllowed) + } + }) + + mux.HandleFunc("/api/nodes/scheduling/qwen", func(w http.ResponseWriter, r *http.Request) { + switch r.Method { + case http.MethodGet: + _ = json.NewEncoder(w).Encode(map[string]any{"model_name": "qwen", "spread_all": true}) + case http.MethodDelete: + w.WriteHeader(http.StatusNoContent) + default: + http.Error(w, "method", http.StatusMethodNotAllowed) + } + }) + return httptest.NewServer(mux) } @@ -197,8 +234,51 @@ var _ = Describe("httpapi.Client against the LocalAI admin REST surface", func() Expect(bs[0].Installed).To(BeTrue()) }) }) + + Describe("Scheduling", func() { + It("lists scheduling configs", func() { + out, err := c.ListScheduling(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(out).To(HaveLen(1)) + Expect(out[0].ModelName).To(Equal("qwen")) + Expect(out[0].MinReplicas).To(Equal(1)) + Expect(schedulingJSONKeys(&out[0])).ToNot(Or( + HaveKey("id"), + HaveKey("unsatisfiable_until"), + HaveKey("unsatisfiable_ticks"), + HaveKey("created_at"), + HaveKey("updated_at"), + )) + }) + + It("gets one scheduling config", func() { + out, err := c.GetScheduling(ctx, "qwen") + Expect(err).ToNot(HaveOccurred()) + Expect(out.ModelName).To(Equal("qwen")) + Expect(out.SpreadAll).To(BeTrue()) + }) + + It("sets a scheduling config", func() { + out, err := c.SetScheduling(ctx, localaitools.SetSchedulingRequest{ModelName: "qwen", MinReplicas: 1, MaxReplicas: 2}) + Expect(err).ToNot(HaveOccurred()) + Expect(out.ModelName).To(Equal("qwen")) + Expect(out.MaxReplicas).To(Equal(2)) + }) + + It("deletes a scheduling config", func() { + Expect(c.DeleteScheduling(ctx, "qwen")).To(Succeed()) + }) + }) }) +func schedulingJSONKeys(config *localaitools.ModelSchedulingConfig) map[string]any { + var out map[string]any + b, err := json.Marshal(config) + Expect(err).ToNot(HaveOccurred()) + Expect(json.Unmarshal(b, &out)).To(Succeed()) + return out +} + var _ = Describe("Model aliases", func() { Describe("ListAliases", func() { It("passes the GET /api/aliases payload through unchanged", func() { diff --git a/pkg/mcp/localaitools/httpapi/routes.go b/pkg/mcp/localaitools/httpapi/routes.go index 2b8a107ddc28..fca25efad741 100644 --- a/pkg/mcp/localaitools/httpapi/routes.go +++ b/pkg/mcp/localaitools/httpapi/routes.go @@ -24,6 +24,7 @@ const ( routeBackendsKnown = "/backends/known" routeBackendsApply = "/backends/apply" routeNodes = "/api/nodes" + routeScheduling = "/api/nodes/scheduling" routeVRAMEstimate = "/api/models/vram-estimate" routeBranding = "/api/branding" routeSettings = "/api/settings" @@ -66,3 +67,7 @@ func routeVoiceProfileDelete(id string) string { func routeNodeVRAMBudget(id string) string { return "/api/nodes/" + url.PathEscape(id) + "/vram-budget" } + +func routeModelScheduling(modelName string) string { + return "/api/nodes/scheduling/" + url.PathEscape(modelName) +} diff --git a/pkg/mcp/localaitools/inproc/client.go b/pkg/mcp/localaitools/inproc/client.go index d3054f614ba0..759cd5bbadec 100644 --- a/pkg/mcp/localaitools/inproc/client.go +++ b/pkg/mcp/localaitools/inproc/client.go @@ -23,6 +23,8 @@ import ( "github.com/mudler/LocalAI/core/schema" "github.com/mudler/LocalAI/core/services/galleryop" "github.com/mudler/LocalAI/core/services/modeladmin" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/nodes/prefixcache" "github.com/mudler/LocalAI/core/services/routing/billing" "github.com/mudler/LocalAI/core/services/routing/pii" "github.com/mudler/LocalAI/core/services/routing/router" @@ -47,6 +49,7 @@ type Client struct { ConfigLoader *config.ModelConfigLoader ModelLoader *model.ModelLoader Gallery *galleryop.GalleryService + NodeRegistry *nodes.NodeRegistry VoiceProfiles *voiceprofile.Store // StatsRecorder and FallbackUser are optional — they back the @@ -75,13 +78,18 @@ type Client struct { // except ModelLoader (used only for SystemInfo's loaded-models report and // best-effort ShutdownModel calls during config edits) and the stats // fields (StatsRecorder, FallbackUser) which gate get_usage_stats. -func New(appConfig *config.ApplicationConfig, systemState *system.SystemState, cl *config.ModelConfigLoader, ml *model.ModelLoader, gs *galleryop.GalleryService) *Client { +func New(appConfig *config.ApplicationConfig, systemState *system.SystemState, cl *config.ModelConfigLoader, ml *model.ModelLoader, gs *galleryop.GalleryService, registries ...*nodes.NodeRegistry) *Client { + var registry *nodes.NodeRegistry + if len(registries) > 0 { + registry = registries[0] + } return &Client{ AppConfig: appConfig, SystemState: systemState, ConfigLoader: cl, ModelLoader: ml, Gallery: gs, + NodeRegistry: registry, VoiceProfiles: voiceprofile.NewStore(appConfig.DataPath), modelAdmin: modeladmin.NewConfigService(cl, appConfig), } @@ -501,6 +509,121 @@ func (c *Client) ListNodes(_ context.Context) ([]localaitools.Node, error) { return []localaitools.Node{}, nil } +func (c *Client) ListScheduling(ctx context.Context) ([]localaitools.ModelSchedulingConfig, error) { + if c.NodeRegistry == nil { + return []localaitools.ModelSchedulingConfig{}, nil + } + configs, err := c.NodeRegistry.ListModelSchedulings(ctx) + if err != nil { + return nil, err + } + return localaitools.SchedulingConfigsFromNodes(configs), nil +} + +func (c *Client) GetScheduling(ctx context.Context, modelName string) (*localaitools.ModelSchedulingConfig, error) { + if modelName == "" { + return nil, errors.New("model_name is required") + } + if c.NodeRegistry == nil { + return nil, errors.New("model scheduling is only available in distributed mode") + } + config, err := c.NodeRegistry.GetModelScheduling(ctx, modelName) + if err != nil { + return nil, err + } + if config == nil { + return nil, nil + } + out := localaitools.SchedulingConfigFromNode(*config) + return &out, nil +} + +func (c *Client) SetScheduling(ctx context.Context, req localaitools.SetSchedulingRequest) (*localaitools.ModelSchedulingConfig, error) { + if req.ModelName == "" { + return nil, errors.New("model_name is required") + } + if c.NodeRegistry == nil { + return nil, errors.New("model scheduling is only available in distributed mode") + } + + existing, err := c.NodeRegistry.GetModelScheduling(ctx, req.ModelName) + if err != nil { + return nil, fmt.Errorf("load existing scheduling config: %w", err) + } + routePolicy := "" + absThr := 0 + relThr := 0.0 + minMatch := 0.0 + if existing != nil { + routePolicy = existing.RoutePolicy + absThr = existing.BalanceAbsThreshold + relThr = existing.BalanceRelThreshold + minMatch = existing.MinPrefixMatch + } + if req.RoutePolicy != nil { + routePolicy = *req.RoutePolicy + } + if req.BalanceAbsThreshold != nil { + absThr = *req.BalanceAbsThreshold + } + if req.BalanceRelThreshold != nil { + relThr = *req.BalanceRelThreshold + } + if req.MinPrefixMatch != nil { + minMatch = *req.MinPrefixMatch + } + if req.SpreadAll && (req.MinReplicas != 0 || req.MaxReplicas != 0) { + return nil, errors.New("spread_all and min_replicas/max_replicas are mutually exclusive") + } + if req.MinReplicas < 0 { + return nil, errors.New("min_replicas must be >= 0") + } + if req.MaxReplicas < 0 { + return nil, errors.New("max_replicas must be >= 0") + } + if req.MaxReplicas > 0 && req.MinReplicas > req.MaxReplicas { + return nil, errors.New("min_replicas must be <= max_replicas") + } + if err := prefixcache.ValidateThresholds(routePolicy, absThr, relThr, minMatch); err != nil { + return nil, err + } + + var selectorJSON string + if len(req.NodeSelector) > 0 { + b, err := json.Marshal(req.NodeSelector) + if err != nil { + return nil, fmt.Errorf("invalid node_selector: %w", err) + } + selectorJSON = string(b) + } + config := &nodes.ModelSchedulingConfig{ + ModelName: req.ModelName, + NodeSelector: selectorJSON, + MinReplicas: req.MinReplicas, + MaxReplicas: req.MaxReplicas, + SpreadAll: req.SpreadAll, + RoutePolicy: routePolicy, + BalanceAbsThreshold: absThr, + BalanceRelThreshold: relThr, + MinPrefixMatch: minMatch, + } + if err := c.NodeRegistry.SetModelScheduling(ctx, config); err != nil { + return nil, err + } + out := localaitools.SchedulingConfigFromNode(*config) + return &out, nil +} + +func (c *Client) DeleteScheduling(ctx context.Context, modelName string) error { + if modelName == "" { + return errors.New("model_name is required") + } + if c.NodeRegistry == nil { + return errors.New("model scheduling is only available in distributed mode") + } + return c.NodeRegistry.DeleteModelScheduling(ctx, modelName) +} + 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 diff --git a/pkg/mcp/localaitools/inproc/client_test.go b/pkg/mcp/localaitools/inproc/client_test.go index e385897c772d..38a05a8a7247 100644 --- a/pkg/mcp/localaitools/inproc/client_test.go +++ b/pkg/mcp/localaitools/inproc/client_test.go @@ -2,6 +2,7 @@ package inproc import ( "context" + "encoding/json" "errors" "os" "path/filepath" @@ -13,8 +14,11 @@ import ( "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/gallery" "github.com/mudler/LocalAI/core/services/galleryop" + "github.com/mudler/LocalAI/core/services/nodes" localaitools "github.com/mudler/LocalAI/pkg/mcp/localaitools" "github.com/mudler/LocalAI/pkg/system" + "gorm.io/driver/sqlite" + "gorm.io/gorm" ) // Regression spec for the bug we fixed when channel sends were @@ -124,3 +128,88 @@ var _ = Describe("inproc.Client model aliases", func() { }) }) }) + +var _ = Describe("inproc.Client model scheduling", func() { + var ( + ctx context.Context + registry *nodes.NodeRegistry + c *Client + stringp = func(s string) *string { return &s } + floatp = func(f float64) *float64 { return &f } + ) + + BeforeEach(func() { + ctx = context.Background() + db, err := gorm.Open(sqlite.Open(filepath.Join(GinkgoT().TempDir(), "nodes.db")), &gorm.Config{}) + Expect(err).ToNot(HaveOccurred()) + registry, err = nodes.NewNodeRegistry(db) + Expect(err).ToNot(HaveOccurred()) + + tempDir := GinkgoT().TempDir() + systemState, err := system.GetSystemState(system.WithModelPath(tempDir)) + Expect(err).ToNot(HaveOccurred()) + appConfig := config.NewApplicationConfig(config.WithSystemState(systemState)) + c = New(appConfig, systemState, config.NewModelConfigLoader(tempDir), nil, nil, registry) + }) + + It("sets, lists, gets, merges, and deletes scheduling configs through the registry", func() { + created, err := c.SetScheduling(ctx, localaitools.SetSchedulingRequest{ + ModelName: "qwen", + NodeSelector: map[string]string{"gpu": "nvidia"}, + MinReplicas: 1, + MaxReplicas: 2, + RoutePolicy: stringp("prefix_cache"), + MinPrefixMatch: floatp(0.4), + }) + Expect(err).ToNot(HaveOccurred()) + Expect(created.ModelName).To(Equal("qwen")) + Expect(created.NodeSelector).To(Equal(`{"gpu":"nvidia"}`)) + Expect(created.RoutePolicy).To(Equal("prefix_cache")) + Expect(created.MinPrefixMatch).To(Equal(0.4)) + Expect(schedulingJSONKeys(created)).ToNot(Or( + HaveKey("id"), + HaveKey("unsatisfiable_until"), + HaveKey("unsatisfiable_ticks"), + HaveKey("created_at"), + HaveKey("updated_at"), + )) + + listed, err := c.ListScheduling(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(listed).To(HaveLen(1)) + Expect(listed[0].ModelName).To(Equal("qwen")) + + updated, err := c.SetScheduling(ctx, localaitools.SetSchedulingRequest{ + ModelName: "qwen", + MinReplicas: 2, + MaxReplicas: 3, + }) + Expect(err).ToNot(HaveOccurred()) + Expect(updated.MinReplicas).To(Equal(2)) + Expect(updated.MaxReplicas).To(Equal(3)) + Expect(updated.RoutePolicy).To(Equal("prefix_cache")) + Expect(updated.MinPrefixMatch).To(Equal(0.4)) + + got, err := c.GetScheduling(ctx, "qwen") + Expect(err).ToNot(HaveOccurred()) + Expect(got).ToNot(BeNil()) + Expect(got.MaxReplicas).To(Equal(3)) + + Expect(c.DeleteScheduling(ctx, "qwen")).To(Succeed()) + missing, err := c.GetScheduling(ctx, "qwen") + Expect(err).ToNot(HaveOccurred()) + Expect(missing).To(BeNil()) + }) +}) + +func schedulingJSONKeys(config *localaitools.ModelSchedulingConfig) map[string]any { + var out map[string]any + Expect(json.Unmarshal([]byte(mustMarshal(config)), &out)).To(Succeed()) + return out +} + +func mustMarshal(v any) string { + b, err := json.Marshal(v) + Expect(err).ToNot(HaveOccurred()) + return string(b) +} diff --git a/pkg/mcp/localaitools/prompts/10_safety.md b/pkg/mcp/localaitools/prompts/10_safety.md index 3f338f32083c..bd5941808bb2 100644 --- a/pkg/mcp/localaitools/prompts/10_safety.md +++ b/pkg/mcp/localaitools/prompts/10_safety.md @@ -2,7 +2,7 @@ These rules are non-negotiable. The user trusts you to operate their server without unintended changes. -1. **Confirm before mutating.** Before calling any of these tools — `install_model`, `import_model_uri`, `delete_model`, `install_backend`, `upgrade_backend`, `edit_model_config`, `reload_models`, `load_model`, `toggle_model_state`, `toggle_model_pinned`, `create_voice_profile`, `delete_voice_profile` — first state in plain language what you are about to do (which tool, which target, which arguments) and wait for the user's explicit confirmation in the next turn. "Yes", "do it", "go ahead", "proceed" all count as confirmation. Anything else does not. +1. **Confirm before mutating.** Before calling any of these tools — `install_model`, `import_model_uri`, `delete_model`, `install_backend`, `upgrade_backend`, `edit_model_config`, `reload_models`, `load_model`, `toggle_model_state`, `toggle_model_pinned`, `create_voice_profile`, `delete_voice_profile`, `set_node_vram_budget`, `set_scheduling`, `delete_scheduling` — first state in plain language what you are about to do (which tool, which target, which arguments) and wait for the user's explicit confirmation in the next turn. "Yes", "do it", "go ahead", "proceed" all count as confirmation. Anything else does not. 2. **Disambiguate before mutating.** If the user's request is ambiguous (several gallery candidates match, the model name has multiple installed versions, the backend has variants), present the candidates as a numbered list and ask the user to pick before calling any mutating tool. diff --git a/pkg/mcp/localaitools/prompts/20_tools.md b/pkg/mcp/localaitools/prompts/20_tools.md index a6cdc9b396f7..0a2d7dddcac4 100644 --- a/pkg/mcp/localaitools/prompts/20_tools.md +++ b/pkg/mcp/localaitools/prompts/20_tools.md @@ -14,6 +14,8 @@ The MCP `tools/list` endpoint also exposes the full input schema for each of the - `vram_estimate` — Estimate VRAM use for a model under a given config. - `system_info` — LocalAI version, paths, distributed flag, loaded models, installed backends. - `list_nodes` — List federated worker nodes (only useful in distributed mode). +- `list_scheduling` — List distributed per-model scheduling configs. +- `get_scheduling` — Read the distributed scheduling config for one model. - `list_voice_profiles` — List reusable voice-cloning profiles and their stable TTS voice URIs. ## Mutating (require user confirmation per safety rule 1) @@ -30,3 +32,6 @@ The MCP `tools/list` endpoint also exposes the full input schema for each of the - `toggle_model_pinned` — Pin or unpin a model (`action`: `pin` or `unpin`). - `create_voice_profile` — Save a consent-confirmed base64 PCM-WAV reference and exact transcript for reuse in TTS. - `delete_voice_profile` — Permanently delete a saved voice profile by UUID. +- `set_node_vram_budget` — Set or clear a federated node's VRAM budget override. +- `set_scheduling` — Create or update a distributed per-model scheduling config. +- `delete_scheduling` — Remove a distributed per-model scheduling config. diff --git a/pkg/mcp/localaitools/prompts/skills/manage_distributed_scheduling.md b/pkg/mcp/localaitools/prompts/skills/manage_distributed_scheduling.md new file mode 100644 index 000000000000..bbe5b4ee0c1c --- /dev/null +++ b/pkg/mcp/localaitools/prompts/skills/manage_distributed_scheduling.md @@ -0,0 +1,9 @@ +# Skill: Manage distributed scheduling + +Use this when the user asks to inspect, set, or remove distributed per-model scheduling rules. + +1. Call `system_info` first. If `distributed` is false, explain that scheduling tools only affect distributed deployments. +2. For inspection, call `list_scheduling` or `get_scheduling` and summarize model name, selector, replica bounds, spread-all, and route policy fields. +3. Before calling `set_scheduling` or `delete_scheduling`, follow safety rule 1 and wait for explicit confirmation. +4. For `set_scheduling`, never combine `spread_all: true` with non-zero `min_replicas` or `max_replicas`. +5. After a confirmed mutation, call `get_scheduling` for that model and summarize the persisted config. diff --git a/pkg/mcp/localaitools/prompts/skills/system_status.md b/pkg/mcp/localaitools/prompts/skills/system_status.md index 0ee7e6f7b3dd..7f7146344a79 100644 --- a/pkg/mcp/localaitools/prompts/skills/system_status.md +++ b/pkg/mcp/localaitools/prompts/skills/system_status.md @@ -4,11 +4,12 @@ Use this when the user asks "what's installed?", "what's running?", "show status 1. Call `system_info` for version, paths, distributed flag, loaded models, installed backends. 2. Call `list_installed_models` (no capability filter) for the full installed-model inventory. -3. If `system_info.distributed` is true, also call `list_nodes` and report worker health. +3. If `system_info.distributed` is true, also call `list_nodes` and `list_scheduling`, then report worker health and any per-model scheduling rules. 4. Present a concise summary: - **Version & mode** (`distributed: true|false`) - **Installed models** (count + list, each with name and capabilities) - **Installed backends** (count + list) - **Loaded right now** (from `loaded_models`) - **Workers** (only when distributed) + - **Scheduling rules** (only when distributed) 5. Do not call mutating tools in this skill. diff --git a/pkg/mcp/localaitools/scheduling.go b/pkg/mcp/localaitools/scheduling.go new file mode 100644 index 000000000000..c5fc80aa3f35 --- /dev/null +++ b/pkg/mcp/localaitools/scheduling.go @@ -0,0 +1,25 @@ +package localaitools + +import "github.com/mudler/LocalAI/core/services/nodes" + +func SchedulingConfigFromNode(config nodes.ModelSchedulingConfig) ModelSchedulingConfig { + return ModelSchedulingConfig{ + ModelName: config.ModelName, + NodeSelector: config.NodeSelector, + MinReplicas: config.MinReplicas, + MaxReplicas: config.MaxReplicas, + SpreadAll: config.SpreadAll, + RoutePolicy: config.RoutePolicy, + BalanceAbsThreshold: config.BalanceAbsThreshold, + BalanceRelThreshold: config.BalanceRelThreshold, + MinPrefixMatch: config.MinPrefixMatch, + } +} + +func SchedulingConfigsFromNodes(configs []nodes.ModelSchedulingConfig) []ModelSchedulingConfig { + out := make([]ModelSchedulingConfig, 0, len(configs)) + for _, config := range configs { + out = append(out, SchedulingConfigFromNode(config)) + } + return out +} diff --git a/pkg/mcp/localaitools/server.go b/pkg/mcp/localaitools/server.go index 70bf6c15e8ef..8a6c64b868a6 100644 --- a/pkg/mcp/localaitools/server.go +++ b/pkg/mcp/localaitools/server.go @@ -47,6 +47,7 @@ func NewServer(client LocalAIClient, opts Options) *mcp.Server { registerBackendTools(srv, client, opts) registerConfigTools(srv, client, opts) registerSystemTools(srv, client, opts) + registerSchedulingTools(srv, client, opts) registerStateTools(srv, client, opts) registerBrandingTools(srv, client, opts) registerVoiceProfileTools(srv, client, opts) diff --git a/pkg/mcp/localaitools/server_test.go b/pkg/mcp/localaitools/server_test.go index a47b84ba3bc3..ce1ede915950 100644 --- a/pkg/mcp/localaitools/server_test.go +++ b/pkg/mcp/localaitools/server_test.go @@ -92,6 +92,7 @@ var expectedFullCatalog = sortedStrings( ToolListInstalledModels, ToolListKnownBackends, ToolListNodes, + ToolListScheduling, ToolListVoiceProfiles, ToolLoadModel, ToolReloadModels, @@ -105,6 +106,9 @@ var expectedFullCatalog = sortedStrings( ToolCreateVoiceProfile, ToolDeleteVoiceProfile, ToolSetNodeVRAMBudget, + ToolSetScheduling, + ToolDeleteScheduling, + ToolGetScheduling, ) // expectedReadOnlyCatalog is the tool set when DisableMutating=true. Sorted. @@ -123,6 +127,8 @@ var expectedReadOnlyCatalog = sortedStrings( ToolListInstalledModels, ToolListKnownBackends, ToolListNodes, + ToolListScheduling, + ToolGetScheduling, ToolListVoiceProfiles, ToolSystemInfo, ToolVRAMEstimate, @@ -165,6 +171,10 @@ var _ = Describe("Tool dispatch", func() { {ToolListKnownBackends, struct{}{}, "ListKnownBackends"}, {ToolSystemInfo, struct{}{}, "SystemInfo"}, {ToolListNodes, struct{}{}, "ListNodes"}, + {ToolListScheduling, struct{}{}, "ListScheduling"}, + {ToolGetScheduling, DeleteSchedulingRequest{ModelName: "qwen"}, "GetScheduling"}, + {ToolSetScheduling, SetSchedulingRequest{ModelName: "qwen", MinReplicas: 1, MaxReplicas: 2}, "SetScheduling"}, + {ToolDeleteScheduling, DeleteSchedulingRequest{ModelName: "qwen"}, "DeleteScheduling"}, {ToolListVoiceProfiles, struct{}{}, "ListVoiceProfiles"}, {ToolInstallModel, InstallModelRequest{ModelName: "test/foo"}, "InstallModel"}, {ToolImportModelURI, ImportModelURIRequest{URI: "Qwen/Qwen3-4B-GGUF"}, "ImportModelURI"}, @@ -219,6 +229,40 @@ var _ = Describe("Tool error surfacing", func() { }) }) +var _ = Describe("Scheduling tool behavior", func() { + It("round-trips list/get/set/delete over the MCP surface", func() { + fc := &fakeClient{ + listScheduling: func() ([]ModelSchedulingConfig, error) { + return []ModelSchedulingConfig{{ModelName: "qwen", MinReplicas: 1, MaxReplicas: 2}}, nil + }, + getScheduling: func(modelName string) (*ModelSchedulingConfig, error) { + return &ModelSchedulingConfig{ModelName: modelName, SpreadAll: true}, nil + }, + setScheduling: func(req SetSchedulingRequest) (*ModelSchedulingConfig, error) { + return &ModelSchedulingConfig{ModelName: req.ModelName, MinReplicas: req.MinReplicas, MaxReplicas: req.MaxReplicas}, nil + }, + } + ctx, sess, done := connectInMemory(fc, Options{}) + DeferCleanup(done) + + listRes := callTool(ctx, sess, ToolListScheduling, struct{}{}) + Expect(listRes.IsError).To(BeFalse(), resultText(listRes)) + Expect(resultText(listRes)).To(ContainSubstring(`"model_name": "qwen"`)) + + getRes := callTool(ctx, sess, ToolGetScheduling, DeleteSchedulingRequest{ModelName: "qwen"}) + Expect(getRes.IsError).To(BeFalse(), resultText(getRes)) + Expect(resultText(getRes)).To(ContainSubstring(`"spread_all": true`)) + + setRes := callTool(ctx, sess, ToolSetScheduling, SetSchedulingRequest{ModelName: "qwen", MinReplicas: 1, MaxReplicas: 2}) + Expect(setRes.IsError).To(BeFalse(), resultText(setRes)) + Expect(resultText(setRes)).To(ContainSubstring(`"max_replicas": 2`)) + + deleteRes := callTool(ctx, sess, ToolDeleteScheduling, DeleteSchedulingRequest{ModelName: "qwen"}) + Expect(deleteRes.IsError).To(BeFalse(), resultText(deleteRes)) + Expect(resultText(deleteRes)).To(ContainSubstring(`"deleted": "qwen"`)) + }) +}) + var _ = Describe("Argument validation", func() { type validationCase struct { desc string @@ -235,6 +279,9 @@ var _ = Describe("Argument validation", func() { {"toggle_model_state rejects unknown action", ToolToggleModelState, map[string]any{"name": "foo", "action": "noop"}, "action must be one of"}, {"edit_model_config rejects empty patch", ToolEditModelConfig, map[string]any{"name": "foo", "patch": map[string]any{}}, "patch is required"}, {"create_voice_profile requires consent", ToolCreateVoiceProfile, CreateVoiceProfileRequest{Name: "Voice", Transcript: "words", AudioBase64: "UklGRg=="}, "consent_confirmed must be true"}, + {"set_scheduling requires model_name", ToolSetScheduling, SetSchedulingRequest{}, "model_name is required"}, + {"set_scheduling rejects invalid replica range", ToolSetScheduling, SetSchedulingRequest{ModelName: "qwen", MinReplicas: 3, MaxReplicas: 1}, "min_replicas must be <= max_replicas"}, + {"delete_scheduling requires model_name", ToolDeleteScheduling, DeleteSchedulingRequest{}, "model_name is required"}, } for _, c := range cases { diff --git a/pkg/mcp/localaitools/tools.go b/pkg/mcp/localaitools/tools.go index 0d6890e34680..89d99432bde3 100644 --- a/pkg/mcp/localaitools/tools.go +++ b/pkg/mcp/localaitools/tools.go @@ -17,6 +17,8 @@ const ( ToolListKnownBackends = "list_known_backends" ToolSystemInfo = "system_info" ToolListNodes = "list_nodes" + ToolListScheduling = "list_scheduling" + ToolGetScheduling = "get_scheduling" ToolVRAMEstimate = "vram_estimate" ToolGetBranding = "get_branding" ToolGetUsageStats = "get_usage_stats" @@ -42,6 +44,8 @@ const ( ToolCreateVoiceProfile = "create_voice_profile" ToolDeleteVoiceProfile = "delete_voice_profile" ToolSetNodeVRAMBudget = "set_node_vram_budget" + ToolSetScheduling = "set_scheduling" + ToolDeleteScheduling = "delete_scheduling" // 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_scheduling.go b/pkg/mcp/localaitools/tools_scheduling.go new file mode 100644 index 000000000000..ed317ddc789b --- /dev/null +++ b/pkg/mcp/localaitools/tools_scheduling.go @@ -0,0 +1,77 @@ +package localaitools + +import ( + "context" + + "github.com/modelcontextprotocol/go-sdk/mcp" +) + +func registerSchedulingTools(s *mcp.Server, client LocalAIClient, opts Options) { + mcp.AddTool(s, &mcp.Tool{ + Name: ToolListScheduling, + Description: "List distributed per-model scheduling configs (only meaningful in distributed mode).", + }, func(ctx context.Context, _ *mcp.CallToolRequest, _ struct{}) (*mcp.CallToolResult, any, error) { + configs, err := client.ListScheduling(ctx) + if err != nil { + return errorResult(err), nil, nil + } + return jsonResult(configs), nil, nil + }) + + mcp.AddTool(s, &mcp.Tool{ + Name: ToolGetScheduling, + Description: "Get the distributed scheduling config for one model, or null when none is configured.", + }, func(ctx context.Context, _ *mcp.CallToolRequest, args DeleteSchedulingRequest) (*mcp.CallToolResult, any, error) { + if args.ModelName == "" { + return errorResultf("model_name is required"), nil, nil + } + config, err := client.GetScheduling(ctx, args.ModelName) + if err != nil { + return errorResult(err), nil, nil + } + return jsonResult(config), nil, nil + }) + + if opts.DisableMutating { + return + } + + mcp.AddTool(s, &mcp.Tool{ + Name: ToolSetScheduling, + Description: "Create or update a distributed per-model scheduling config. Requires user confirmation per safety rule 1.", + }, func(ctx context.Context, _ *mcp.CallToolRequest, args SetSchedulingRequest) (*mcp.CallToolResult, any, error) { + if args.ModelName == "" { + return errorResultf("model_name is required"), nil, nil + } + if args.SpreadAll && (args.MinReplicas != 0 || args.MaxReplicas != 0) { + return errorResultf("spread_all and min_replicas/max_replicas are mutually exclusive"), nil, nil + } + if args.MinReplicas < 0 { + return errorResultf("min_replicas must be >= 0"), nil, nil + } + if args.MaxReplicas < 0 { + return errorResultf("max_replicas must be >= 0"), nil, nil + } + if args.MaxReplicas > 0 && args.MinReplicas > args.MaxReplicas { + return errorResultf("min_replicas must be <= max_replicas"), nil, nil + } + config, err := client.SetScheduling(ctx, args) + if err != nil { + return errorResult(err), nil, nil + } + return jsonResult(config), nil, nil + }) + + mcp.AddTool(s, &mcp.Tool{ + Name: ToolDeleteScheduling, + Description: "Delete a distributed per-model scheduling config. Requires user confirmation per safety rule 1.", + }, func(ctx context.Context, _ *mcp.CallToolRequest, args DeleteSchedulingRequest) (*mcp.CallToolResult, any, error) { + if args.ModelName == "" { + return errorResultf("model_name is required"), nil, nil + } + if err := client.DeleteScheduling(ctx, args.ModelName); err != nil { + return errorResult(err), nil, nil + } + return jsonResult(map[string]string{"deleted": args.ModelName}), nil, nil + }) +}