Compare commits

..

1 Commits

Author SHA1 Message Date
dependabot[bot]
0de0f8fcbb chore(deps): bump the pip group across 4 directories with 1 update
Bumps the pip group with 1 update in the /backend/python/ace-step directory: torch.
Bumps the pip group with 1 update in the /backend/python/llama-cpp-quantization directory: torch.
Bumps the pip group with 1 update in the /backend/python/sglang directory: torch.
Bumps the pip group with 1 update in the /backend/python/vllm-omni directory: torch.


Updates `torch` from 2.10.0+rocm7.0 to 2.12.0+cpu

Updates `torch` from 2.10.0+rocm7.0 to 2.12.0+cpu

Updates `torch` from 2.10.0+rocm7.0 to 2.12.0+cpu

Updates `torch` from 2.10.0+rocm7.0 to 2.12.0+cpu

Updates `torch` from 2.10.0 to 2.12.0+cpu

Updates `torch` from 2.10.0 to 2.12.0+cpu

Updates `torch` from 2.10.0 to 2.12.0+cpu

Updates `torch` from 2.10.0 to 2.12.0+cpu

Updates `torch` from 2.9.0 to 2.12.0+cpu

Updates `torch` from 2.9.0 to 2.12.0+cpu

Updates `torch` from 2.9.0 to 2.12.0+cpu

Updates `torch` from 2.9.0 to 2.12.0+cpu

Updates `torch` from 2.7.0 to 2.12.0+cu130

Updates `torch` from 2.7.0 to 2.12.0+cu130

Updates `torch` from 2.7.0 to 2.12.0+cu130

Updates `torch` from 2.7.0 to 2.12.0+cu130

---
updated-dependencies:
- dependency-name: torch
  dependency-version: 2.12.0+cpu
  dependency-type: direct:production
- dependency-name: torch
  dependency-version: 2.12.0+cpu
  dependency-type: direct:production
- dependency-name: torch
  dependency-version: 2.12.0+cpu
  dependency-type: direct:production
- dependency-name: torch
  dependency-version: 2.12.0+cu130
  dependency-type: direct:production
...

Signed-off-by: dependabot[bot] <support@github.com>
2026-06-27 08:10:29 +00:00
51 changed files with 225 additions and 2608 deletions

View File

@@ -4991,6 +4991,9 @@ includeDarwin:
- backend: "qwen-tts"
tag-suffix: "-metal-darwin-arm64-qwen-tts"
build-type: "mps"
- backend: "fish-speech"
tag-suffix: "-metal-darwin-arm64-fish-speech"
build-type: "mps"
- backend: "voxcpm"
tag-suffix: "-metal-darwin-arm64-voxcpm"
build-type: "mps"

View File

@@ -13,14 +13,8 @@ if [ "$(uname)" != "Darwin" ]; then
fi
if [ "$(uname)" = "Darwin" ]; then
# macOS: single fallback variant (Metal/Accelerate). The cmake build emits a
# Mach-O named .so, but tolerate .dylib too — pick whichever exists so the Go
# loader doesn't panic on a hardcoded name that isn't on disk.
if [ -e "$CURDIR/libgowhisper-fallback.dylib" ]; then
LIBRARY="$CURDIR/libgowhisper-fallback.dylib"
else
LIBRARY="$CURDIR/libgowhisper-fallback.so"
fi
# macOS: single dylib variant (Metal or Accelerate)
LIBRARY="$CURDIR/libgowhisper-fallback.dylib"
export DYLD_LIBRARY_PATH="$CURDIR"/lib:$DYLD_LIBRARY_PATH
else
LIBRARY="$CURDIR/libgowhisper-fallback.so"

View File

@@ -1356,6 +1356,7 @@
intel: "intel-fish-speech"
amd: "rocm-fish-speech"
nvidia-l4t: "nvidia-l4t-fish-speech"
metal: "metal-fish-speech"
default: "cpu-fish-speech"
nvidia-cuda-13: "cuda13-fish-speech"
nvidia-cuda-12: "cuda12-fish-speech"
@@ -4869,6 +4870,7 @@
intel: "intel-fish-speech-development"
amd: "rocm-fish-speech-development"
nvidia-l4t: "nvidia-l4t-fish-speech-development"
metal: "metal-fish-speech-development"
default: "cpu-fish-speech-development"
nvidia-cuda-13: "cuda13-fish-speech-development"
nvidia-cuda-12: "cuda12-fish-speech-development"
@@ -4944,6 +4946,16 @@
uri: "quay.io/go-skynet/local-ai-backends:master-nvidia-l4t-cuda-13-arm64-fish-speech"
mirrors:
- localai/localai-backends:master-nvidia-l4t-cuda-13-arm64-fish-speech
- !!merge <<: *fish-speech
name: "metal-fish-speech"
uri: "quay.io/go-skynet/local-ai-backends:latest-metal-darwin-arm64-fish-speech"
mirrors:
- localai/localai-backends:latest-metal-darwin-arm64-fish-speech
- !!merge <<: *fish-speech
name: "metal-fish-speech-development"
uri: "quay.io/go-skynet/local-ai-backends:master-metal-darwin-arm64-fish-speech"
mirrors:
- localai/localai-backends:master-metal-darwin-arm64-fish-speech
## faster-qwen3-tts
- !!merge <<: *faster-qwen3-tts
name: "faster-qwen3-tts-development"

View File

@@ -1,5 +1,5 @@
--extra-index-url https://download.pytorch.org/whl/rocm7.0
torch==2.10.0+rocm7.0
torch==2.12.0+cpu
torchaudio
torchvision

View File

@@ -0,0 +1,2 @@
torch
torchaudio

View File

@@ -1,5 +1,5 @@
--extra-index-url https://download.pytorch.org/whl/cpu
torch==2.10.0
torch==2.12.0+cpu
transformers>=4.56.2
huggingface-hub>=1.3.0
sentencepiece

View File

@@ -1,4 +1,4 @@
torch==2.10.0
torch==2.12.0+cpu
transformers>=4.56.2
huggingface-hub>=1.3.0
sentencepiece

View File

@@ -1,6 +1,6 @@
--extra-index-url https://download.pytorch.org/whl/cpu
accelerate
torch==2.9.0
torch==2.12.0+cpu
torchvision
torchaudio
transformers

View File

@@ -6,7 +6,7 @@
# for cublas12 so uv consults this index alongside PyPI.
--extra-index-url https://download.pytorch.org/whl/cu128
accelerate
torch==2.9.1
torch==2.12.0+cpu
torchvision
torchaudio
transformers

View File

@@ -1,4 +1,4 @@
accelerate
torch==2.7.0
torch==2.12.0+cu130
transformers
bitsandbytes

View File

@@ -37,8 +37,6 @@ func (a *Application) RestartAgentJobService() error {
if d.JobStore != nil {
agentJobService.SetDistributedJobStore(d.JobStore)
}
// Keep agent tasks consistent across replicas (same client the dispatcher uses).
agentJobService.SetTaskSyncNATS(d.Nats)
}
// Start the service

View File

@@ -604,10 +604,6 @@ func (a *Application) StartAgentPool() {
usm.SetJobDBStore(s)
}
}
// Keep per-user agent tasks consistent across replicas (nil in standalone).
if d := a.Distributed(); d != nil {
usm.SetJobSyncNATS(d.Nats)
}
aps.SetUserServicesManager(usm)
a.agentPoolService.Store(aps)

View File

@@ -280,9 +280,6 @@ func New(opts ...config.AppOption) (*Application, error) {
if application.agentJobService != nil {
application.agentJobService.SetDistributedBackends(distSvc.Dispatcher)
application.agentJobService.SetDistributedJobStore(distSvc.JobStore)
// Keep agent tasks consistent across replicas (jobs already sync via the
// dispatcher + DB read-through). Same NATS client the dispatcher uses.
application.agentJobService.SetTaskSyncNATS(distSvc.Nats)
}
// Wire skill store into AgentPoolService (wired at pool start time via closure)
// The actual wiring happens in StartAgentPool since the pool doesn't exist yet.

View File

@@ -12,12 +12,14 @@ package config
// these; config never imports backend.
const (
// DefaultContextSize is the fallback context window when none is configured
// or estimable from the model. It is also the fallback for a GGUF whose
// metadata yields no usable estimate or that the parser cannot read at all
// (e.g. a quant type it does not know, such as NVFP4): a model-agnostic
// safe default beats a tiny, surprising window that truncates real prompts.
// or estimable from the model.
DefaultContextSize = 4096
// GGUFFallbackContextSize is the context window for a GGUF model whose
// metadata yields no usable estimate (see guessGGUFFromFile). Deliberately
// smaller than DefaultContextSize to stay conservative on memory there.
GGUFFallbackContextSize = 1024
// DefaultNGPULayers means "offload all layers"; the backend (fit_params)
// clamps to what actually fits in device memory.
DefaultNGPULayers = 99999999

View File

@@ -33,7 +33,7 @@ func guessGGUFFromFile(cfg *ModelConfig, f *gguf.GGUFFile, defaultCtx int) {
cSize := int(ctxSize)
cfg.ContextSize = &cSize
} else {
defaultCtx = DefaultContextSize
defaultCtx = GGUFFallbackContextSize
cfg.ContextSize = &defaultCtx
}
}

View File

@@ -34,7 +34,7 @@ func llamaCppDefaults(cfg *ModelConfig, modelPath string) {
// Default context size if not set, regardless of whether GGUF parsing succeeds
defer func() {
if cfg.ContextSize == nil {
ctx := DefaultContextSize
ctx := GGUFFallbackContextSize
cfg.ContextSize = &ctx
}
}()

View File

@@ -248,11 +248,7 @@ var _ = Describe("Backend hooks and parser defaults", func() {
}
cfg.SetDefaults(ModelPath(dir))
// An unreadable/unparseable GGUF (e.g. a quant type the parser does
// not know, such as NVFP4) yields no estimate, so the hook must fall
// back to DefaultContextSize rather than a tiny, surprising value.
Expect(cfg.ContextSize).NotTo(BeNil())
Expect(*cfg.ContextSize).To(Equal(DefaultContextSize))
})
})

View File

@@ -23,10 +23,8 @@ import (
"github.com/mudler/LocalAI/core/application"
"github.com/mudler/LocalAI/core/schema"
"github.com/mudler/LocalAI/core/services/distributed"
"github.com/mudler/LocalAI/core/services/finetune"
"github.com/mudler/LocalAI/core/services/galleryop"
"github.com/mudler/LocalAI/core/services/messaging"
"github.com/mudler/LocalAI/core/services/nodes"
"github.com/mudler/LocalAI/core/services/quantization"
@@ -402,45 +400,25 @@ func API(application *application.Application) (*echo.Echo, error) {
routes.RegisterAgentPoolRoutes(e, application, agentsMw, skillsMw, collectionsMw)
// Fine-tuning routes
fineTuningMw := auth.RequireFeature(application.AuthDB(), auth.FeatureFineTuning)
// In distributed mode pass the shared NATS client + PostgreSQL store so
// fine-tune jobs stay consistent across replicas (the SyncedMap broadcasts
// mutations and hydrates from the DB); standalone passes nil for both.
var ftNats messaging.MessagingClient
var ftStore *distributed.FineTuneStore
if d := application.Distributed(); d != nil {
ftNats = d.Nats
if d.DistStores != nil && d.DistStores.FineTune != nil {
ftStore = d.DistStores.FineTune
}
}
ftService := finetune.NewFineTuneService(
application.ApplicationConfig(),
application.ModelLoader(),
application.ModelConfigLoader(),
ftNats,
ftStore,
)
if d := application.Distributed(); d != nil {
ftService.SetNATSClient(d.Nats)
if d.DistStores != nil && d.DistStores.FineTune != nil {
ftService.SetFineTuneStore(d.DistStores.FineTune)
}
}
routes.RegisterFineTuningRoutes(e, ftService, application.ApplicationConfig(), fineTuningMw)
// Quantization routes
quantizationMw := auth.RequireFeature(application.AuthDB(), auth.FeatureQuantization)
// In distributed mode pass the shared NATS client + PostgreSQL store so
// quantization jobs stay consistent across replicas (the SyncedMap broadcasts
// mutations and hydrates from the DB); standalone passes nil for both.
var quantNats messaging.MessagingClient
var quantStore *distributed.QuantStore
if d := application.Distributed(); d != nil {
quantNats = d.Nats
if d.DistStores != nil && d.DistStores.Quant != nil {
quantStore = d.DistStores.Quant
}
}
qService := quantization.NewQuantizationService(
application.ApplicationConfig(),
application.ModelLoader(),
application.ModelConfigLoader(),
quantNats,
quantStore,
)
routes.RegisterQuantizationRoutes(e, qService, application.ApplicationConfig(), quantizationMw)

View File

@@ -25,7 +25,6 @@ import (
"github.com/mudler/LocalAI/core/http/auth"
"github.com/mudler/LocalAI/core/schema"
"github.com/mudler/LocalAI/core/services/galleryop"
"github.com/mudler/LocalAI/core/services/messaging"
"github.com/mudler/LocalAI/core/services/nodes"
"github.com/mudler/LocalAI/core/services/nodes/prefixcache"
"github.com/mudler/LocalAI/pkg/httpclient"
@@ -551,23 +550,12 @@ func DeleteBackendOnNodeEndpoint(unloader nodes.NodeCommandSender) echo.HandlerF
}
// ListBackendsOnNodeEndpoint lists installed backends on a worker node via NATS.
func ListBackendsOnNodeEndpoint(unloader nodes.NodeCommandSender, registry *nodes.NodeRegistry) echo.HandlerFunc {
func ListBackendsOnNodeEndpoint(unloader nodes.NodeCommandSender) echo.HandlerFunc {
return func(c echo.Context) error {
nodeID := c.Param("id")
// Agent-type workers don't run backends and never subscribe to the
// nodes.<id>.backend.list NATS subject, so the request would hang
// until timeout with "no responders". Their backend list is simply
// empty. Mirror the aggregate-list guard in managers_distributed.go
// (skip nodes whose NodeType is set and not "backend") so the
// single-node and cluster-wide views stay consistent.
if node, err := registry.Get(c.Request().Context(), nodeID); err == nil {
if node.NodeType != "" && node.NodeType != nodes.NodeTypeBackend {
return c.JSON(http.StatusOK, []messaging.NodeBackendInfo{})
}
}
if unloader == nil {
return c.JSON(http.StatusServiceUnavailable, nodeError(http.StatusServiceUnavailable, "NATS not configured"))
}
nodeID := c.Param("id")
reply, err := unloader.ListBackends(nodeID)
if err != nil {
xlog.Error("Failed to list backends on node", "node", nodeID, "error", err)

View File

@@ -1,103 +0,0 @@
package localai
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"github.com/labstack/echo/v4"
"github.com/mudler/LocalAI/core/services/messaging"
"github.com/mudler/LocalAI/core/services/nodes"
"github.com/mudler/LocalAI/core/services/testutil"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
)
// stubNodeCommandSender records whether ListBackends was invoked so the test can
// assert the endpoint short-circuits (no NATS request) for agent-type nodes.
type stubNodeCommandSender struct {
listBackendsCalled bool
}
func (s *stubNodeCommandSender) InstallBackend(_, _, _, _, _, _, _ string, _ int, _ string, _ func(messaging.BackendInstallProgressEvent)) (*messaging.BackendInstallReply, error) {
return &messaging.BackendInstallReply{}, nil
}
func (s *stubNodeCommandSender) UpgradeBackend(_, _, _, _, _, _ string, _ int, _ string, _ func(messaging.BackendInstallProgressEvent)) (*messaging.BackendUpgradeReply, error) {
return &messaging.BackendUpgradeReply{}, nil
}
func (s *stubNodeCommandSender) DeleteBackend(_, _ string) (*messaging.BackendDeleteReply, error) {
return &messaging.BackendDeleteReply{Success: true}, nil
}
func (s *stubNodeCommandSender) ListBackends(_ string) (*messaging.BackendListReply, error) {
s.listBackendsCalled = true
return &messaging.BackendListReply{Backends: []messaging.NodeBackendInfo{{Name: "llama-cpp"}}}, nil
}
func (s *stubNodeCommandSender) StopBackend(_, _ string) error { return nil }
func (s *stubNodeCommandSender) UnloadModelOnNode(_, _ string) error { return nil }
var _ = Describe("ListBackendsOnNodeEndpoint", func() {
var registry *nodes.NodeRegistry
BeforeEach(func() {
db := testutil.SetupTestDB()
var err error
registry, err = nodes.NewNodeRegistry(db)
Expect(err).ToNot(HaveOccurred())
})
callEndpoint := func(unloader nodes.NodeCommandSender, nodeID string) *httptest.ResponseRecorder {
e := echo.New()
req := httptest.NewRequest(http.MethodGet, "/", nil)
rec := httptest.NewRecorder()
c := e.NewContext(req, rec)
c.SetParamNames("id")
c.SetParamValues(nodeID)
handler := ListBackendsOnNodeEndpoint(unloader, registry)
Expect(handler(c)).To(Succeed())
return rec
}
It("returns an empty list for an agent node without issuing a NATS request", func() {
ctx := context.Background()
node := &nodes.BackendNode{Name: "agent-1", NodeType: nodes.NodeTypeAgent}
Expect(registry.Register(ctx, node, true)).To(Succeed())
stub := &stubNodeCommandSender{}
rec := callEndpoint(stub, node.ID)
Expect(rec.Code).To(Equal(http.StatusOK))
Expect(stub.listBackendsCalled).To(BeFalse(),
"agent workers don't subscribe to backend.list; the endpoint must not issue the doomed NATS request")
var list []messaging.NodeBackendInfo
Expect(json.Unmarshal(rec.Body.Bytes(), &list)).To(Succeed())
Expect(list).To(BeEmpty())
// Must be `[]`, not `null`, so the UI can render it.
Expect(rec.Body.String()).To(ContainSubstring("[]"))
})
It("consults the unloader (NATS) for a backend node", func() {
ctx := context.Background()
node := &nodes.BackendNode{Name: "backend-1", NodeType: nodes.NodeTypeBackend, Address: "10.0.0.1:50051"}
Expect(registry.Register(ctx, node, true)).To(Succeed())
stub := &stubNodeCommandSender{}
rec := callEndpoint(stub, node.ID)
Expect(rec.Code).To(Equal(http.StatusOK))
Expect(stub.listBackendsCalled).To(BeTrue(),
"backend nodes must still be queried over NATS")
var list []messaging.NodeBackendInfo
Expect(json.Unmarshal(rec.Body.Bytes(), &list)).To(Succeed())
Expect(list).To(HaveLen(1))
Expect(list[0].Name).To(Equal("llama-cpp"))
})
})

View File

@@ -88,7 +88,7 @@ func RegisterNodeAdminRoutes(e *echo.Echo, registry *nodes.NodeRegistry, unloade
admin.POST("/:id/approve", localai.ApproveNodeEndpoint(registry, authDB, hmacSecret, natsCfg))
// Backend management on workers
admin.GET("/:id/backends", localai.ListBackendsOnNodeEndpoint(unloader, registry))
admin.GET("/:id/backends", localai.ListBackendsOnNodeEndpoint(unloader))
admin.POST("/:id/backends/install", localai.InstallBackendOnNodeEndpoint(unloader, galleryService, opcache, appConfig))
admin.POST("/:id/backends/delete", localai.DeleteBackendOnNodeEndpoint(unloader))

View File

@@ -30,8 +30,6 @@ import (
mcpTools "github.com/mudler/LocalAI/core/http/endpoints/mcp"
"github.com/mudler/LocalAI/core/schema"
"github.com/mudler/LocalAI/core/services/jobs"
"github.com/mudler/LocalAI/core/services/messaging"
"github.com/mudler/LocalAI/core/services/syncstate"
"github.com/mudler/LocalAI/core/templates"
"github.com/mudler/LocalAI/pkg/httpclient"
"github.com/mudler/LocalAI/pkg/model"
@@ -45,18 +43,8 @@ type AgentJobService struct {
configLoader *config.ModelConfigLoader
evaluator *templates.Evaluator
// tasks is the cross-replica task store: an in-memory map kept consistent
// across replicas via NATS, with read-through to the configured persister
// (file in standalone, PostgreSQL in distributed). Unlike jobs - which already
// converge via the dispatcher + DB read-through - tasks previously read
// in-memory only, so ListTasks went stale on non-originating replicas.
tasks *syncstate.SyncedMap[string, schema.Task]
// taskNats is the distributed NATS client backing the tasks SyncedMap. It is
// not available at construction time, so it is injected via SetTaskSyncNATS
// during distributed wiring; nil keeps tasks in-memory-only (standalone).
taskNats messaging.MessagingClient
// Storage (in-memory primary, persister for secondary persistence)
tasks *xsync.SyncedMap[string, schema.Task]
jobs *xsync.SyncedMap[string, schema.Job]
persister JobPersister
userID string // Scoping: empty for global (main service), set for per-user instances
@@ -108,31 +96,6 @@ func (s *AgentJobService) SetDistributedJobStore(store *jobs.JobStore) {
s.persister = &dbJobPersister{store: store}
}
// SetTaskSyncNATS wires the distributed NATS client used to keep agent *tasks*
// consistent across replicas (jobs already converge via the dispatcher + DB
// read-through, so they are left untouched). The client is not available when the
// service is constructed, so it is injected here during distributed wiring and the
// tasks SyncedMap is rebuilt to pick it up. It is always called before Start /
// hydrate, while the map is still empty, so rebuilding loses no state. Passing nil
// (standalone) keeps the map in-memory-only with no broadcast.
func (s *AgentJobService) SetTaskSyncNATS(nats messaging.MessagingClient) {
s.taskNats = nats
s.buildTasksMap()
}
// buildTasksMap (re)constructs the cross-replica tasks SyncedMap from the current
// taskNats. The Store adapter reads s.persister/s.userID live, so a persister swap
// (SetDistributedJobStore) needs no rebuild; only the NATS client, fixed at
// New-time, forces one - hence SetTaskSyncNATS calls this.
func (s *AgentJobService) buildTasksMap() {
s.tasks = syncstate.New(syncstate.Config[string, schema.Task]{
Name: "agent.tasks",
Key: func(t schema.Task) string { return t.ID },
Nats: s.taskNats,
Store: &taskStoreAdapter{svc: s},
})
}
// Dispatcher returns the distributed dispatcher (nil if not in distributed mode).
func (s *AgentJobService) Dispatcher() DistributedDispatcher {
return s.dispatcher
@@ -143,6 +106,13 @@ func (s *AgentJobService) DBStore() *jobs.JobStore {
return s.rawDBStore
}
// saveTasks persists tasks via the configured persister (file or DB).
func (s *AgentJobService) saveTasks(task schema.Task) {
if err := s.persister.SaveTask(s.userID, task); err != nil {
xlog.Warn("Failed to persist task", "error", err, "task_id", task.ID)
}
}
// saveJobs persists jobs via the configured persister (file or DB).
func (s *AgentJobService) saveJobs(job schema.Job) {
if err := s.persister.SaveJob(s.userID, job); err != nil {
@@ -159,8 +129,18 @@ func (s *AgentJobService) LoadFromDB() {
// loadFromPersister loads tasks and jobs from the configured persister into memory.
func (s *AgentJobService) loadFromPersister() {
if err := s.hydrateTasks(s.appConfig.Context); err != nil {
if tasks, err := s.persister.LoadTasks(s.userID); err != nil {
xlog.Warn("Failed to load tasks from persister", "error", err)
} else {
for _, task := range tasks {
s.tasks.Set(task.ID, task)
if task.Enabled && task.Cron != "" {
if err := s.ScheduleCronTask(task); err != nil {
xlog.Warn("Failed to schedule cron task on load", "error", err, "task_id", task.ID)
}
}
}
xlog.Info("Loaded tasks from persister", "count", len(tasks))
}
if loadedJobs, err := s.persister.LoadJobs(s.userID); err != nil {
@@ -173,27 +153,6 @@ func (s *AgentJobService) loadFromPersister() {
}
}
// hydrateTasks loads tasks into the cross-replica SyncedMap and (re)schedules
// cron entries for enabled tasks. Hydration goes through the SyncedMap's Store
// read-through (Start), not Set, so it neither re-persists nor re-broadcasts the
// loaded tasks. Each service instance hydrates exactly once: the main service via
// Start -> loadFromPersister, per-user services via LoadFromDB or LoadTasksFromFile.
func (s *AgentJobService) hydrateTasks(ctx context.Context) error {
if err := s.tasks.Start(ctx); err != nil {
return err
}
tasks := s.tasks.List()
for _, task := range tasks {
if task.Enabled && task.Cron != "" {
if err := s.ScheduleCronTask(task); err != nil {
xlog.Warn("Failed to schedule cron task on load", "error", err, "task_id", task.ID)
}
}
}
xlog.Info("Loaded tasks from persister", "count", len(tasks))
return nil
}
// JobExecution represents a job to be executed
type JobExecution struct {
Job schema.Job
@@ -241,19 +200,21 @@ func NewAgentJobServiceWithPaths(
) *AgentJobService {
retentionDays := cmp.Or(appConfig.AgentJobRetentionDays, 30)
tasks := xsync.NewSyncedMap[string, schema.Task]()
jobsMap := xsync.NewSyncedMap[string, schema.Job]()
s := &AgentJobService{
return &AgentJobService{
appConfig: appConfig,
modelLoader: modelLoader,
configLoader: configLoader,
evaluator: evaluator,
tasks: tasks,
jobs: jobsMap,
persister: &fileJobPersister{
tasks: tasks,
jobs: jobsMap,
tasksFile: tasksFile,
jobsFile: jobsFile,
taskSet: make(map[string]schema.Task),
},
jobQueue: make(chan JobExecution, 100), // Buffer for 100 jobs
cancellations: xsync.NewSyncedMap[string, context.CancelFunc](),
@@ -261,17 +222,25 @@ func NewAgentJobServiceWithPaths(
cronEntries: xsync.NewSyncedMap[string, cron.EntryID](),
retentionDays: retentionDays,
}
// Build the cross-replica tasks map standalone (nil NATS); SetTaskSyncNATS
// rebuilds it with the distributed client once that is available, before Start.
s.buildTasksMap()
return s
}
// LoadTasksFromFile loads tasks from the persister into the in-memory map
// and schedules cron entries. Named "FromFile" for backward compat; in DB
// mode it loads from the database.
func (s *AgentJobService) LoadTasksFromFile() error {
return s.hydrateTasks(s.appConfig.Context)
tasks, err := s.persister.LoadTasks(s.userID)
if err != nil {
return err
}
for _, task := range tasks {
s.tasks.Set(task.ID, task)
if task.Enabled && task.Cron != "" {
if err := s.ScheduleCronTask(task); err != nil {
xlog.Warn("Failed to schedule cron task on load", "error", err, "task_id", task.ID)
}
}
}
return nil
}
// SaveTasksToFile flushes the current tasks map via the persister. File
@@ -324,12 +293,8 @@ func (s *AgentJobService) CreateTask(task schema.Task) (string, error) {
task.Enabled = true // Default to enabled
}
// Store task: Set updates the in-memory map, write-throughs to the persister
// (file or DB), and broadcasts the create to peer replicas. Background ctx
// because CreateTask carries no request ctx (mirrors the finetune service).
if err := s.tasks.Set(context.Background(), task); err != nil {
return "", fmt.Errorf("failed to persist task: %w", err)
}
// Store task
s.tasks.Set(id, task)
// Schedule cron if enabled and has cron expression
if task.Enabled && task.Cron != "" {
@@ -338,15 +303,16 @@ func (s *AgentJobService) CreateTask(task schema.Task) (string, error) {
}
}
s.saveTasks(task)
return id, nil
}
// UpdateTask updates an existing task
func (s *AgentJobService) UpdateTask(id string, task schema.Task) error {
existing, ok := s.tasks.Get(id)
if !ok {
if !s.tasks.Exists(id) {
return fmt.Errorf("%w: %s", ErrTaskNotFound, id)
}
existing := s.tasks.Get(id)
// Preserve ID and CreatedAt
task.ID = id
@@ -358,10 +324,8 @@ func (s *AgentJobService) UpdateTask(id string, task schema.Task) error {
s.UnscheduleCronTask(id)
}
// Store updated task: write-through + broadcast (see CreateTask).
if err := s.tasks.Set(context.Background(), task); err != nil {
return fmt.Errorf("failed to persist task: %w", err)
}
// Store updated task
s.tasks.Set(id, task)
// Schedule new cron if enabled and has cron expression
if task.Enabled && task.Cron != "" {
@@ -370,22 +334,24 @@ func (s *AgentJobService) UpdateTask(id string, task schema.Task) error {
}
}
s.saveTasks(task)
return nil
}
// DeleteTask deletes a task
func (s *AgentJobService) DeleteTask(id string) error {
if _, ok := s.tasks.Get(id); !ok {
if !s.tasks.Exists(id) {
return fmt.Errorf("%w: %s", ErrTaskNotFound, id)
}
// Unschedule cron
s.UnscheduleCronTask(id)
// Delete removes from the in-memory map, deletes from the persister, and
// broadcasts the removal to peer replicas.
if err := s.tasks.Delete(context.Background(), id); err != nil {
xlog.Warn("Failed to delete task from store", "error", err, "task_id", id)
// Remove from memory
s.tasks.Delete(id)
if err := s.persister.DeleteTask(id); err != nil {
xlog.Warn("Failed to delete task from persister", "error", err, "task_id", id)
}
return nil
@@ -393,8 +359,8 @@ func (s *AgentJobService) DeleteTask(id string) error {
// GetTask retrieves a task by ID
func (s *AgentJobService) GetTask(id string) (*schema.Task, error) {
task, ok := s.tasks.Get(id)
if !ok {
task := s.tasks.Get(id)
if task.ID == "" {
return nil, fmt.Errorf("%w: %s", ErrTaskNotFound, id)
}
return &task, nil
@@ -402,7 +368,7 @@ func (s *AgentJobService) GetTask(id string) (*schema.Task, error) {
// ListTasks returns all tasks, sorted by creation date (newest first)
func (s *AgentJobService) ListTasks() []schema.Task {
tasks := s.tasks.List()
tasks := s.tasks.Values()
// Sort by CreatedAt descending (newest first), then by Name for stability
slices.SortFunc(tasks, func(a, b schema.Task) int {
if a.CreatedAt.Equal(b.CreatedAt) {
@@ -431,8 +397,8 @@ func (s *AgentJobService) buildPrompt(templateStr string, params map[string]stri
// ExecuteJob creates and queues a job for execution
// multimedia can be nil for backward compatibility
func (s *AgentJobService) ExecuteJob(taskID string, params map[string]string, triggeredBy string, multimedia *schema.MultimediaAttachment) (string, error) {
task, ok := s.tasks.Get(taskID)
if !ok {
task := s.tasks.Get(taskID)
if task.ID == "" {
return "", fmt.Errorf("%w: %s", ErrTaskNotFound, taskID)
}
@@ -1485,12 +1451,6 @@ func (s *AgentJobService) Stop() error {
if s.cronScheduler != nil {
s.cronScheduler.Stop()
}
// Release the tasks SyncedMap subscription / background workers.
if s.tasks != nil {
if err := s.tasks.Close(); err != nil {
xlog.Warn("Error closing tasks sync map", "error", err)
}
}
xlog.Info("AgentJobService stopped")
return nil
}

View File

@@ -14,38 +14,24 @@ import (
)
// fileJobPersister persists tasks and jobs to JSON files.
//
// Jobs serialize the service's in-memory jobs syncmap on each save (bulk write).
// Tasks are kept in this persister's own taskSet map instead: the tasks SyncedMap
// calls SaveTask/DeleteTask while holding its internal lock (write-through), so
// reading back the SyncedMap here would re-enter that lock and deadlock. The
// self-contained taskSet, seeded by LoadTasks, lets a per-task write rewrite the
// whole bulk file without touching the SyncedMap.
//
// Runtime reads (GetJob/ListJobs) return nil (the in-memory state is the
// authoritative source); LoadTasks/LoadJobs bootstrap state at startup.
// It holds references to the service's syncmaps and serializes the entire
// map contents on each save (bulk write). Reads at runtime return nil
// (the in-memory map is the authoritative source); LoadTasks/LoadJobs
// are used only at startup to bootstrap the syncmaps.
type fileJobPersister struct {
tasks *xsync.SyncedMap[string, schema.Task]
jobs *xsync.SyncedMap[string, schema.Job]
tasksFile string
jobsFile string
mu sync.Mutex
// taskSet is the persister's own view of all tasks, seeded by LoadTasks and
// updated by SaveTask/DeleteTask. The bulk JSON file is rewritten from it.
taskSet map[string]schema.Task
}
func (p *fileJobPersister) SaveTask(_ string, task schema.Task) error {
p.mu.Lock()
defer p.mu.Unlock()
p.taskSet[task.ID] = task
return p.writeTasksLocked()
func (p *fileJobPersister) SaveTask(_ string, _ schema.Task) error {
return p.saveTasksToFile()
}
func (p *fileJobPersister) DeleteTask(taskID string) error {
p.mu.Lock()
defer p.mu.Unlock()
delete(p.taskSet, taskID)
return p.writeTasksLocked()
func (p *fileJobPersister) DeleteTask(_ string) error {
return p.saveTasksToFile()
}
func (p *fileJobPersister) SaveJob(_ string, _ schema.Job) error {
@@ -57,9 +43,7 @@ func (p *fileJobPersister) DeleteJob(_ string) error {
}
func (p *fileJobPersister) FlushTasks() error {
p.mu.Lock()
defer p.mu.Unlock()
return p.writeTasksLocked()
return p.saveTasksToFile()
}
func (p *fileJobPersister) FlushJobs() error {
@@ -99,12 +83,6 @@ func (p *fileJobPersister) LoadTasks(_ string) ([]schema.Task, error) {
return nil, fmt.Errorf("failed to parse tasks file: %w", err)
}
// Seed the in-memory set so subsequent per-task SaveTask/DeleteTask merge into
// (rather than overwrite) the persisted tasks when the bulk file is rewritten.
for _, t := range tf.Tasks {
p.taskSet[t.ID] = t
}
xlog.Info("Loaded tasks from file", "count", len(tf.Tasks))
return tf.Tasks, nil
}
@@ -140,19 +118,18 @@ func (p *fileJobPersister) CleanupOldJobs(_ time.Duration) (int64, error) {
return 0, nil // cleanup handled via in-memory filtering
}
// writeTasksLocked serializes the persister's task set to the JSON file. Callers
// must hold p.mu.
func (p *fileJobPersister) writeTasksLocked() error {
// saveTasksToFile serializes the entire tasks map to the JSON file.
func (p *fileJobPersister) saveTasksToFile() error {
if p.tasksFile == "" {
return nil
}
tasks := make([]schema.Task, 0, len(p.taskSet))
for _, t := range p.taskSet {
tasks = append(tasks, t)
}
p.mu.Lock()
defer p.mu.Unlock()
tf := schema.TasksFile{Tasks: tasks}
tf := schema.TasksFile{
Tasks: p.tasks.Values(),
}
data, err := json.MarshalIndent(tf, "", " ")
if err != nil {

View File

@@ -20,26 +20,28 @@ var _ = Describe("JobPersister", func() {
Context("fileJobPersister", func() {
var (
p *fileJobPersister
tasks *xsync.SyncedMap[string, schema.Task]
jobsMap *xsync.SyncedMap[string, schema.Job]
tmpDir string
)
BeforeEach(func() {
tmpDir = GinkgoT().TempDir()
tasks = xsync.NewSyncedMap[string, schema.Task]()
jobsMap = xsync.NewSyncedMap[string, schema.Job]()
p = &fileJobPersister{
tasks: tasks,
jobs: jobsMap,
tasksFile: filepath.Join(tmpDir, "tasks.json"),
jobsFile: filepath.Join(tmpDir, "jobs.json"),
// taskSet is the persister's own task view (decoupled from the tasks
// SyncedMap to avoid re-entering its lock during write-through).
taskSet: make(map[string]schema.Task),
}
})
It("SaveTask writes all tasks to file", func() {
Expect(p.SaveTask("", schema.Task{ID: "t1", Name: "Task One", Model: "m", Prompt: "p"})).To(Succeed())
Expect(p.SaveTask("", schema.Task{ID: "t2", Name: "Task Two", Model: "m", Prompt: "p"})).To(Succeed())
tasks.Set("t1", schema.Task{ID: "t1", Name: "Task One", Model: "m", Prompt: "p"})
tasks.Set("t2", schema.Task{ID: "t2", Name: "Task Two", Model: "m", Prompt: "p"})
Expect(p.SaveTask("", schema.Task{})).To(Succeed())
// Verify file contents
data, err := os.ReadFile(p.tasksFile)
@@ -50,9 +52,11 @@ var _ = Describe("JobPersister", func() {
})
It("DeleteTask writes updated tasks to file", func() {
Expect(p.SaveTask("", schema.Task{ID: "t1", Name: "Keep"})).To(Succeed())
Expect(p.SaveTask("", schema.Task{ID: "t2", Name: "Delete"})).To(Succeed())
tasks.Set("t1", schema.Task{ID: "t1", Name: "Keep"})
tasks.Set("t2", schema.Task{ID: "t2", Name: "Delete"})
// Simulate deletion from memory (caller does this before calling persister)
tasks.Delete("t2")
Expect(p.DeleteTask("t2")).To(Succeed())
data, err := os.ReadFile(p.tasksFile)

View File

@@ -1,152 +0,0 @@
package agentpool
// White-box tests (package agentpool) so a spec can build two AgentJobService
// instances sharing one in-memory bus and assert that agent *tasks* converge
// across replicas - the bug this migration fixes (ListTasks used to read
// in-memory only, so a task created on replica A was invisible on replica B).
// Jobs are deliberately untouched here: they already converge via the dispatcher
// + DB read-through.
import (
"context"
"time"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"github.com/mudler/LocalAI/core/config"
"github.com/mudler/LocalAI/core/schema"
"github.com/mudler/LocalAI/core/services/messaging"
"github.com/mudler/LocalAI/core/services/syncstate"
"github.com/mudler/LocalAI/core/services/testutil"
"github.com/mudler/LocalAI/pkg/system"
)
// newTaskSyncService builds an AgentJobService wired to the given bus and a
// throwaway data dir (so the file persister has somewhere to write). Model/config
// loaders are nil because the task sync paths under test never touch them.
func newTaskSyncService(bus messaging.MessagingClient) *AgentJobService {
tmpDir := GinkgoT().TempDir()
sysState := &system.SystemState{}
sysState.Model.ModelsPath = tmpDir
appConfig := config.NewApplicationConfig(
config.WithDynamicConfigDir(tmpDir),
config.WithContext(context.Background()),
)
appConfig.SystemState = sysState
svc := NewAgentJobServiceWithPaths(appConfig, nil, nil, nil,
// Distinct per-replica files so the file persister write-through never
// crosses replicas: convergence here must be proven via the bus alone.
tmpDir+"/tasks.json", tmpDir+"/jobs.json")
svc.SetTaskSyncNATS(bus)
return svc
}
var _ = Describe("AgentJobService task cross-replica sync", func() {
Describe("two replicas sharing one bus", func() {
var (
bus *testutil.FakeBus
a, b *AgentJobService
)
BeforeEach(func() {
// One shared bus, two replicas: exactly the distributed topology where a
// round-robin request may land on a replica that did not originate the
// change.
bus = testutil.NewFakeBus()
a = newTaskSyncService(bus)
b = newTaskSyncService(bus)
// Start hydrates (empty here) and subscribes both replicas to deltas.
Expect(a.Start(context.Background())).To(Succeed())
Expect(b.Start(context.Background())).To(Succeed())
})
AfterEach(func() {
Expect(a.Stop()).To(Succeed())
Expect(b.Stop()).To(Succeed())
})
It("makes a task created on A visible via B's GetTask and ListTasks", func() {
id, err := a.CreateTask(schema.Task{Name: "Shared", Model: "m", Prompt: "p"})
Expect(err).NotTo(HaveOccurred())
got, err := b.GetTask(id)
Expect(err).NotTo(HaveOccurred(), "B must see a task A just created")
Expect(got.Name).To(Equal("Shared"))
listed := b.ListTasks()
Expect(listed).To(HaveLen(1))
Expect(listed[0].ID).To(Equal(id))
})
It("propagates a task update from A to B", func() {
id, err := a.CreateTask(schema.Task{Name: "Before", Model: "m", Prompt: "p"})
Expect(err).NotTo(HaveOccurred())
Expect(a.UpdateTask(id, schema.Task{Name: "After", Model: "m", Prompt: "p"})).To(Succeed())
got, err := b.GetTask(id)
Expect(err).NotTo(HaveOccurred())
Expect(got.Name).To(Equal("After"), "an update on A must be visible on B")
})
It("removes a task from B when it is deleted on A", func() {
id, err := a.CreateTask(schema.Task{Name: "Doomed", Model: "m", Prompt: "p"})
Expect(err).NotTo(HaveOccurred())
_, err = b.GetTask(id)
Expect(err).NotTo(HaveOccurred(), "precondition: B must have the task before the delete")
Expect(a.DeleteTask(id)).To(Succeed())
_, err = b.GetTask(id)
Expect(err).To(HaveOccurred(), "a delete on A must remove the task from B")
Expect(b.ListTasks()).To(BeEmpty())
})
It("does not re-broadcast a delta it received (echo-loop guard)", func() {
subject := messaging.SubjectSyncStateDelta("agent.tasks")
_, err := a.CreateTask(schema.Task{Name: "Once", Model: "m", Prompt: "p"})
Expect(err).NotTo(HaveOccurred())
// Exactly one publish: A's create. B applies it without re-publishing,
// otherwise this would be 2+ and a real bus would storm.
Expect(bus.PublishCount(subject)).To(Equal(1))
})
})
Describe("ListTasks ordering and scoping", func() {
var svc *AgentJobService
BeforeEach(func() {
svc = newTaskSyncService(testutil.NewFakeBus())
Expect(svc.Start(context.Background())).To(Succeed())
})
AfterEach(func() { Expect(svc.Stop()).To(Succeed()) })
It("sorts newest-first, breaking ties by name", func() {
// CreateTask stamps CreatedAt with time.Now(); space them out so ordering
// is deterministic rather than relying on the sub-millisecond gap.
oldID, err := svc.CreateTask(schema.Task{Name: "Old", Model: "m", Prompt: "p"})
Expect(err).NotTo(HaveOccurred())
time.Sleep(5 * time.Millisecond)
newID, err := svc.CreateTask(schema.Task{Name: "New", Model: "m", Prompt: "p"})
Expect(err).NotTo(HaveOccurred())
listed := svc.ListTasks()
Expect(listed).To(HaveLen(2))
Expect(listed[0].ID).To(Equal(newID), "newest first")
Expect(listed[1].ID).To(Equal(oldID))
})
})
Describe("compile-time adapter contract", func() {
It("satisfies syncstate.Store for tasks", func() {
// Mirrors the var assertion in task_syncstore.go; keeps the type
// referenced from a spec so drift surfaces here too.
var _ syncstate.Store[string, schema.Task] = (*taskStoreAdapter)(nil)
Expect(&taskStoreAdapter{}).ToNot(BeNil())
})
})
})

View File

@@ -1,47 +0,0 @@
package agentpool
import (
"context"
"github.com/mudler/LocalAI/core/schema"
"github.com/mudler/LocalAI/core/services/syncstate"
)
// taskStoreAdapter bridges the existing JobPersister (file- or DB-backed) to the
// generic syncstate.Store the tasks SyncedMap consumes. Only tasks are migrated:
// jobs already converge across replicas via the dispatcher (NATS) plus the DB
// read-through in ListJobs/GetJob, whereas ListTasks read in-memory only and so
// went stale on replicas that did not originate the change.
//
// The adapter reads svc.persister and svc.userID live (rather than capturing
// them) because both are configured by setters - SetDistributedJobStore swaps the
// file persister for the DB one, SetUserID scopes per-user queries - AFTER the
// service, and thus this adapter, is constructed. Reading them at call time means
// the SyncedMap never has to be rebuilt when the persister is swapped.
//
// The SyncedMap value type is schema.Task: the exact shape ListTasks returns, so
// reads need no conversion and REST responses are provably unchanged.
type taskStoreAdapter struct {
svc *AgentJobService
}
// compile-time assertion that the adapter satisfies the component's Store.
var _ syncstate.Store[string, schema.Task] = (*taskStoreAdapter)(nil)
// List hydrates the map from durable storage on Start/reconnect: the file's task
// list (standalone) or every task row (DB / distributed).
func (a *taskStoreAdapter) List(_ context.Context) ([]schema.Task, error) {
return a.svc.persister.LoadTasks(a.svc.userID)
}
// Upsert write-through persists a single task created/updated locally; the
// SyncedMap then broadcasts the delta to peers.
func (a *taskStoreAdapter) Upsert(_ context.Context, task schema.Task) error {
return a.svc.persister.SaveTask(a.svc.userID, task)
}
// Delete write-through removes a task locally; the SyncedMap then broadcasts the
// removal to peers.
func (a *taskStoreAdapter) Delete(_ context.Context, id string) error {
return a.svc.persister.DeleteTask(id)
}

View File

@@ -7,7 +7,6 @@ import (
"github.com/mudler/LocalAGI/webui/collections"
"github.com/mudler/LocalAI/core/config"
"github.com/mudler/LocalAI/core/services/jobs"
"github.com/mudler/LocalAI/core/services/messaging"
"github.com/mudler/LocalAI/core/templates"
"github.com/mudler/LocalAI/pkg/model"
"github.com/mudler/xlog"
@@ -29,9 +28,6 @@ type UserServicesManager struct {
// Shared distributed backends (set once, inherited by per-user job services)
jobDispatcher DistributedDispatcher
jobDBStore *jobs.JobStore
// jobNats keeps per-user agent tasks consistent across replicas (nil in
// standalone). Inherited by each per-user AgentJobService.
jobNats messaging.MessagingClient
}
// NewUserServicesManager creates a new UserServicesManager.
@@ -166,10 +162,6 @@ func (m *UserServicesManager) GetJobs(userID string) (*AgentJobService, error) {
if m.jobDispatcher != nil {
svc.SetDistributedBackends(m.jobDispatcher)
}
// Inherit the NATS client so per-user tasks broadcast across replicas. Must be
// set before the hydrate below (LoadFromDB / LoadTasksFromFile) so the tasks
// SyncedMap is rebuilt with the client while it is still empty.
svc.SetTaskSyncNATS(m.jobNats)
if m.jobDBStore != nil {
svc.SetDistributedJobStore(m.jobDBStore)
// Load tasks/jobs from DB immediately (per-user services skip Start())
@@ -197,12 +189,6 @@ func (m *UserServicesManager) SetJobDBStore(s *jobs.JobStore) {
m.jobDBStore = s
}
// SetJobSyncNATS sets the NATS client used to keep per-user agent tasks consistent
// across replicas.
func (m *UserServicesManager) SetJobSyncNATS(nats messaging.MessagingClient) {
m.jobNats = nats
}
// ListAllUserIDs returns all user IDs that have scoped data directories.
func (m *UserServicesManager) ListAllUserIDs() ([]string, error) {
return m.storage.ListUserDirs()

View File

@@ -8,7 +8,6 @@ import (
"github.com/google/uuid"
"github.com/mudler/LocalAI/core/services/advisorylock"
"gorm.io/gorm"
"gorm.io/gorm/clause"
)
// FineTuneJobRecord tracks fine-tune jobs in PostgreSQL.
@@ -81,34 +80,6 @@ func (s *FineTuneStore) List(userID string) ([]FineTuneJobRecord, error) {
return jobs, q.Find(&jobs).Error
}
// ListAll returns every fine-tune job across all users. The SyncedMap that backs
// FineTuneService is a single global map (the REST API filters by user at read
// time), so hydrate needs the full set rather than the per-user List above.
func (s *FineTuneStore) ListAll() ([]FineTuneJobRecord, error) {
var jobs []FineTuneJobRecord
return jobs, s.db.Order("created_at DESC").Find(&jobs).Error
}
// Upsert idempotently inserts or fully replaces a job row by primary key. The
// SyncedMap write-through path issues a single Set per mutation regardless of
// whether the job already exists, so it needs one create-or-update primitive
// (Create alone fails on a duplicate key, UpdateStatus alone misses new rows and
// only touches a few columns).
func (s *FineTuneStore) Upsert(job *FineTuneJobRecord) error {
if job.ID == "" {
job.ID = uuid.New().String()
}
now := time.Now()
if job.CreatedAt.IsZero() {
job.CreatedAt = now
}
job.UpdatedAt = now
return s.db.Clauses(clause.OnConflict{
Columns: []clause.Column{{Name: "id"}},
UpdateAll: true,
}).Create(job).Error
}
// UpdateStatus updates the status and message of a fine-tune job.
func (s *FineTuneStore) UpdateStatus(id, status, message string) error {
return s.db.Model(&FineTuneJobRecord{}).Where("id = ?", id).Updates(map[string]any{

View File

@@ -1,13 +0,0 @@
package distributed_test
import (
"testing"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
)
func TestDistributed(t *testing.T) {
RegisterFailHandler(Fail)
RunSpecs(t, "Distributed Suite")
}

View File

@@ -1,61 +0,0 @@
package distributed_test
import (
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"github.com/mudler/LocalAI/core/services/distributed"
"github.com/mudler/LocalAI/core/services/testutil"
)
var _ = Describe("FineTuneStore", func() {
var store *distributed.FineTuneStore
BeforeEach(func() {
db := testutil.SetupTestDB()
var err error
store, err = distributed.NewFineTuneStore(db)
Expect(err).ToNot(HaveOccurred())
})
Describe("ListAll", func() {
It("returns jobs across all users (unlike per-user List)", func() {
Expect(store.Create(&distributed.FineTuneJobRecord{ID: "j1", UserID: "u1", Status: "queued"})).To(Succeed())
Expect(store.Create(&distributed.FineTuneJobRecord{ID: "j2", UserID: "u2", Status: "queued"})).To(Succeed())
all, err := store.ListAll()
Expect(err).ToNot(HaveOccurred())
Expect(all).To(HaveLen(2))
perUser, err := store.List("u1")
Expect(err).ToNot(HaveOccurred())
Expect(perUser).To(HaveLen(1), "List stays per-user")
})
})
Describe("Upsert", func() {
It("inserts a new row", func() {
Expect(store.Upsert(&distributed.FineTuneJobRecord{ID: "up-1", UserID: "u1", Status: "queued"})).To(Succeed())
got, err := store.Get("up-1")
Expect(err).ToNot(HaveOccurred())
Expect(got.Status).To(Equal("queued"))
})
It("idempotently updates an existing row on a repeated key", func() {
Expect(store.Upsert(&distributed.FineTuneJobRecord{ID: "up-2", UserID: "u1", Status: "queued"})).To(Succeed())
// Second Upsert with the same primary key must update, not error on a
// duplicate-key violation (this is the SyncedMap write-through contract).
Expect(store.Upsert(&distributed.FineTuneJobRecord{ID: "up-2", UserID: "u1", Status: "completed", Message: "done"})).To(Succeed())
got, err := store.Get("up-2")
Expect(err).ToNot(HaveOccurred())
Expect(got.Status).To(Equal("completed"))
Expect(got.Message).To(Equal("done"))
all, err := store.ListAll()
Expect(err).ToNot(HaveOccurred())
Expect(all).To(HaveLen(1), "upsert must not create a duplicate")
})
})
})

View File

@@ -11,7 +11,6 @@ import (
type Stores struct {
Gallery *GalleryStore
FineTune *FineTuneStore
Quant *QuantStore
Skills *SkillStore
}
@@ -27,21 +26,15 @@ func InitStores(db *gorm.DB) (*Stores, error) {
return nil, fmt.Errorf("fine-tune store: %w", err)
}
quant, err := NewQuantStore(db)
if err != nil {
return nil, fmt.Errorf("quantization store: %w", err)
}
skills, err := NewSkillStore(db)
if err != nil {
return nil, fmt.Errorf("skills store: %w", err)
}
xlog.Info("Distributed stores initialized (Gallery, FineTune, Quant, Skills)")
xlog.Info("Distributed stores initialized (Gallery, FineTune, Skills)")
return &Stores{
Gallery: gallery,
FineTune: ft,
Quant: quant,
Skills: skills,
}, nil
}

View File

@@ -1,105 +0,0 @@
package distributed
import (
"context"
"fmt"
"time"
"github.com/google/uuid"
"github.com/mudler/LocalAI/core/services/advisorylock"
"gorm.io/gorm"
"gorm.io/gorm/clause"
)
// QuantJobRecord tracks quantization jobs in PostgreSQL. The columns mirror the
// API shape (schema.QuantizationJob); the structured Config and ExtraOptions are
// serialized into JSON text columns so a record fully reconstructs the job.
type QuantJobRecord struct {
ID string `gorm:"primaryKey;size:36" json:"id"`
UserID string `gorm:"index;size:36" json:"user_id,omitempty"`
Model string `gorm:"size:255" json:"model"`
Backend string `gorm:"size:64" json:"backend"`
ModelID string `gorm:"size:255" json:"model_id,omitempty"`
QuantizationType string `gorm:"size:32" json:"quantization_type"`
Status string `gorm:"index;size:32;default:queued" json:"status"` // queued, downloading, converting, quantizing, completed, failed, stopped
Message string `gorm:"type:text" json:"message,omitempty"`
OutputDir string `gorm:"size:512" json:"output_dir,omitempty"`
OutputFile string `gorm:"size:512" json:"output_file,omitempty"`
ConfigJSON string `gorm:"column:config;type:text" json:"-"`
ExtraOptsJSON string `gorm:"column:extra_options;type:text" json:"-"`
ImportStatus string `gorm:"size:32" json:"import_status,omitempty"`
ImportMessage string `gorm:"type:text" json:"import_message,omitempty"`
ImportModelName string `gorm:"size:255" json:"import_model_name,omitempty"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
func (QuantJobRecord) TableName() string { return "quantization_jobs" }
// QuantStore manages quantization job state in PostgreSQL.
type QuantStore struct {
db *gorm.DB
}
// NewQuantStore creates a new QuantStore and auto-migrates.
// Uses a PostgreSQL advisory lock to prevent concurrent migration races
// when multiple instances (frontend + workers) start at the same time.
func NewQuantStore(db *gorm.DB) (*QuantStore, error) {
if err := advisorylock.WithLockCtx(context.Background(), db, advisorylock.KeySchemaMigrate, func() error {
return db.AutoMigrate(&QuantJobRecord{})
}); err != nil {
return nil, fmt.Errorf("migrating quantization_jobs: %w", err)
}
return &QuantStore{db: db}, nil
}
// Create stores a new quantization job.
func (s *QuantStore) Create(job *QuantJobRecord) error {
if job.ID == "" {
job.ID = uuid.New().String()
}
job.CreatedAt = time.Now()
job.UpdatedAt = job.CreatedAt
return s.db.Create(job).Error
}
// Get retrieves a quantization job by ID.
func (s *QuantStore) Get(id string) (*QuantJobRecord, error) {
var job QuantJobRecord
if err := s.db.First(&job, "id = ?", id).Error; err != nil {
return nil, err
}
return &job, nil
}
// ListAll returns every quantization job across all users. The SyncedMap that
// backs QuantizationService is a single global map (the REST API filters by user
// at read time), so hydrate needs the full set.
func (s *QuantStore) ListAll() ([]QuantJobRecord, error) {
var jobs []QuantJobRecord
return jobs, s.db.Order("created_at DESC").Find(&jobs).Error
}
// Upsert idempotently inserts or fully replaces a job row by primary key. The
// SyncedMap write-through path issues a single Set per mutation regardless of
// whether the job already exists, so it needs one create-or-update primitive
// (Create alone fails on a duplicate key).
func (s *QuantStore) Upsert(job *QuantJobRecord) error {
if job.ID == "" {
job.ID = uuid.New().String()
}
now := time.Now()
if job.CreatedAt.IsZero() {
job.CreatedAt = now
}
job.UpdatedAt = now
return s.db.Clauses(clause.OnConflict{
Columns: []clause.Column{{Name: "id"}},
UpdateAll: true,
}).Create(job).Error
}
// Delete removes a quantization job.
func (s *QuantStore) Delete(id string) error {
return s.db.Where("id = ?", id).Delete(&QuantJobRecord{}).Error
}

View File

@@ -1,57 +0,0 @@
package distributed_test
import (
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"github.com/mudler/LocalAI/core/services/distributed"
"github.com/mudler/LocalAI/core/services/testutil"
)
var _ = Describe("QuantStore", func() {
var store *distributed.QuantStore
BeforeEach(func() {
db := testutil.SetupTestDB()
var err error
store, err = distributed.NewQuantStore(db)
Expect(err).ToNot(HaveOccurred())
})
Describe("ListAll", func() {
It("returns jobs across all users", func() {
Expect(store.Create(&distributed.QuantJobRecord{ID: "j1", UserID: "u1", Status: "queued"})).To(Succeed())
Expect(store.Create(&distributed.QuantJobRecord{ID: "j2", UserID: "u2", Status: "queued"})).To(Succeed())
all, err := store.ListAll()
Expect(err).ToNot(HaveOccurred())
Expect(all).To(HaveLen(2))
})
})
Describe("Upsert", func() {
It("inserts a new row", func() {
Expect(store.Upsert(&distributed.QuantJobRecord{ID: "up-1", UserID: "u1", Status: "queued"})).To(Succeed())
got, err := store.Get("up-1")
Expect(err).ToNot(HaveOccurred())
Expect(got.Status).To(Equal("queued"))
})
It("idempotently updates an existing row on a repeated key", func() {
Expect(store.Upsert(&distributed.QuantJobRecord{ID: "up-2", UserID: "u1", Status: "queued"})).To(Succeed())
// Second Upsert with the same primary key must update, not error on a
// duplicate-key violation (this is the SyncedMap write-through contract).
Expect(store.Upsert(&distributed.QuantJobRecord{ID: "up-2", UserID: "u1", Status: "completed", Message: "done"})).To(Succeed())
got, err := store.Get("up-2")
Expect(err).ToNot(HaveOccurred())
Expect(got.Status).To(Equal("completed"))
Expect(got.Message).To(Equal("done"))
all, err := store.ListAll()
Expect(err).ToNot(HaveOccurred())
Expect(all).To(HaveLen(1), "upsert must not create a duplicate")
})
})
})

View File

@@ -1,13 +0,0 @@
package finetune
import (
"testing"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
)
func TestFinetune(t *testing.T) {
RegisterFailHandler(Fail)
RunSpecs(t, "Finetune Suite")
}

View File

@@ -19,7 +19,6 @@ import (
"github.com/mudler/LocalAI/core/schema"
"github.com/mudler/LocalAI/core/services/distributed"
"github.com/mudler/LocalAI/core/services/messaging"
"github.com/mudler/LocalAI/core/services/syncstate"
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
"github.com/mudler/LocalAI/pkg/model"
"github.com/mudler/LocalAI/pkg/utils"
@@ -33,63 +32,44 @@ type FineTuneService struct {
modelLoader *model.ModelLoader
configLoader *config.ModelConfigLoader
// mu serializes the read-modify-write of job values. The SyncedMap guards its
// own map structure, but a job is a pointer mutated in place (e.g. the export
// goroutine), so the service still needs a lock to keep those field updates
// and the subsequent Set atomic with respect to readers.
mu sync.Mutex
mu sync.Mutex
jobs map[string]*schema.FineTuneJob
// jobs is the cross-replica job store: an in-memory map kept consistent across
// replicas via NATS, optionally read-through to PostgreSQL in distributed mode.
jobs *syncstate.SyncedMap[string, *schema.FineTuneJob]
// Distributed mode (nil when not in distributed mode)
natsClient messaging.Publisher
fineTuneStore *distributed.FineTuneStore
}
// NewFineTuneService creates a new FineTuneService. In distributed mode pass the
// shared NATS client and PostgreSQL store so jobs stay consistent across
// replicas; pass nil for both in standalone mode, where the disk Loader hydrates
// the map and there is nothing to broadcast.
// SetNATSClient sets the NATS client for distributed progress publishing.
func (s *FineTuneService) SetNATSClient(nc messaging.Publisher) {
s.mu.Lock()
defer s.mu.Unlock()
s.natsClient = nc
}
// SetFineTuneStore sets the PostgreSQL fine-tune store for distributed persistence.
func (s *FineTuneService) SetFineTuneStore(store *distributed.FineTuneStore) {
s.mu.Lock()
defer s.mu.Unlock()
s.fineTuneStore = store
}
// NewFineTuneService creates a new FineTuneService.
func NewFineTuneService(
appConfig *config.ApplicationConfig,
modelLoader *model.ModelLoader,
configLoader *config.ModelConfigLoader,
nats messaging.MessagingClient,
store *distributed.FineTuneStore,
) *FineTuneService {
s := &FineTuneService{
appConfig: appConfig,
modelLoader: modelLoader,
configLoader: configLoader,
jobs: make(map[string]*schema.FineTuneJob),
}
// Only attach a Store interface when a concrete store exists, otherwise the
// SyncedMap would see a non-nil interface wrapping a nil pointer and try to
// hydrate/write through a nil DB.
var syncStore syncstate.Store[string, *schema.FineTuneJob]
if store != nil {
syncStore = &fineTuneStoreAdapter{store: store}
}
s.jobs = syncstate.New(syncstate.Config[string, *schema.FineTuneJob]{
Name: "finetune.jobs",
Key: func(j *schema.FineTuneJob) string { return j.ID },
Nats: nats,
Store: syncStore,
Loader: s.loadJobsFromDisk, // ignored when Store is set (distributed mode)
})
// Hydrate + subscribe. A hydrate failure must not take the server down: log
// and continue degraded (standalone), mirroring the OpCache wiring.
if err := s.jobs.Start(appConfig.Context); err != nil {
xlog.Warn("FineTune SyncedMap start failed; running degraded", "error", err)
}
s.loadAllJobs()
return s
}
// Close releases the SyncedMap subscription and background workers.
func (s *FineTuneService) Close() error {
return s.jobs.Close()
}
// fineTuneBaseDir returns the base directory for fine-tune job data.
func (s *FineTuneService) fineTuneBaseDir() string {
return filepath.Join(s.appConfig.DataPath, "fine-tune")
@@ -120,18 +100,15 @@ func (s *FineTuneService) saveJobState(job *schema.FineTuneJob) {
}
}
// loadJobsFromDisk scans the fine-tune directory for persisted jobs and returns
// them. It is the SyncedMap Loader used in standalone mode (no DB); the returned
// slice hydrates the map on Start.
func (s *FineTuneService) loadJobsFromDisk(_ context.Context) ([]*schema.FineTuneJob, error) {
// loadAllJobs scans the fine-tune directory for persisted jobs and loads them.
func (s *FineTuneService) loadAllJobs() {
baseDir := s.fineTuneBaseDir()
entries, err := os.ReadDir(baseDir)
if err != nil {
// Directory doesn't exist yet — that's fine, start empty.
return nil, nil
// Directory doesn't exist yet — that's fine
return
}
var jobs []*schema.FineTuneJob
for _, entry := range entries {
if !entry.IsDir() {
continue
@@ -160,13 +137,12 @@ func (s *FineTuneService) loadJobsFromDisk(_ context.Context) ([]*schema.FineTun
job.ExportMessage = "Server restarted while export was running"
}
jobs = append(jobs, &job)
s.jobs[job.ID] = &job
}
if len(jobs) > 0 {
xlog.Info("Loaded persisted fine-tune jobs", "count", len(jobs))
if len(s.jobs) > 0 {
xlog.Info("Loaded persisted fine-tune jobs", "count", len(s.jobs))
}
return jobs, nil
}
// StartJob starts a new fine-tuning job.
@@ -260,14 +236,28 @@ func (s *FineTuneService) StartJob(ctx context.Context, userID string, req schem
CreatedAt: time.Now().UTC().Format(time.RFC3339),
Config: &req,
}
// Set write-through persists to PostgreSQL (distributed) and broadcasts to
// peer replicas; the disk state.json is written separately for restart
// recovery / standalone hydrate.
if err := s.jobs.Set(ctx, job); err != nil {
return nil, fmt.Errorf("failed to persist job: %w", err)
}
s.jobs[jobID] = job
s.saveJobState(job)
// Persist to PostgreSQL in distributed mode
if s.fineTuneStore != nil {
configJSON, _ := json.Marshal(req)
extraJSON, _ := json.Marshal(req.ExtraOptions)
s.fineTuneStore.Create(&distributed.FineTuneJobRecord{
ID: jobID,
UserID: userID,
Model: req.Model,
Backend: backendName,
ModelID: modelID,
TrainingType: req.TrainingType,
TrainingMethod: req.TrainingMethod,
Status: "queued",
OutputDir: outputDir,
ConfigJSON: string(configJSON),
ExtraOptsJSON: string(extraJSON),
})
}
return &schema.FineTuneJobResponse{
ID: jobID,
Status: "queued",
@@ -280,7 +270,7 @@ func (s *FineTuneService) GetJob(userID, jobID string) (*schema.FineTuneJob, err
s.mu.Lock()
defer s.mu.Unlock()
job, ok := s.jobs.Get(jobID)
job, ok := s.jobs[jobID]
if !ok {
return nil, fmt.Errorf("job not found: %s", jobID)
}
@@ -296,7 +286,7 @@ func (s *FineTuneService) ListJobs(userID string) []*schema.FineTuneJob {
defer s.mu.Unlock()
var result []*schema.FineTuneJob
for _, job := range s.jobs.List() {
for _, job := range s.jobs {
if userID == "" || job.UserID == userID {
result = append(result, job)
}
@@ -312,7 +302,7 @@ func (s *FineTuneService) ListJobs(userID string) []*schema.FineTuneJob {
// StopJob stops a running fine-tuning job.
func (s *FineTuneService) StopJob(ctx context.Context, userID, jobID string, saveCheckpoint bool) error {
s.mu.Lock()
job, ok := s.jobs.Get(jobID)
job, ok := s.jobs[jobID]
if !ok {
s.mu.Unlock()
return fmt.Errorf("job not found: %s", jobID)
@@ -333,10 +323,10 @@ func (s *FineTuneService) StopJob(ctx context.Context, userID, jobID string, sav
s.mu.Lock()
job.Status = "stopped"
job.Message = "Training stopped by user"
if err := s.jobs.Set(ctx, job); err != nil {
xlog.Warn("Failed to persist stopped job", "job_id", jobID, "error", err)
}
s.saveJobState(job)
if s.fineTuneStore != nil {
s.fineTuneStore.UpdateStatus(jobID, "stopped", "Training stopped by user")
}
s.mu.Unlock()
return nil
@@ -345,7 +335,7 @@ func (s *FineTuneService) StopJob(ctx context.Context, userID, jobID string, sav
// DeleteJob removes a fine-tuning job and its associated data from disk.
func (s *FineTuneService) DeleteJob(userID, jobID string) error {
s.mu.Lock()
job, ok := s.jobs.Get(jobID)
job, ok := s.jobs[jobID]
if !ok {
s.mu.Unlock()
return fmt.Errorf("job not found: %s", jobID)
@@ -370,10 +360,9 @@ func (s *FineTuneService) DeleteJob(userID, jobID string) error {
}
exportModelName := job.ExportModelName
// Delete write-through removes the DB row (distributed) and broadcasts the
// removal to peer replicas. DeleteJob has no ctx, so use Background.
if err := s.jobs.Delete(context.Background(), jobID); err != nil {
xlog.Warn("Failed to delete job from store", "job_id", jobID, "error", err)
delete(s.jobs, jobID)
if s.fineTuneStore != nil {
s.fineTuneStore.Delete(jobID)
}
s.mu.Unlock()
@@ -409,7 +398,7 @@ func (s *FineTuneService) DeleteJob(userID, jobID string) error {
// StreamProgress opens a gRPC progress stream and calls the callback for each update.
func (s *FineTuneService) StreamProgress(ctx context.Context, userID, jobID string, callback func(event *schema.FineTuneProgressEvent)) error {
s.mu.Lock()
job, ok := s.jobs.Get(jobID)
job, ok := s.jobs[jobID]
if !ok {
s.mu.Unlock()
return fmt.Errorf("job not found: %s", jobID)
@@ -438,7 +427,7 @@ func (s *FineTuneService) StreamProgress(ctx context.Context, userID, jobID stri
}, func(update *pb.FineTuneProgressUpdate) {
// Update job status and persist
s.mu.Lock()
if j, ok := s.jobs.Get(jobID); ok {
if j, ok := s.jobs[jobID]; ok {
// Don't let progress updates overwrite terminal states
isTerminal := j.Status == "stopped" || j.Status == "completed" || j.Status == "failed"
if !isTerminal {
@@ -447,10 +436,10 @@ func (s *FineTuneService) StreamProgress(ctx context.Context, userID, jobID stri
if update.Message != "" {
j.Message = update.Message
}
if err := s.jobs.Set(ctx, j); err != nil {
xlog.Warn("Failed to persist progress update", "job_id", jobID, "error", err)
}
s.saveJobState(j)
if s.fineTuneStore != nil {
s.fineTuneStore.UpdateStatus(jobID, j.Status, j.Message)
}
}
s.mu.Unlock()
@@ -485,7 +474,7 @@ func (s *FineTuneService) StreamProgress(ctx context.Context, userID, jobID stri
// ListCheckpoints lists checkpoints for a job.
func (s *FineTuneService) ListCheckpoints(ctx context.Context, userID, jobID string) ([]*pb.CheckpointInfo, error) {
s.mu.Lock()
job, ok := s.jobs.Get(jobID)
job, ok := s.jobs[jobID]
if !ok {
s.mu.Unlock()
return nil, fmt.Errorf("job not found: %s", jobID)
@@ -531,7 +520,7 @@ func sanitizeModelName(s string) string {
// ExportModel starts an async model export from a checkpoint and returns the intended model name immediately.
func (s *FineTuneService) ExportModel(ctx context.Context, userID, jobID string, req schema.ExportRequest) (string, error) {
s.mu.Lock()
job, ok := s.jobs.Get(jobID)
job, ok := s.jobs[jobID]
if !ok {
s.mu.Unlock()
return "", fmt.Errorf("job not found: %s", jobID)
@@ -583,9 +572,6 @@ func (s *FineTuneService) ExportModel(ctx context.Context, userID, jobID string,
job.ExportStatus = "exporting"
job.ExportMessage = ""
job.ExportModelName = ""
if err := s.jobs.Set(ctx, job); err != nil {
xlog.Warn("Failed to persist export start", "job_id", jobID, "error", err)
}
s.saveJobState(job)
s.mu.Unlock()
@@ -676,30 +662,24 @@ func (s *FineTuneService) ExportModel(ctx context.Context, userID, jobID string,
xlog.Info("Model exported and registered", "job_id", jobID, "model_name", modelName, "format", req.ExportFormat)
// Runs after the HTTP request returns, so use Background rather than the
// (now likely cancelled) request ctx for the write-through.
s.mu.Lock()
job.ExportStatus = "completed"
job.ExportModelName = modelName
job.ExportMessage = ""
if err := s.jobs.Set(context.Background(), job); err != nil {
xlog.Warn("Failed to persist export completion", "job_id", jobID, "error", err)
}
s.saveJobState(job)
if s.fineTuneStore != nil {
s.fineTuneStore.UpdateExportStatus(jobID, "completed", "", modelName)
}
s.mu.Unlock()
}()
return modelName, nil
}
// setExportMessage updates the export message and persists the job state. Called
// from the background export goroutine, so it uses Background for write-through.
// setExportMessage updates the export message and persists the job state.
func (s *FineTuneService) setExportMessage(job *schema.FineTuneJob, msg string) {
s.mu.Lock()
job.ExportMessage = msg
if err := s.jobs.Set(context.Background(), job); err != nil {
xlog.Warn("Failed to persist export message", "job_id", job.ID, "error", err)
}
s.saveJobState(job)
s.mu.Unlock()
}
@@ -707,7 +687,7 @@ func (s *FineTuneService) setExportMessage(job *schema.FineTuneJob, msg string)
// GetExportedModelPath returns the path to the exported model directory and its name.
func (s *FineTuneService) GetExportedModelPath(userID, jobID string) (string, string, error) {
s.mu.Lock()
job, ok := s.jobs.Get(jobID)
job, ok := s.jobs[jobID]
if !ok {
s.mu.Unlock()
return "", "", fmt.Errorf("job not found: %s", jobID)
@@ -743,10 +723,10 @@ func (s *FineTuneService) setExportFailed(job *schema.FineTuneJob, message strin
s.mu.Lock()
job.ExportStatus = "failed"
job.ExportMessage = message
if err := s.jobs.Set(context.Background(), job); err != nil {
xlog.Warn("Failed to persist export failure", "job_id", job.ID, "error", err)
}
s.saveJobState(job)
if s.fineTuneStore != nil {
s.fineTuneStore.UpdateExportStatus(job.ID, "failed", message, "")
}
s.mu.Unlock()
}

View File

@@ -1,185 +0,0 @@
package finetune
// White-box tests (package finetune) so a spec can drive the service's internal
// SyncedMap the same way StartJob does (via jobs.Set) without standing up a
// training backend, then assert the cross-replica reads (GetJob/ListJobs) and
// the adapter conversions that keep REST responses byte-for-byte unchanged.
import (
"context"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"github.com/mudler/LocalAI/core/config"
"github.com/mudler/LocalAI/core/schema"
"github.com/mudler/LocalAI/core/services/distributed"
"github.com/mudler/LocalAI/core/services/testutil"
)
// newTestService builds a standalone FineTuneService wired to the given bus. The
// model/config loaders are nil because the read/sync paths under test never touch
// them; the data dir is a throwaway temp dir so the disk Loader finds nothing.
func newTestService(bus *testutil.FakeBus) *FineTuneService {
appConfig := &config.ApplicationConfig{
Context: context.Background(),
DataPath: GinkgoT().TempDir(),
}
return NewFineTuneService(appConfig, nil, nil, bus, nil)
}
var _ = Describe("FineTuneService", func() {
ctx := context.Background()
Describe("cross-replica job visibility", func() {
var (
bus *testutil.FakeBus
a, b *FineTuneService
)
BeforeEach(func() {
// One shared bus, two replicas: exactly the distributed topology where
// a round-robin request may land on a replica that did not originate
// the change.
bus = testutil.NewFakeBus()
a = newTestService(bus)
b = newTestService(bus)
})
AfterEach(func() {
Expect(a.Close()).To(Succeed())
Expect(b.Close()).To(Succeed())
})
It("makes a job created on A visible via B's GetJob and ListJobs", func() {
job := &schema.FineTuneJob{ID: "job-1", UserID: "user-1", Status: "queued", CreatedAt: "2026-06-27T10:00:00Z"}
// StartJob persists via jobs.Set; drive that directly to avoid a backend.
Expect(a.jobs.Set(ctx, job)).To(Succeed())
got, err := b.GetJob("user-1", "job-1")
Expect(err).ToNot(HaveOccurred(), "B must see a job A just created")
Expect(got.Status).To(Equal("queued"))
listed := b.ListJobs("user-1")
Expect(listed).To(HaveLen(1))
Expect(listed[0].ID).To(Equal("job-1"))
})
It("removes a job from B when it is deleted on A", func() {
job := &schema.FineTuneJob{ID: "job-2", UserID: "user-1", Status: "completed", CreatedAt: "2026-06-27T10:00:00Z"}
Expect(a.jobs.Set(ctx, job)).To(Succeed())
_, err := b.GetJob("user-1", "job-2")
Expect(err).ToNot(HaveOccurred(), "precondition: B must have the job before the delete")
Expect(a.jobs.Delete(ctx, "job-2")).To(Succeed())
_, err = b.GetJob("user-1", "job-2")
Expect(err).To(HaveOccurred(), "a delete on A must remove the job from B")
})
It("propagates a status update from A to B", func() {
job := &schema.FineTuneJob{ID: "job-3", UserID: "user-1", Status: "training", CreatedAt: "2026-06-27T10:00:00Z"}
Expect(a.jobs.Set(ctx, job)).To(Succeed())
updated := &schema.FineTuneJob{ID: "job-3", UserID: "user-1", Status: "completed", CreatedAt: "2026-06-27T10:00:00Z"}
Expect(a.jobs.Set(ctx, updated)).To(Succeed())
got, err := b.GetJob("user-1", "job-3")
Expect(err).ToNot(HaveOccurred())
Expect(got.Status).To(Equal("completed"))
})
})
Describe("ListJobs", func() {
var svc *FineTuneService
BeforeEach(func() {
svc = newTestService(testutil.NewFakeBus())
})
AfterEach(func() { Expect(svc.Close()).To(Succeed()) })
It("filters by user and sorts newest-first", func() {
Expect(svc.jobs.Set(ctx, &schema.FineTuneJob{ID: "old", UserID: "u1", CreatedAt: "2026-06-25T10:00:00Z"})).To(Succeed())
Expect(svc.jobs.Set(ctx, &schema.FineTuneJob{ID: "new", UserID: "u1", CreatedAt: "2026-06-27T10:00:00Z"})).To(Succeed())
Expect(svc.jobs.Set(ctx, &schema.FineTuneJob{ID: "other", UserID: "u2", CreatedAt: "2026-06-26T10:00:00Z"})).To(Succeed())
jobs := svc.ListJobs("u1")
Expect(jobs).To(HaveLen(2), "only u1's jobs")
Expect(jobs[0].ID).To(Equal("new"), "newest first")
Expect(jobs[1].ID).To(Equal("old"))
})
It("returns every user's jobs when the userID filter is empty", func() {
Expect(svc.jobs.Set(ctx, &schema.FineTuneJob{ID: "a", UserID: "u1", CreatedAt: "2026-06-25T10:00:00Z"})).To(Succeed())
Expect(svc.jobs.Set(ctx, &schema.FineTuneJob{ID: "b", UserID: "u2", CreatedAt: "2026-06-26T10:00:00Z"})).To(Succeed())
Expect(svc.ListJobs("")).To(HaveLen(2))
})
It("rejects GetJob for a job owned by another user", func() {
Expect(svc.jobs.Set(ctx, &schema.FineTuneJob{ID: "x", UserID: "owner", CreatedAt: "2026-06-25T10:00:00Z"})).To(Succeed())
_, err := svc.GetJob("intruder", "x")
Expect(err).To(HaveOccurred(), "a different user must not read someone else's job")
})
})
Describe("store adapter conversion", func() {
// The SyncedMap value type is *schema.FineTuneJob (the exact REST shape).
// These specs prove the DB adapter round-trips it losslessly, so hydrate
// and write-through in distributed mode keep responses unchanged.
It("round-trips a job through jobToRecord/recordToJob preserving the API shape", func() {
original := &schema.FineTuneJob{
ID: "rt-1",
UserID: "user-1",
Model: "base-model",
Backend: "trl",
ModelID: "trl-finetune-rt-1",
TrainingType: "lora",
TrainingMethod: "sft",
Status: "completed",
Message: "done",
OutputDir: "/data/fine-tune/rt-1",
ExtraOptions: map[string]string{"hf_token": "secret"},
CreatedAt: "2026-06-27T10:00:00Z",
ExportStatus: "completed",
ExportMessage: "",
ExportModelName: "base-model-ft-rt-1",
Config: &schema.FineTuneJobRequest{Model: "base-model", Backend: "trl", DatasetSource: "data.jsonl"},
}
rec := jobToRecord(original)
Expect(rec.ID).To(Equal("rt-1"))
Expect(rec.ConfigJSON).ToNot(BeEmpty(), "structured config must serialize into the JSON column")
Expect(rec.ExtraOptsJSON).ToNot(BeEmpty())
back := recordToJob(rec)
Expect(back.ID).To(Equal(original.ID))
Expect(back.UserID).To(Equal(original.UserID))
Expect(back.Model).To(Equal(original.Model))
Expect(back.Backend).To(Equal(original.Backend))
Expect(back.ModelID).To(Equal(original.ModelID))
Expect(back.TrainingType).To(Equal(original.TrainingType))
Expect(back.TrainingMethod).To(Equal(original.TrainingMethod))
Expect(back.Status).To(Equal(original.Status))
Expect(back.Message).To(Equal(original.Message))
Expect(back.OutputDir).To(Equal(original.OutputDir))
Expect(back.ExportStatus).To(Equal(original.ExportStatus))
Expect(back.ExportModelName).To(Equal(original.ExportModelName))
Expect(back.CreatedAt).To(Equal(original.CreatedAt))
Expect(back.ExtraOptions).To(Equal(original.ExtraOptions))
Expect(back.Config).ToNot(BeNil())
Expect(back.Config.DatasetSource).To(Equal("data.jsonl"))
})
})
Describe("compile-time adapter contract", func() {
It("satisfies syncstate.Store for *distributed.FineTuneStore", func() {
// Guards against drift between the adapter and the component interface;
// the var assertion in syncstore.go covers it at build time, this keeps
// the type referenced from a spec too.
var _ *distributed.FineTuneStore
Expect(&fineTuneStoreAdapter{}).ToNot(BeNil())
})
})
})

View File

@@ -1,114 +0,0 @@
package finetune
import (
"context"
"encoding/json"
"time"
"github.com/mudler/LocalAI/core/schema"
"github.com/mudler/LocalAI/core/services/distributed"
"github.com/mudler/LocalAI/core/services/syncstate"
)
// fineTuneStoreAdapter bridges the distributed PostgreSQL FineTuneStore to the
// generic syncstate.Store the SyncedMap consumes. It is only wired in distributed
// mode; standalone leaves Store nil and hydrates from disk via a Loader instead.
//
// The SyncedMap value type is *schema.FineTuneJob (the exact shape the REST API
// returns) so reads need no conversion and the response JSON is provably
// unchanged. The adapter is the single place that translates between that API
// shape and the DB FineTuneJobRecord.
type fineTuneStoreAdapter struct {
store *distributed.FineTuneStore
}
// compile-time assertion that the adapter satisfies the component's Store.
var _ syncstate.Store[string, *schema.FineTuneJob] = (*fineTuneStoreAdapter)(nil)
func (a *fineTuneStoreAdapter) List(_ context.Context) ([]*schema.FineTuneJob, error) {
records, err := a.store.ListAll()
if err != nil {
return nil, err
}
jobs := make([]*schema.FineTuneJob, 0, len(records))
for i := range records {
jobs = append(jobs, recordToJob(&records[i]))
}
return jobs, nil
}
func (a *fineTuneStoreAdapter) Upsert(_ context.Context, job *schema.FineTuneJob) error {
return a.store.Upsert(jobToRecord(job))
}
func (a *fineTuneStoreAdapter) Delete(_ context.Context, id string) error {
return a.store.Delete(id)
}
// recordToJob maps a persisted DB record back to the API shape, reconstructing
// the structured Config / ExtraOptions from their JSON columns.
func recordToJob(r *distributed.FineTuneJobRecord) *schema.FineTuneJob {
job := &schema.FineTuneJob{
ID: r.ID,
UserID: r.UserID,
Model: r.Model,
Backend: r.Backend,
ModelID: r.ModelID,
TrainingType: r.TrainingType,
TrainingMethod: r.TrainingMethod,
Status: r.Status,
Message: r.Message,
OutputDir: r.OutputDir,
ExportStatus: r.ExportStatus,
ExportMessage: r.ExportMessage,
ExportModelName: r.ExportModelName,
CreatedAt: r.CreatedAt.UTC().Format(time.RFC3339),
}
if r.ExtraOptsJSON != "" {
// Best-effort: a malformed column must not drop the whole job from the API.
_ = json.Unmarshal([]byte(r.ExtraOptsJSON), &job.ExtraOptions)
}
if r.ConfigJSON != "" {
var cfg schema.FineTuneJobRequest
if err := json.Unmarshal([]byte(r.ConfigJSON), &cfg); err == nil {
job.Config = &cfg
}
}
return job
}
// jobToRecord maps the API shape to a DB record for write-through, serializing
// the structured Config / ExtraOptions into their JSON columns. CreatedAt is
// parsed back from the RFC3339 string the service stamps; an unparseable value
// is left zero so FineTuneStore.Upsert stamps "now".
func jobToRecord(job *schema.FineTuneJob) *distributed.FineTuneJobRecord {
rec := &distributed.FineTuneJobRecord{
ID: job.ID,
UserID: job.UserID,
Model: job.Model,
Backend: job.Backend,
ModelID: job.ModelID,
TrainingType: job.TrainingType,
TrainingMethod: job.TrainingMethod,
Status: job.Status,
Message: job.Message,
OutputDir: job.OutputDir,
ExportStatus: job.ExportStatus,
ExportMessage: job.ExportMessage,
ExportModelName: job.ExportModelName,
}
if job.Config != nil {
if data, err := json.Marshal(job.Config); err == nil {
rec.ConfigJSON = string(data)
}
}
if job.ExtraOptions != nil {
if data, err := json.Marshal(job.ExtraOptions); err == nil {
rec.ExtraOptsJSON = string(data)
}
}
if t, err := time.Parse(time.RFC3339, job.CreatedAt); err == nil {
rec.CreatedAt = t
}
return rec
}

View File

@@ -22,14 +22,6 @@ const subscribeConfirmTimeout = 5 * time.Second
type Client struct {
conn *nats.Conn
mu sync.RWMutex
// reconnectCbs are invoked after the underlying connection is
// re-established. nats.go transparently resubscribes existing
// subscriptions on reconnect, but it cannot know that a consumer kept
// derived in-memory state (e.g. syncstate.SyncedMap) that may have drifted
// while the link was down — these callbacks let such consumers re-hydrate.
cbMu sync.Mutex
reconnectCbs []func()
}
// New creates a new NATS client with auto-reconnect.
@@ -39,10 +31,6 @@ func New(url string, opts ...Option) (*Client, error) {
o(&cfg)
}
// Allocate the client up front so the reconnect handler closure can reach
// it; conn is populated after nats.Connect succeeds below.
c := &Client{}
natsOpts := []nats.Option{
nats.RetryOnFailedConnect(true),
nats.MaxReconnects(-1),
@@ -53,7 +41,6 @@ func New(url string, opts ...Option) (*Client, error) {
}),
nats.ReconnectHandler(func(_ *nats.Conn) {
xlog.Info("NATS reconnected")
c.runReconnectCallbacks()
}),
nats.ClosedHandler(func(_ *nats.Conn) {
xlog.Info("NATS connection closed")
@@ -116,33 +103,7 @@ func New(url string, opts ...Option) (*Client, error) {
return nil, fmt.Errorf("connecting to NATS at %s: %w", sanitize.URL(url), err)
}
c.conn = nc
return c, nil
}
// OnReconnect registers a callback invoked after the NATS connection is
// re-established. It is consumed via an optional interface type-assertion
// (interface{ OnReconnect(func()) }) rather than being added to MessagingClient,
// so the messaging abstraction stays minimal and standalone/test clients are not
// forced to implement reconnect semantics. A nil callback is ignored.
func (c *Client) OnReconnect(cb func()) {
if cb == nil {
return
}
c.cbMu.Lock()
c.reconnectCbs = append(c.reconnectCbs, cb)
c.cbMu.Unlock()
}
// runReconnectCallbacks invokes registered reconnect callbacks. It copies the
// slice under the lock so a callback that (re)registers cannot deadlock.
func (c *Client) runReconnectCallbacks() {
c.cbMu.Lock()
cbs := append([]func(){}, c.reconnectCbs...)
c.cbMu.Unlock()
for _, cb := range cbs {
cb()
}
return &Client{conn: nc}, nil
}
// Publish marshals data as JSON and publishes it to the given subject.

View File

@@ -380,20 +380,6 @@ func SubjectCacheInvalidateCollection(name string) string {
return "cache.invalidate.collections." + sanitizeSubjectToken(name)
}
// SyncedMap State Sync (Pub/Sub — broadcast to all frontends)
//
// The reusable syncstate.SyncedMap component publishes a {op,key,value} delta on
// this subject whenever a replica mutates a piece of cross-replica in-memory
// state. Peers subscribe and apply the delta to their own map, so a round-robin
// API request that lands on a replica which did not originate the change still
// sees it. Convergence on (re)connect is done by re-hydrating from the durable
// source, so no request/reply snapshot subject is needed here.
func SubjectSyncStateDelta(name string) string {
return subjectSyncStatePrefix + sanitizeSubjectToken(name) + ".delta"
}
const subjectSyncStatePrefix = "state."
// Prefix-Cache Routing Sync (Pub/Sub - broadcast to all frontends)
//
// Frontends share prefix-cache observations so a request routed to any replica

View File

@@ -1,13 +0,0 @@
package quantization
import (
"testing"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
)
func TestQuantization(t *testing.T) {
RegisterFailHandler(Fail)
RunSpecs(t, "Quantization Suite")
}

View File

@@ -17,9 +17,6 @@ import (
"github.com/mudler/LocalAI/core/config"
"github.com/mudler/LocalAI/core/gallery/importers"
"github.com/mudler/LocalAI/core/schema"
"github.com/mudler/LocalAI/core/services/distributed"
"github.com/mudler/LocalAI/core/services/messaging"
"github.com/mudler/LocalAI/core/services/syncstate"
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
"github.com/mudler/LocalAI/pkg/model"
"github.com/mudler/LocalAI/pkg/utils"
@@ -33,63 +30,26 @@ type QuantizationService struct {
modelLoader *model.ModelLoader
configLoader *config.ModelConfigLoader
// mu serializes the read-modify-write of job values. The SyncedMap guards its
// own map structure, but a job is a pointer mutated in place (e.g. the import
// goroutine), so the service still needs a lock to keep those field updates and
// the subsequent Set atomic with respect to readers.
mu sync.Mutex
// jobs is the cross-replica job store: an in-memory map kept consistent across
// replicas via NATS, optionally read-through to PostgreSQL in distributed mode.
jobs *syncstate.SyncedMap[string, *schema.QuantizationJob]
mu sync.Mutex
jobs map[string]*schema.QuantizationJob
}
// NewQuantizationService creates a new QuantizationService. In distributed mode
// pass the shared NATS client and PostgreSQL store so jobs stay consistent across
// replicas; pass nil for both in standalone mode, where the disk Loader hydrates
// the map and there is nothing to broadcast.
// NewQuantizationService creates a new QuantizationService.
func NewQuantizationService(
appConfig *config.ApplicationConfig,
modelLoader *model.ModelLoader,
configLoader *config.ModelConfigLoader,
nats messaging.MessagingClient,
store *distributed.QuantStore,
) *QuantizationService {
s := &QuantizationService{
appConfig: appConfig,
modelLoader: modelLoader,
configLoader: configLoader,
jobs: make(map[string]*schema.QuantizationJob),
}
// Only attach a Store interface when a concrete store exists, otherwise the
// SyncedMap would see a non-nil interface wrapping a nil pointer and try to
// hydrate/write through a nil DB.
var syncStore syncstate.Store[string, *schema.QuantizationJob]
if store != nil {
syncStore = &quantStoreAdapter{store: store}
}
s.jobs = syncstate.New(syncstate.Config[string, *schema.QuantizationJob]{
Name: "quant.jobs",
Key: func(j *schema.QuantizationJob) string { return j.ID },
Nats: nats,
Store: syncStore,
Loader: s.loadJobsFromDisk, // ignored when Store is set (distributed mode)
})
// Hydrate + subscribe. A hydrate failure must not take the server down: log and
// continue degraded (standalone), mirroring the FineTune/OpCache wiring.
if err := s.jobs.Start(appConfig.Context); err != nil {
xlog.Warn("Quantization SyncedMap start failed; running degraded", "error", err)
}
s.loadAllJobs()
return s
}
// Close releases the SyncedMap subscription and background workers.
func (s *QuantizationService) Close() error {
return s.jobs.Close()
}
// quantizationBaseDir returns the base directory for quantization job data.
func (s *QuantizationService) quantizationBaseDir() string {
return filepath.Join(s.appConfig.DataPath, "quantization")
@@ -120,18 +80,15 @@ func (s *QuantizationService) saveJobState(job *schema.QuantizationJob) {
}
}
// loadJobsFromDisk scans the quantization directory for persisted jobs and
// returns them. It is the SyncedMap Loader used in standalone mode (no DB); the
// returned slice hydrates the map on Start.
func (s *QuantizationService) loadJobsFromDisk(_ context.Context) ([]*schema.QuantizationJob, error) {
// loadAllJobs scans the quantization directory for persisted jobs and loads them.
func (s *QuantizationService) loadAllJobs() {
baseDir := s.quantizationBaseDir()
entries, err := os.ReadDir(baseDir)
if err != nil {
// Directory doesn't exist yet — that's fine, start empty.
return nil, nil
// Directory doesn't exist yet — that's fine
return
}
var jobs []*schema.QuantizationJob
for _, entry := range entries {
if !entry.IsDir() {
continue
@@ -160,13 +117,12 @@ func (s *QuantizationService) loadJobsFromDisk(_ context.Context) ([]*schema.Qua
job.ImportMessage = "Server restarted while import was running"
}
jobs = append(jobs, &job)
s.jobs[job.ID] = &job
}
if len(jobs) > 0 {
xlog.Info("Loaded persisted quantization jobs", "count", len(jobs))
if len(s.jobs) > 0 {
xlog.Info("Loaded persisted quantization jobs", "count", len(s.jobs))
}
return jobs, nil
}
// StartJob starts a new quantization job.
@@ -232,12 +188,7 @@ func (s *QuantizationService) StartJob(ctx context.Context, userID string, req s
CreatedAt: time.Now().UTC().Format(time.RFC3339),
Config: &req,
}
// Set write-through persists to PostgreSQL (distributed) and broadcasts to
// peer replicas; the disk state.json is written separately for restart
// recovery / standalone hydrate.
if err := s.jobs.Set(ctx, job); err != nil {
return nil, fmt.Errorf("failed to persist job: %w", err)
}
s.jobs[jobID] = job
s.saveJobState(job)
return &schema.QuantizationJobResponse{
@@ -252,7 +203,7 @@ func (s *QuantizationService) GetJob(userID, jobID string) (*schema.Quantization
s.mu.Lock()
defer s.mu.Unlock()
job, ok := s.jobs.Get(jobID)
job, ok := s.jobs[jobID]
if !ok {
return nil, fmt.Errorf("job not found: %s", jobID)
}
@@ -268,7 +219,7 @@ func (s *QuantizationService) ListJobs(userID string) []*schema.QuantizationJob
defer s.mu.Unlock()
var result []*schema.QuantizationJob
for _, job := range s.jobs.List() {
for _, job := range s.jobs {
if userID == "" || job.UserID == userID {
result = append(result, job)
}
@@ -284,7 +235,7 @@ func (s *QuantizationService) ListJobs(userID string) []*schema.QuantizationJob
// StopJob stops a running quantization job.
func (s *QuantizationService) StopJob(ctx context.Context, userID, jobID string) error {
s.mu.Lock()
job, ok := s.jobs.Get(jobID)
job, ok := s.jobs[jobID]
if !ok {
s.mu.Unlock()
return fmt.Errorf("job not found: %s", jobID)
@@ -305,9 +256,6 @@ func (s *QuantizationService) StopJob(ctx context.Context, userID, jobID string)
s.mu.Lock()
job.Status = "stopped"
job.Message = "Quantization stopped by user"
if err := s.jobs.Set(ctx, job); err != nil {
xlog.Warn("Failed to persist stopped job", "job_id", jobID, "error", err)
}
s.saveJobState(job)
s.mu.Unlock()
@@ -317,7 +265,7 @@ func (s *QuantizationService) StopJob(ctx context.Context, userID, jobID string)
// DeleteJob removes a quantization job and its associated data from disk.
func (s *QuantizationService) DeleteJob(userID, jobID string) error {
s.mu.Lock()
job, ok := s.jobs.Get(jobID)
job, ok := s.jobs[jobID]
if !ok {
s.mu.Unlock()
return fmt.Errorf("job not found: %s", jobID)
@@ -341,11 +289,7 @@ func (s *QuantizationService) DeleteJob(userID, jobID string) error {
}
importModelName := job.ImportModelName
// Delete write-through removes the DB row (distributed) and broadcasts the
// removal to peer replicas. DeleteJob has no ctx, so use Background.
if err := s.jobs.Delete(context.Background(), jobID); err != nil {
xlog.Warn("Failed to delete job from store", "job_id", jobID, "error", err)
}
delete(s.jobs, jobID)
s.mu.Unlock()
// Remove job directory (state.json, output files)
@@ -380,7 +324,7 @@ func (s *QuantizationService) DeleteJob(userID, jobID string) error {
// StreamProgress opens a gRPC progress stream and calls the callback for each update.
func (s *QuantizationService) StreamProgress(ctx context.Context, userID, jobID string, callback func(event *schema.QuantizationProgressEvent)) error {
s.mu.Lock()
job, ok := s.jobs.Get(jobID)
job, ok := s.jobs[jobID]
if !ok {
s.mu.Unlock()
return fmt.Errorf("job not found: %s", jobID)
@@ -409,7 +353,7 @@ func (s *QuantizationService) StreamProgress(ctx context.Context, userID, jobID
}, func(update *pb.QuantizationProgressUpdate) {
// Update job status and persist
s.mu.Lock()
if j, ok := s.jobs.Get(jobID); ok {
if j, ok := s.jobs[jobID]; ok {
// Don't let progress updates overwrite terminal states
isTerminal := j.Status == "stopped" || j.Status == "completed" || j.Status == "failed"
if !isTerminal {
@@ -421,9 +365,6 @@ func (s *QuantizationService) StreamProgress(ctx context.Context, userID, jobID
if update.OutputFile != "" {
j.OutputFile = update.OutputFile
}
if err := s.jobs.Set(ctx, j); err != nil {
xlog.Warn("Failed to persist progress update", "job_id", jobID, "error", err)
}
s.saveJobState(j)
}
s.mu.Unlock()
@@ -458,7 +399,7 @@ func sanitizeQuantModelName(s string) string {
// ImportModel imports a quantized model into LocalAI asynchronously.
func (s *QuantizationService) ImportModel(ctx context.Context, userID, jobID string, req schema.QuantizationImportRequest) (string, error) {
s.mu.Lock()
job, ok := s.jobs.Get(jobID)
job, ok := s.jobs[jobID]
if !ok {
s.mu.Unlock()
return "", fmt.Errorf("job not found: %s", jobID)
@@ -518,9 +459,6 @@ func (s *QuantizationService) ImportModel(ctx context.Context, userID, jobID str
job.ImportStatus = "importing"
job.ImportMessage = ""
job.ImportModelName = ""
if err := s.jobs.Set(ctx, job); err != nil {
xlog.Warn("Failed to persist import start", "job_id", jobID, "error", err)
}
s.saveJobState(job)
s.mu.Unlock()
@@ -576,15 +514,10 @@ func (s *QuantizationService) ImportModel(ctx context.Context, userID, jobID str
xlog.Info("Quantized model imported and registered", "job_id", jobID, "model_name", modelName)
// Runs after the HTTP request returns, so use Background rather than the
// (now likely cancelled) request ctx for the write-through.
s.mu.Lock()
job.ImportStatus = "completed"
job.ImportModelName = modelName
job.ImportMessage = ""
if err := s.jobs.Set(context.Background(), job); err != nil {
xlog.Warn("Failed to persist import completion", "job_id", jobID, "error", err)
}
s.saveJobState(job)
s.mu.Unlock()
}()
@@ -592,14 +525,10 @@ func (s *QuantizationService) ImportModel(ctx context.Context, userID, jobID str
return modelName, nil
}
// setImportMessage updates the import message and persists the job state. Called
// from the background import goroutine, so it uses Background for write-through.
// setImportMessage updates the import message and persists the job state.
func (s *QuantizationService) setImportMessage(job *schema.QuantizationJob, msg string) {
s.mu.Lock()
job.ImportMessage = msg
if err := s.jobs.Set(context.Background(), job); err != nil {
xlog.Warn("Failed to persist import message", "job_id", job.ID, "error", err)
}
s.saveJobState(job)
s.mu.Unlock()
}
@@ -610,9 +539,6 @@ func (s *QuantizationService) setImportFailed(job *schema.QuantizationJob, messa
s.mu.Lock()
job.ImportStatus = "failed"
job.ImportMessage = message
if err := s.jobs.Set(context.Background(), job); err != nil {
xlog.Warn("Failed to persist import failure", "job_id", job.ID, "error", err)
}
s.saveJobState(job)
s.mu.Unlock()
}
@@ -620,7 +546,7 @@ func (s *QuantizationService) setImportFailed(job *schema.QuantizationJob, messa
// GetOutputPath returns the path to the quantized model file and a download name.
func (s *QuantizationService) GetOutputPath(userID, jobID string) (string, string, error) {
s.mu.Lock()
job, ok := s.jobs.Get(jobID)
job, ok := s.jobs[jobID]
if !ok {
s.mu.Unlock()
return "", "", fmt.Errorf("job not found: %s", jobID)

View File

@@ -1,187 +0,0 @@
package quantization
// White-box tests (package quantization) so a spec can drive the service's
// internal SyncedMap the same way StartJob does (via jobs.Set) without standing
// up a quantization backend, then assert the cross-replica reads
// (GetJob/ListJobs) and the adapter conversions that keep REST responses
// byte-for-byte unchanged.
import (
"context"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"github.com/mudler/LocalAI/core/config"
"github.com/mudler/LocalAI/core/schema"
"github.com/mudler/LocalAI/core/services/distributed"
"github.com/mudler/LocalAI/core/services/testutil"
)
// newTestService builds a standalone QuantizationService wired to the given bus.
// The model/config loaders are nil because the read/sync paths under test never
// touch them; the data dir is a throwaway temp dir so the disk Loader finds
// nothing.
func newTestService(bus *testutil.FakeBus) *QuantizationService {
appConfig := &config.ApplicationConfig{
Context: context.Background(),
DataPath: GinkgoT().TempDir(),
}
return NewQuantizationService(appConfig, nil, nil, bus, nil)
}
var _ = Describe("QuantizationService", func() {
ctx := context.Background()
Describe("cross-replica job visibility", func() {
var (
bus *testutil.FakeBus
a, b *QuantizationService
)
BeforeEach(func() {
// One shared bus, two replicas: exactly the distributed topology where a
// round-robin request may land on a replica that did not originate the
// change.
bus = testutil.NewFakeBus()
a = newTestService(bus)
b = newTestService(bus)
})
AfterEach(func() {
Expect(a.Close()).To(Succeed())
Expect(b.Close()).To(Succeed())
})
It("makes a job created on A visible via B's GetJob and ListJobs", func() {
job := &schema.QuantizationJob{ID: "job-1", UserID: "user-1", Status: "queued", CreatedAt: "2026-06-27T10:00:00Z"}
// StartJob persists via jobs.Set; drive that directly to avoid a backend.
Expect(a.jobs.Set(ctx, job)).To(Succeed())
got, err := b.GetJob("user-1", "job-1")
Expect(err).ToNot(HaveOccurred(), "B must see a job A just created")
Expect(got.Status).To(Equal("queued"))
listed := b.ListJobs("user-1")
Expect(listed).To(HaveLen(1))
Expect(listed[0].ID).To(Equal("job-1"))
})
It("removes a job from B when it is deleted on A", func() {
job := &schema.QuantizationJob{ID: "job-2", UserID: "user-1", Status: "completed", CreatedAt: "2026-06-27T10:00:00Z"}
Expect(a.jobs.Set(ctx, job)).To(Succeed())
_, err := b.GetJob("user-1", "job-2")
Expect(err).ToNot(HaveOccurred(), "precondition: B must have the job before the delete")
Expect(a.jobs.Delete(ctx, "job-2")).To(Succeed())
_, err = b.GetJob("user-1", "job-2")
Expect(err).To(HaveOccurred(), "a delete on A must remove the job from B")
})
It("propagates a status update from A to B", func() {
job := &schema.QuantizationJob{ID: "job-3", UserID: "user-1", Status: "quantizing", CreatedAt: "2026-06-27T10:00:00Z"}
Expect(a.jobs.Set(ctx, job)).To(Succeed())
updated := &schema.QuantizationJob{ID: "job-3", UserID: "user-1", Status: "completed", CreatedAt: "2026-06-27T10:00:00Z"}
Expect(a.jobs.Set(ctx, updated)).To(Succeed())
got, err := b.GetJob("user-1", "job-3")
Expect(err).ToNot(HaveOccurred())
Expect(got.Status).To(Equal("completed"))
})
})
Describe("ListJobs", func() {
var svc *QuantizationService
BeforeEach(func() {
svc = newTestService(testutil.NewFakeBus())
})
AfterEach(func() { Expect(svc.Close()).To(Succeed()) })
It("filters by user and sorts newest-first", func() {
Expect(svc.jobs.Set(ctx, &schema.QuantizationJob{ID: "old", UserID: "u1", CreatedAt: "2026-06-25T10:00:00Z"})).To(Succeed())
Expect(svc.jobs.Set(ctx, &schema.QuantizationJob{ID: "new", UserID: "u1", CreatedAt: "2026-06-27T10:00:00Z"})).To(Succeed())
Expect(svc.jobs.Set(ctx, &schema.QuantizationJob{ID: "other", UserID: "u2", CreatedAt: "2026-06-26T10:00:00Z"})).To(Succeed())
jobs := svc.ListJobs("u1")
Expect(jobs).To(HaveLen(2), "only u1's jobs")
Expect(jobs[0].ID).To(Equal("new"), "newest first")
Expect(jobs[1].ID).To(Equal("old"))
})
It("returns every user's jobs when the userID filter is empty", func() {
Expect(svc.jobs.Set(ctx, &schema.QuantizationJob{ID: "a", UserID: "u1", CreatedAt: "2026-06-25T10:00:00Z"})).To(Succeed())
Expect(svc.jobs.Set(ctx, &schema.QuantizationJob{ID: "b", UserID: "u2", CreatedAt: "2026-06-26T10:00:00Z"})).To(Succeed())
Expect(svc.ListJobs("")).To(HaveLen(2))
})
It("rejects GetJob for a job owned by another user", func() {
Expect(svc.jobs.Set(ctx, &schema.QuantizationJob{ID: "x", UserID: "owner", CreatedAt: "2026-06-25T10:00:00Z"})).To(Succeed())
_, err := svc.GetJob("intruder", "x")
Expect(err).To(HaveOccurred(), "a different user must not read someone else's job")
})
})
Describe("store adapter conversion", func() {
// The SyncedMap value type is *schema.QuantizationJob (the exact REST shape).
// These specs prove the DB adapter round-trips it losslessly, so hydrate and
// write-through in distributed mode keep responses unchanged.
It("round-trips a job through jobToRecord/recordToJob preserving the API shape", func() {
original := &schema.QuantizationJob{
ID: "rt-1",
UserID: "user-1",
Model: "base-model",
Backend: "llama-cpp-quantization",
ModelID: "llama-cpp-quantization-quantize-rt-1",
QuantizationType: "q4_k_m",
Status: "completed",
Message: "done",
OutputDir: "/data/quantization/rt-1",
OutputFile: "/data/quantization/rt-1/model.gguf",
ExtraOptions: map[string]string{"hf_token": "secret"},
CreatedAt: "2026-06-27T10:00:00Z",
ImportStatus: "completed",
ImportMessage: "",
ImportModelName: "base-model-q4_k_m-rt-1",
Config: &schema.QuantizationJobRequest{Model: "base-model", Backend: "llama-cpp-quantization", QuantizationType: "q4_k_m"},
}
rec := jobToRecord(original)
Expect(rec.ID).To(Equal("rt-1"))
Expect(rec.ConfigJSON).ToNot(BeEmpty(), "structured config must serialize into the JSON column")
Expect(rec.ExtraOptsJSON).ToNot(BeEmpty())
back := recordToJob(rec)
Expect(back.ID).To(Equal(original.ID))
Expect(back.UserID).To(Equal(original.UserID))
Expect(back.Model).To(Equal(original.Model))
Expect(back.Backend).To(Equal(original.Backend))
Expect(back.ModelID).To(Equal(original.ModelID))
Expect(back.QuantizationType).To(Equal(original.QuantizationType))
Expect(back.Status).To(Equal(original.Status))
Expect(back.Message).To(Equal(original.Message))
Expect(back.OutputDir).To(Equal(original.OutputDir))
Expect(back.OutputFile).To(Equal(original.OutputFile))
Expect(back.ImportStatus).To(Equal(original.ImportStatus))
Expect(back.ImportModelName).To(Equal(original.ImportModelName))
Expect(back.CreatedAt).To(Equal(original.CreatedAt))
Expect(back.ExtraOptions).To(Equal(original.ExtraOptions))
Expect(back.Config).ToNot(BeNil())
Expect(back.Config.QuantizationType).To(Equal("q4_k_m"))
})
})
Describe("compile-time adapter contract", func() {
It("satisfies syncstate.Store for *distributed.QuantStore", func() {
// Guards against drift between the adapter and the component interface;
// the var assertion in syncstore.go covers it at build time, this keeps
// the type referenced from a spec too.
var _ *distributed.QuantStore
Expect(&quantStoreAdapter{}).ToNot(BeNil())
})
})
})

View File

@@ -1,114 +0,0 @@
package quantization
import (
"context"
"encoding/json"
"time"
"github.com/mudler/LocalAI/core/schema"
"github.com/mudler/LocalAI/core/services/distributed"
"github.com/mudler/LocalAI/core/services/syncstate"
)
// quantStoreAdapter bridges the distributed PostgreSQL QuantStore to the generic
// syncstate.Store the SyncedMap consumes. It is only wired in distributed mode;
// standalone leaves Store nil and hydrates from disk via a Loader instead.
//
// The SyncedMap value type is *schema.QuantizationJob (the exact shape the REST
// API returns) so reads need no conversion and the response JSON is provably
// unchanged. The adapter is the single place that translates between that API
// shape and the DB QuantJobRecord.
type quantStoreAdapter struct {
store *distributed.QuantStore
}
// compile-time assertion that the adapter satisfies the component's Store.
var _ syncstate.Store[string, *schema.QuantizationJob] = (*quantStoreAdapter)(nil)
func (a *quantStoreAdapter) List(_ context.Context) ([]*schema.QuantizationJob, error) {
records, err := a.store.ListAll()
if err != nil {
return nil, err
}
jobs := make([]*schema.QuantizationJob, 0, len(records))
for i := range records {
jobs = append(jobs, recordToJob(&records[i]))
}
return jobs, nil
}
func (a *quantStoreAdapter) Upsert(_ context.Context, job *schema.QuantizationJob) error {
return a.store.Upsert(jobToRecord(job))
}
func (a *quantStoreAdapter) Delete(_ context.Context, id string) error {
return a.store.Delete(id)
}
// recordToJob maps a persisted DB record back to the API shape, reconstructing
// the structured Config / ExtraOptions from their JSON columns.
func recordToJob(r *distributed.QuantJobRecord) *schema.QuantizationJob {
job := &schema.QuantizationJob{
ID: r.ID,
UserID: r.UserID,
Model: r.Model,
Backend: r.Backend,
ModelID: r.ModelID,
QuantizationType: r.QuantizationType,
Status: r.Status,
Message: r.Message,
OutputDir: r.OutputDir,
OutputFile: r.OutputFile,
ImportStatus: r.ImportStatus,
ImportMessage: r.ImportMessage,
ImportModelName: r.ImportModelName,
CreatedAt: r.CreatedAt.UTC().Format(time.RFC3339),
}
if r.ExtraOptsJSON != "" {
// Best-effort: a malformed column must not drop the whole job from the API.
_ = json.Unmarshal([]byte(r.ExtraOptsJSON), &job.ExtraOptions)
}
if r.ConfigJSON != "" {
var cfg schema.QuantizationJobRequest
if err := json.Unmarshal([]byte(r.ConfigJSON), &cfg); err == nil {
job.Config = &cfg
}
}
return job
}
// jobToRecord maps the API shape to a DB record for write-through, serializing
// the structured Config / ExtraOptions into their JSON columns. CreatedAt is
// parsed back from the RFC3339 string the service stamps; an unparseable value is
// left zero so QuantStore.Upsert stamps "now".
func jobToRecord(job *schema.QuantizationJob) *distributed.QuantJobRecord {
rec := &distributed.QuantJobRecord{
ID: job.ID,
UserID: job.UserID,
Model: job.Model,
Backend: job.Backend,
ModelID: job.ModelID,
QuantizationType: job.QuantizationType,
Status: job.Status,
Message: job.Message,
OutputDir: job.OutputDir,
OutputFile: job.OutputFile,
ImportStatus: job.ImportStatus,
ImportMessage: job.ImportMessage,
ImportModelName: job.ImportModelName,
}
if job.Config != nil {
if data, err := json.Marshal(job.Config); err == nil {
rec.ConfigJSON = string(data)
}
}
if job.ExtraOptions != nil {
if data, err := json.Marshal(job.ExtraOptions); err == nil {
rec.ExtraOptsJSON = string(data)
}
}
if t, err := time.Parse(time.RFC3339, job.CreatedAt); err == nil {
rec.CreatedAt = t
}
return rec
}

View File

@@ -1,289 +0,0 @@
// Package syncstate provides SyncedMap, a reusable cross-replica in-memory map.
//
// LocalAI in distributed mode runs multiple frontend replicas behind a
// round-robin load balancer. Several features keep process-local in-memory state
// that is surfaced to the HTTP/UI API; without cross-replica sync a poll that
// lands on a replica which did not originate a change sees stale or missing data.
// SyncedMap collapses the three legs each feature otherwise hand-wires - an
// in-memory map, a NATS broadcast/apply path, and optional durable read-through -
// into one well-tested component so cross-replica consistency is a configuration
// choice rather than a bespoke re-implementation.
package syncstate
import (
"context"
"sync"
"time"
"github.com/mudler/LocalAI/core/services/messaging"
"github.com/mudler/xlog"
)
// Op values carried on the wire and passed to OnApply.
const (
opSet = "set"
opDelete = "delete"
)
// Store is optional durable backing for a SyncedMap. In distributed mode it is a
// single shared DB, so the apply path (a delta received from a peer) updates
// memory only and never re-writes the Store.
type Store[K comparable, V any] interface {
List(ctx context.Context) ([]V, error)
Upsert(ctx context.Context, v V) error
Delete(ctx context.Context, k K) error
}
// Config configures a SyncedMap.
type Config[K comparable, V any] struct {
Name string // subject namespace, e.g. "finetune.jobs"
Key func(V) K // extract the key from a value
Nats messaging.MessagingClient // nil => standalone: in-memory only, no broadcast/subscribe
Store Store[K, V] // optional read-through persistence
Loader func(ctx context.Context) ([]V, error) // source when there is no Store (e.g. disk reload)
OnApply func(op string, k K, v V) // optional hook after an applied change (e.g. ShutdownModel)
Reconcile time.Duration // optional periodic re-hydrate; 0 = off
}
// delta is the JSON wire envelope broadcast on every local mutation. Value is
// omitempty so a delete carries only op+key.
type delta[K comparable, V any] struct {
Op string `json:"op"`
Key K `json:"key"`
Value V `json:"value,omitempty"`
}
// SyncedMap is a cross-replica in-memory map. A local write (Set/Delete) updates
// memory, the optional durable Store, then broadcasts a delta to peers. A peer's
// delta updates memory only and fires OnApply - it never re-broadcasts and never
// writes the Store. That structural split is the echo-loop guard (same pattern as
// galleryop.mergeStatus / OpCache.applyStart): receiving your own broadcast just
// re-applies an idempotent value to memory, so there is no storm and no
// double-write.
type SyncedMap[K comparable, V any] struct {
cfg Config[K, V]
mu sync.RWMutex
data map[K]V
sub Subscription
// lifeCtx outlives Start's argument: a reconnect callback or reconcile tick
// can fire long after Start returns, so they must not be tied to a ctx the
// caller may cancel. Close cancels it.
lifeCtx context.Context
cancel context.CancelFunc
wg sync.WaitGroup
}
// Subscription is the subset of messaging.Subscription the component holds onto.
type Subscription = messaging.Subscription
// New constructs a SyncedMap. Call Start to hydrate and begin syncing.
func New[K comparable, V any](cfg Config[K, V]) *SyncedMap[K, V] {
return &SyncedMap[K, V]{cfg: cfg, data: make(map[K]V)}
}
func (m *SyncedMap[K, V]) subject() string {
return messaging.SubjectSyncStateDelta(m.cfg.Name)
}
// Start hydrates from the source, subscribes for peer deltas, registers a
// reconnect re-hydrate (when the client supports it), and starts the optional
// reconcile ticker.
func (m *SyncedMap[K, V]) Start(ctx context.Context) error {
if err := m.hydrate(ctx); err != nil {
return err
}
// The cancel func is stored on the struct and invoked in Close (covered by
// tests); lifeCtx must outlive Start to drive the reconnect/reconcile
// goroutines, so it cannot be cancelled or deferred within this scope.
m.lifeCtx, m.cancel = context.WithCancel(context.Background()) // #nosec G118 -- cancel is invoked in Close()
if m.cfg.Nats != nil {
sub, err := messaging.SubscribeJSON(m.cfg.Nats, m.subject(), m.apply)
if err != nil {
return err
}
m.sub = sub
// nats.go transparently resubscribes on reconnect, but it cannot know we
// kept derived in-memory state that may have drifted while the link was
// down, so re-hydrate from the durable source. Detected via an optional
// interface so MessagingClient itself stays minimal; standalone/test
// clients without the method simply fall back to the reconcile ticker.
if r, ok := m.cfg.Nats.(interface{ OnReconnect(func()) }); ok {
r.OnReconnect(func() {
if err := m.hydrate(m.lifeCtx); err != nil {
xlog.Warn("syncstate: reconnect re-hydrate failed", "name", m.cfg.Name, "error", err)
}
})
}
}
if m.cfg.Reconcile > 0 {
m.wg.Add(1)
go m.reconcileLoop()
}
return nil
}
// Close unsubscribes and stops the reconcile ticker.
func (m *SyncedMap[K, V]) Close() error {
if m.cancel != nil {
m.cancel()
}
m.wg.Wait()
if m.sub != nil {
return m.sub.Unsubscribe()
}
return nil
}
// Set updates the value locally, writes through the Store, then broadcasts.
// Per the data-flow contract the Store write happens under the lock so memory and
// durable state move together; the broadcast is best-effort after unlocking.
func (m *SyncedMap[K, V]) Set(ctx context.Context, v V) error {
k := m.cfg.Key(v)
m.mu.Lock()
m.data[k] = v
if m.cfg.Store != nil {
if err := m.cfg.Store.Upsert(ctx, v); err != nil {
m.mu.Unlock()
return err
}
}
m.mu.Unlock()
m.publish(opSet, k, v)
return nil
}
// Delete removes the key locally, deletes it from the Store, then broadcasts.
func (m *SyncedMap[K, V]) Delete(ctx context.Context, k K) error {
m.mu.Lock()
delete(m.data, k)
if m.cfg.Store != nil {
if err := m.cfg.Store.Delete(ctx, k); err != nil {
m.mu.Unlock()
return err
}
}
m.mu.Unlock()
var zero V
m.publish(opDelete, k, zero)
return nil
}
// Get returns the value for k and whether it was present.
func (m *SyncedMap[K, V]) Get(k K) (V, bool) {
m.mu.RLock()
defer m.mu.RUnlock()
v, ok := m.data[k]
return v, ok
}
// List returns a snapshot slice of all values.
func (m *SyncedMap[K, V]) List() []V {
m.mu.RLock()
defer m.mu.RUnlock()
out := make([]V, 0, len(m.data))
for _, v := range m.data {
out = append(out, v)
}
return out
}
// Snapshot returns a copy of the underlying map.
func (m *SyncedMap[K, V]) Snapshot() map[K]V {
m.mu.RLock()
defer m.mu.RUnlock()
out := make(map[K]V, len(m.data))
for k, v := range m.data {
out[k] = v
}
return out
}
// publish broadcasts a delta. Standalone (nil Nats) is a strict no-op.
func (m *SyncedMap[K, V]) publish(op string, k K, v V) {
if m.cfg.Nats == nil {
return
}
if err := m.cfg.Nats.Publish(m.subject(), delta[K, V]{Op: op, Key: k, Value: v}); err != nil {
xlog.Warn("syncstate: failed to broadcast delta", "name", m.cfg.Name, "op", op, "error", err)
}
}
// apply handles a peer's delta: memory-only update plus OnApply. It deliberately
// never writes the Store nor re-publishes - that is the echo-loop guard.
func (m *SyncedMap[K, V]) apply(d delta[K, V]) {
switch d.Op {
case opSet:
m.mu.Lock()
m.data[d.Key] = d.Value
m.mu.Unlock()
case opDelete:
m.mu.Lock()
delete(m.data, d.Key)
m.mu.Unlock()
default:
xlog.Warn("syncstate: ignoring delta with unknown op", "name", m.cfg.Name, "op", d.Op)
return
}
if m.cfg.OnApply != nil {
m.cfg.OnApply(d.Op, d.Key, d.Value)
}
}
// hydrate replaces the whole map from the durable source: Store if present, else
// Loader. With neither, a late joiner starts empty and catches up via deltas
// (acceptable only for ephemeral state).
func (m *SyncedMap[K, V]) hydrate(ctx context.Context) error {
var (
vals []V
err error
)
switch {
case m.cfg.Store != nil:
vals, err = m.cfg.Store.List(ctx)
case m.cfg.Loader != nil:
vals, err = m.cfg.Loader(ctx)
default:
return nil
}
if err != nil {
return err
}
m.replaceAll(vals)
return nil
}
// replaceAll atomically swaps the map contents for the given values, keyed via
// cfg.Key.
func (m *SyncedMap[K, V]) replaceAll(vals []V) {
next := make(map[K]V, len(vals))
for _, v := range vals {
next[m.cfg.Key(v)] = v
}
m.mu.Lock()
m.data = next
m.mu.Unlock()
}
// reconcileLoop periodically re-hydrates to repair silent drift (missed deltas).
func (m *SyncedMap[K, V]) reconcileLoop() {
defer m.wg.Done()
t := time.NewTicker(m.cfg.Reconcile)
defer t.Stop()
for {
select {
case <-m.lifeCtx.Done():
return
case <-t.C:
if err := m.hydrate(m.lifeCtx); err != nil {
xlog.Warn("syncstate: reconcile re-hydrate failed", "name", m.cfg.Name, "error", err)
}
}
}
}

View File

@@ -1,13 +0,0 @@
package syncstate_test
import (
"testing"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
)
func TestSyncstate(t *testing.T) {
RegisterFailHandler(Fail)
RunSpecs(t, "Syncstate Suite")
}

View File

@@ -1,291 +0,0 @@
package syncstate_test
import (
"context"
"sync"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"github.com/mudler/LocalAI/core/services/messaging"
"github.com/mudler/LocalAI/core/services/syncstate"
"github.com/mudler/LocalAI/core/services/testutil"
)
// job is a minimal JSON-serializable value stand-in for the real cross-replica
// records (finetune/quant/agent jobs) the component is built for.
type job struct {
ID string `json:"id"`
Status string `json:"status"`
}
func jobKey(j *job) string { return j.ID }
const stateName = "test.jobs"
func deltaSubject() string { return messaging.SubjectSyncStateDelta(stateName) }
// fakeStore is an in-memory Store that records call counts so specs can assert
// the write-through-vs-apply split (local writes hit the Store; applied deltas
// must not).
type fakeStore struct {
mu sync.Mutex
data map[string]*job
upsertCalls int
deleteCalls int
listCalls int
}
func newFakeStore(seed ...*job) *fakeStore {
s := &fakeStore{data: map[string]*job{}}
for _, j := range seed {
s.data[j.ID] = j
}
return s
}
func (s *fakeStore) List(_ context.Context) ([]*job, error) {
s.mu.Lock()
defer s.mu.Unlock()
s.listCalls++
out := make([]*job, 0, len(s.data))
for _, j := range s.data {
out = append(out, j)
}
return out, nil
}
func (s *fakeStore) Upsert(_ context.Context, j *job) error {
s.mu.Lock()
defer s.mu.Unlock()
s.upsertCalls++
s.data[j.ID] = j
return nil
}
func (s *fakeStore) Delete(_ context.Context, k string) error {
s.mu.Lock()
defer s.mu.Unlock()
s.deleteCalls++
delete(s.data, k)
return nil
}
// add simulates a peer replica writing to the shared DB out-of-band (e.g. while
// this replica was partitioned), so a re-hydrate can be observed to pick it up.
func (s *fakeStore) add(j *job) {
s.mu.Lock()
defer s.mu.Unlock()
s.data[j.ID] = j
}
func (s *fakeStore) counts() (upsert, del, list int) {
s.mu.Lock()
defer s.mu.Unlock()
return s.upsertCalls, s.deleteCalls, s.listCalls
}
var _ = Describe("SyncedMap", func() {
ctx := context.Background()
Describe("cross-replica delta propagation", func() {
var (
bus *testutil.FakeBus
a, b *syncstate.SyncedMap[string, *job]
)
BeforeEach(func() {
bus = testutil.NewFakeBus()
a = syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey, Nats: bus})
b = syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey, Nats: bus})
Expect(a.Start(ctx)).To(Succeed())
Expect(b.Start(ctx)).To(Succeed())
})
AfterEach(func() {
Expect(a.Close()).To(Succeed())
Expect(b.Close()).To(Succeed())
})
It("propagates a Set on A to B", func() {
Expect(a.Set(ctx, &job{ID: "1", Status: "running"})).To(Succeed())
got, ok := b.Get("1")
Expect(ok).To(BeTrue(), "replica B should see the value A just set")
Expect(got.Status).To(Equal("running"))
})
It("prunes a Delete on A from B", func() {
Expect(a.Set(ctx, &job{ID: "1", Status: "running"})).To(Succeed())
_, present := b.Get("1")
Expect(present).To(BeTrue(), "precondition: B must have the value before the delete")
Expect(a.Delete(ctx, "1")).To(Succeed())
_, ok := b.Get("1")
Expect(ok).To(BeFalse(), "a delete on A must remove the key from B")
})
})
Describe("hydration", func() {
It("hydrates on Start from a preloaded Store", func() {
store := newFakeStore(&job{ID: "x", Status: "done"})
m := syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey, Store: store})
Expect(m.Start(ctx)).To(Succeed())
got, ok := m.Get("x")
Expect(ok).To(BeTrue(), "Start must populate the map from the Store")
Expect(got.Status).To(Equal("done"))
})
It("uses the Loader when Store is nil", func() {
m := syncstate.New(syncstate.Config[string, *job]{
Name: stateName,
Key: jobKey,
Loader: func(_ context.Context) ([]*job, error) {
return []*job{{ID: "l", Status: "loaded"}}, nil
},
})
Expect(m.Start(ctx)).To(Succeed())
got, ok := m.Get("l")
Expect(ok).To(BeTrue(), "Loader output must hydrate the map when there is no Store")
Expect(got.Status).To(Equal("loaded"))
})
})
Describe("echo-loop guard", func() {
It("applies its own broadcast once and does not re-publish", func() {
bus := testutil.NewFakeBus()
a := syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey, Nats: bus})
b := syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey, Nats: bus})
Expect(a.Start(ctx)).To(Succeed())
Expect(b.Start(ctx)).To(Succeed())
defer func() {
Expect(a.Close()).To(Succeed())
Expect(b.Close()).To(Succeed())
}()
Expect(a.Set(ctx, &job{ID: "e", Status: "running"})).To(Succeed())
// One local write must produce exactly one broadcast: A and B both
// receive it and apply to memory, but the apply path never re-publishes.
Expect(bus.PublishCount(deltaSubject())).To(Equal(1),
"the apply path must not re-broadcast, otherwise replicas storm")
Expect(a.List()).To(HaveLen(1), "A must not double-store its own echo")
_, ok := b.Get("e")
Expect(ok).To(BeTrue())
})
})
Describe("Store write-through vs apply", func() {
It("writes the Store on local Set/Delete but not on an applied delta", func() {
bus := testutil.NewFakeBus()
storeA := newFakeStore()
storeB := newFakeStore()
a := syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey, Nats: bus, Store: storeA})
b := syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey, Nats: bus, Store: storeB})
Expect(a.Start(ctx)).To(Succeed())
Expect(b.Start(ctx)).To(Succeed())
defer func() {
Expect(a.Close()).To(Succeed())
Expect(b.Close()).To(Succeed())
}()
Expect(a.Set(ctx, &job{ID: "w", Status: "running"})).To(Succeed())
upA, _, _ := storeA.counts()
upB, _, _ := storeB.counts()
Expect(upA).To(Equal(1), "local Set must write through to its own Store")
Expect(upB).To(Equal(0), "the apply path must never write the peer's Store")
Expect(a.Delete(ctx, "w")).To(Succeed())
_, delA, _ := storeA.counts()
_, delB, _ := storeB.counts()
Expect(delA).To(Equal(1), "local Delete must delete from its own Store")
Expect(delB).To(Equal(0), "the apply path must never delete from the peer's Store")
})
})
Describe("OnApply hook", func() {
It("fires with the correct op and key on an applied delta", func() {
bus := testutil.NewFakeBus()
var (
mu sync.Mutex
ops []string
keys []string
)
a := syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey, Nats: bus})
b := syncstate.New(syncstate.Config[string, *job]{
Name: stateName, Key: jobKey, Nats: bus,
OnApply: func(op string, k string, _ *job) {
mu.Lock()
ops = append(ops, op)
keys = append(keys, k)
mu.Unlock()
},
})
Expect(a.Start(ctx)).To(Succeed())
Expect(b.Start(ctx)).To(Succeed())
defer func() {
Expect(a.Close()).To(Succeed())
Expect(b.Close()).To(Succeed())
}()
Expect(a.Set(ctx, &job{ID: "o", Status: "running"})).To(Succeed())
Expect(a.Delete(ctx, "o")).To(Succeed())
mu.Lock()
defer mu.Unlock()
Expect(ops).To(Equal([]string{"set", "delete"}))
Expect(keys).To(Equal([]string{"o", "o"}))
})
})
Describe("standalone (nil Nats)", func() {
It("works in-memory with no panic and nothing to broadcast", func() {
m := syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey})
Expect(m.Start(ctx)).To(Succeed())
defer func() { Expect(m.Close()).To(Succeed()) }()
Expect(func() {
Expect(m.Set(ctx, &job{ID: "s", Status: "running"})).To(Succeed())
}).ToNot(Panic())
got, ok := m.Get("s")
Expect(ok).To(BeTrue())
Expect(got.Status).To(Equal("running"))
Expect(m.List()).To(HaveLen(1))
Expect(m.Snapshot()).To(HaveKey("s"))
Expect(m.Delete(ctx, "s")).To(Succeed())
_, ok = m.Get("s")
Expect(ok).To(BeFalse())
})
})
Describe("reconnect re-hydrate", func() {
It("re-reads the source when the messaging client reconnects", func() {
bus := testutil.NewFakeBus()
store := newFakeStore(&job{ID: "init", Status: "running"})
m := syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey, Nats: bus, Store: store})
Expect(m.Start(ctx)).To(Succeed())
defer func() { Expect(m.Close()).To(Succeed()) }()
_, ok := m.Get("init")
Expect(ok).To(BeTrue())
// A peer writes to the shared DB while we are unaware (no delta seen).
store.add(&job{ID: "late", Status: "running"})
_, ok = m.Get("late")
Expect(ok).To(BeFalse(), "the new row should not appear before a re-hydrate")
bus.TriggerReconnect()
_, ok = m.Get("late")
Expect(ok).To(BeTrue(), "reconnect must re-hydrate from the source and pick up drift")
_, _, list := store.counts()
Expect(list).To(Equal(2), "exactly one Start hydrate plus one reconnect re-hydrate")
})
})
})

View File

@@ -1,160 +0,0 @@
package testutil
import (
"encoding/json"
"strings"
"sync"
"time"
"github.com/mudler/LocalAI/core/services/messaging"
)
// FakeBus is an in-memory messaging.MessagingClient that delivers each published
// message synchronously to every registered subscriber whose subject filter
// matches, including NATS-style wildcard subjects (`*` matches exactly one
// token).
//
// Synchronous delivery keeps specs deterministic: the moment Publish returns,
// every matching subscriber's handler has already run, so the spec body can read
// the resulting state without polling. It is the shared test double for every
// cross-replica-sync adopter (gallery, syncstate, ...) so they exercise the same
// delivery semantics. It deliberately depends only on the standard library and
// the messaging package — no test framework — so it is importable anywhere.
type FakeBus struct {
mu sync.Mutex
subs []fakeBusSub
// publishCounts records how many messages were published per subject, so a
// spec can assert the echo-loop guard (an applied delta must not re-publish).
publishCounts map[string]int
// reconnectCbs back the optional OnReconnect/TriggerReconnect pair, letting a
// spec exercise the component's reconnect re-hydrate path without a real
// NATS server.
reconnectCbs []func()
}
type fakeBusSub struct {
subject string
handler func([]byte)
}
// NewFakeBus returns a ready-to-use in-memory bus.
func NewFakeBus() *FakeBus {
return &FakeBus{publishCounts: map[string]int{}}
}
// subjectMatches reports whether a subscription filter matches a concrete
// subject, honoring the single-token `*` wildcard used by NATS.
func subjectMatches(filter, subject string) bool {
if filter == subject {
return true
}
fp := strings.Split(filter, ".")
sp := strings.Split(subject, ".")
if len(fp) != len(sp) {
return false
}
for i := range fp {
if fp[i] == "*" {
continue
}
if fp[i] != sp[i] {
return false
}
}
return true
}
// Publish marshals data as JSON and delivers it synchronously to every matching
// subscriber.
func (b *FakeBus) Publish(subject string, data any) error {
payload, err := json.Marshal(data)
if err != nil {
return err
}
b.mu.Lock()
b.publishCounts[subject]++
subs := append([]fakeBusSub(nil), b.subs...)
b.mu.Unlock()
for _, s := range subs {
if subjectMatches(s.subject, subject) {
s.handler(payload)
}
}
return nil
}
// PublishCount returns how many messages were published on the exact subject.
func (b *FakeBus) PublishCount(subject string) int {
b.mu.Lock()
defer b.mu.Unlock()
return b.publishCounts[subject]
}
type fakeBusSubscription struct {
bus *FakeBus
subRef fakeBusSub
}
func (s *fakeBusSubscription) Unsubscribe() error {
s.bus.mu.Lock()
defer s.bus.mu.Unlock()
for i, candidate := range s.bus.subs {
if candidate.subject == s.subRef.subject {
s.bus.subs = append(s.bus.subs[:i], s.bus.subs[i+1:]...)
return nil
}
}
return nil
}
func (b *FakeBus) Subscribe(subject string, handler func([]byte)) (messaging.Subscription, error) {
sub := fakeBusSub{subject: subject, handler: handler}
b.mu.Lock()
b.subs = append(b.subs, sub)
b.mu.Unlock()
return &fakeBusSubscription{bus: b, subRef: sub}, nil
}
func (b *FakeBus) QueueSubscribe(subject, _ string, handler func([]byte)) (messaging.Subscription, error) {
return b.Subscribe(subject, handler)
}
func (b *FakeBus) QueueSubscribeReply(string, string, func([]byte, func([]byte))) (messaging.Subscription, error) {
return &fakeBusSubscription{bus: b}, nil
}
func (b *FakeBus) SubscribeReply(string, func([]byte, func([]byte))) (messaging.Subscription, error) {
return &fakeBusSubscription{bus: b}, nil
}
func (b *FakeBus) Request(string, []byte, time.Duration) ([]byte, error) {
return nil, nil
}
func (b *FakeBus) IsConnected() bool { return true }
func (b *FakeBus) Close() {}
// OnReconnect mirrors *messaging.Client.OnReconnect so a spec can drive the
// component's reconnect re-hydrate path. The component detects this method via an
// optional interface assertion; implementing it here keeps the fake a faithful
// stand-in for the concrete client.
func (b *FakeBus) OnReconnect(cb func()) {
if cb == nil {
return
}
b.mu.Lock()
b.reconnectCbs = append(b.reconnectCbs, cb)
b.mu.Unlock()
}
// TriggerReconnect runs every registered reconnect callback, simulating a NATS
// reconnect event.
func (b *FakeBus) TriggerReconnect() {
b.mu.Lock()
cbs := append([]func(){}, b.reconnectCbs...)
b.mu.Unlock()
for _, cb := range cbs {
cb()
}
}

View File

@@ -1,3 +1,3 @@
{
"version": "v4.5.5"
"version": "v4.5.2"
}

View File

@@ -579,10 +579,6 @@
icon: https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3.6/Figures/qwen3.6_35b_a3b_score.png
overrides:
backend: llama-cpp
# NVFP4 GGUFs use a quant type the GGUF metadata parser cannot read, so
# context size cannot be auto-derived; set it explicitly (the model trains
# to 262144, 32768 is a safe default operators can raise).
context_size: 32768
function:
automatic_tool_parsing_fallback: true
grammar:
@@ -615,9 +611,6 @@
- gguf
overrides:
backend: llama-cpp
# NVFP4 GGUFs use a quant type the GGUF metadata parser cannot read, so
# context size cannot be auto-derived; set it explicitly.
context_size: 32768
function:
automatic_tool_parsing_fallback: true
grammar:
@@ -645,9 +638,6 @@
icon: https://cdn-uploads.huggingface.co/production/uploads/66309bd090589b7c65950665/sGQKmrMc6L6guMoaB5_Y2.png
overrides:
backend: llama-cpp
# NVFP4 GGUFs use a quant type the GGUF metadata parser cannot read, so
# context size cannot be auto-derived; set it explicitly.
context_size: 32768
function:
automatic_tool_parsing_fallback: true
grammar:
@@ -698,10 +688,6 @@
icon: https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3.6/Figures/qwen3.6_27b_score.png
overrides:
backend: llama-cpp
# NVFP4 GGUFs use a quant type the GGUF metadata parser cannot read, so
# context size cannot be auto-derived; set it explicitly (the model trains
# to 262144, 32768 is a safe default operators can raise).
context_size: 32768
function:
automatic_tool_parsing_fallback: true
grammar:

View File

@@ -19,7 +19,6 @@ func WorkerPermissions(nodeID, nodeType string) (pubAllow, subAllow []string) {
// Keep this list in sync with the subscriptions in core/cli/agent_worker.go.
subAllow = []string{
"agent.execute",
"agent.*.cancel",
"jobs.*.cancel",
"jobs.*.progress",
"jobs.*.result",

View File

@@ -1,161 +0,0 @@
package distributed_test
import (
"context"
"github.com/mudler/LocalAI/core/services/distributed"
"github.com/mudler/LocalAI/core/services/messaging"
"github.com/mudler/LocalAI/core/services/syncstate"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
pgdriver "gorm.io/driver/postgres"
"gorm.io/gorm"
"gorm.io/gorm/logger"
)
// ftSyncStore adapts the real FineTuneStore to syncstate.Store, exactly as the
// finetune service does in production. Defined here (rather than reusing the
// service's unexported adapter) so the e2e exercises the store + component over
// real infrastructure without pulling in backend execution.
type ftSyncStore struct{ s *distributed.FineTuneStore }
func (a ftSyncStore) List(_ context.Context) ([]*distributed.FineTuneJobRecord, error) {
recs, err := a.s.ListAll()
if err != nil {
return nil, err
}
out := make([]*distributed.FineTuneJobRecord, len(recs))
for i := range recs {
r := recs[i]
out[i] = &r
}
return out, nil
}
func (a ftSyncStore) Upsert(_ context.Context, r *distributed.FineTuneJobRecord) error {
return a.s.Upsert(r)
}
func (a ftSyncStore) Delete(_ context.Context, k string) error { return a.s.Delete(k) }
// This suite is the real-infrastructure counterpart to the fake-bus unit tests:
// two SyncedMap instances stand in for two LocalAI frontend replicas, each with
// its OWN NATS connection to a shared NATS server and a SHARED PostgreSQL store -
// the exact distributed-mode invariant (single shared DB, per-replica process
// state). It proves the delta path works over the wire and that a late-joining
// replica recovers via store hydrate (the at-most-once gap a fake bus cannot
// exercise).
var _ = Describe("SyncedMap two-replica sync over real NATS", Label("Distributed"), func() {
var (
infra *TestInfra
ftStore *distributed.FineTuneStore
)
BeforeEach(func() {
infra = SetupInfra("localai_syncstate_dist_test")
db, err := gorm.Open(pgdriver.Open(infra.PGURL), &gorm.Config{
Logger: logger.Default.LogMode(logger.Silent),
})
Expect(err).ToNot(HaveOccurred())
ftStore, err = distributed.NewFineTuneStore(db)
Expect(err).ToNot(HaveOccurred())
})
// newReplica builds an independent "replica": its own NATS client to the
// shared server plus a SyncedMap over the shared store, started (hydrate +
// subscribe) and cleaned up automatically.
newReplica := func() *syncstate.SyncedMap[string, *distributed.FineTuneJobRecord] {
GinkgoHelper()
nc, err := messaging.New(infra.NatsURL)
Expect(err).ToNot(HaveOccurred())
sm := syncstate.New(syncstate.Config[string, *distributed.FineTuneJobRecord]{
Name: "finetune.jobs",
Key: func(r *distributed.FineTuneJobRecord) string { return r.ID },
Nats: nc,
Store: ftSyncStore{s: ftStore},
})
Expect(sm.Start(infra.Ctx)).To(Succeed())
FlushNATS(nc) // ensure the subscription is registered server-side before any publish
DeferCleanup(func() {
_ = sm.Close()
nc.Close()
})
return sm
}
rec := func(id, status string) *distributed.FineTuneJobRecord {
return &distributed.FineTuneJobRecord{
ID: id, UserID: "u1", Model: "m", Backend: "b",
TrainingType: "lora", TrainingMethod: "sft", Status: status,
}
}
It("propagates a create from replica A to replica B over the wire", func() {
a := newReplica()
b := newReplica()
Expect(a.Set(infra.Ctx, rec("job-1", "queued"))).To(Succeed())
Eventually(func() bool { _, ok := b.Get("job-1"); return ok }, "10s", "50ms").
Should(BeTrue(), "replica B must observe the job created on A via NATS")
got, ok := b.Get("job-1")
Expect(ok).To(BeTrue())
Expect(got.Status).To(Equal("queued"))
})
It("propagates an update and a delete across replicas", func() {
a := newReplica()
b := newReplica()
Expect(a.Set(infra.Ctx, rec("job-2", "queued"))).To(Succeed())
Eventually(func() bool { _, ok := b.Get("job-2"); return ok }, "10s", "50ms").Should(BeTrue())
// Update on A -> B reflects the new status.
Expect(a.Set(infra.Ctx, rec("job-2", "training"))).To(Succeed())
Eventually(func() string {
if r, ok := b.Get("job-2"); ok {
return r.Status
}
return ""
}, "10s", "50ms").Should(Equal("training"))
// Delete on A -> B prunes (a reload-from-path could not do this).
Expect(a.Delete(infra.Ctx, "job-2")).To(Succeed())
Eventually(func() bool { _, ok := b.Get("job-2"); return ok }, "10s", "50ms").
Should(BeFalse(), "replica B must drop the job deleted on A")
})
It("hydrates a late-joining replica from the shared store (missed-delta recovery)", func() {
a := newReplica()
// Written (and broadcast) BEFORE replica C exists, so C can never receive
// the delta - it can only learn the job by hydrating from shared Postgres
// on Start. This is the at-most-once gap a fake bus cannot exercise.
Expect(a.Set(infra.Ctx, rec("job-3", "completed"))).To(Succeed())
Eventually(func() (*distributed.FineTuneJobRecord, error) { return ftStore.Get("job-3") }, "10s", "50ms").
ShouldNot(BeNil(), "write-through must reach the shared store first")
c := newReplica() // joins late; Start() hydrates from the store synchronously
got, ok := c.Get("job-3")
Expect(ok).To(BeTrue(), "late replica must recover the job via store hydrate, not a delta")
Expect(got.Status).To(Equal("completed"))
})
It("write-through persists a local Set to the shared PostgreSQL store", func() {
a := newReplica()
Expect(a.Set(infra.Ctx, rec("job-4", "queued"))).To(Succeed())
persisted, err := ftStore.Get("job-4")
Expect(err).ToNot(HaveOccurred())
Expect(persisted.ID).To(Equal("job-4"))
Expect(persisted.Status).To(Equal("queued"))
})
})