Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 28 additions & 0 deletions core/services/messaging/subjects.go
Original file line number Diff line number Diff line change
Expand Up @@ -287,11 +287,39 @@ type BackendStopRequest struct {
Force bool `json:"force,omitempty"`
}

// BackendStopReply is the worker's answer to a backend.stop request.
//
// backend.stop had no reply until this type existed. The controller published
// and returned success as soon as the local publish succeeded, so a stop that
// killed nothing, and a stop that failed outright, both looked identical to a
// stop that worked. An operator calling the unload endpoint got HTTP 200 while
// the backend kept running and holding its VRAM.
type BackendStopReply struct {
Success bool `json:"success"`
Error string `json:"error,omitempty"`

// StoppedProcessKeys names every `modelID#replica` process the worker
// terminated while serving this request.
StoppedProcessKeys []string `json:"stopped_process_keys,omitempty"`

// ReportsStoppedProcesses distinguishes "this worker enumerates what it
// stopped and stopped nothing" from "this worker predates the field", the
// same way BackendDeleteReply does. Both send an empty list and only the
// first is authoritative, so a controller that cannot tell them apart would
// read silence as a completed stop — the exact conclusion this reply exists
// to prevent.
ReportsStoppedProcesses bool `json:"reports_stopped_processes,omitempty"`
}

// SubjectNodeBackendStop tells a worker node to stop its gRPC backend process.
// Equivalent to the local deleteProcess(). The node will:
// 1. Best-effort bounded Free() via gRPC (unless Force is true)
// 2. Kill the backend process
// 3. Can be restarted via another backend.start event.
//
// Request-reply, answered with a BackendStopReply. A worker that predates that
// reply never answers, so the controller must treat a timeout as "unconfirmed"
// rather than "failed" — see RemoteUnloaderAdapter.stopBackend.
func SubjectNodeBackendStop(nodeID string) string {
return subjectNodePrefix + sanitizeSubjectToken(nodeID) + ".backend.stop"
}
Expand Down
58 changes: 55 additions & 3 deletions core/services/nodes/unloader.go
Original file line number Diff line number Diff line change
Expand Up @@ -417,19 +417,71 @@ func (a *RemoteUnloaderAdapter) ListRunningModels(nodeID string) (*messaging.Mod
a.nats, subject, messaging.ModelsRunningRequest{}, 10*time.Second)
}

// backendStopAckTimeout bounds the wait for a worker's backend.stop reply.
//
// A single stop is bounded by the worker's 5s best-effort Free plus the kill,
// so this leaves comfortable headroom. It is also the stall a worker that
// predates the reply imposes on every stop, which is why it is not generous:
// such a worker performs the stop and simply never answers, so the controller
// waits out the full budget before falling back to the old assumption.
const backendStopAckTimeout = 15 * time.Second

// StopBackend tells a worker node to stop a specific gRPC backend process.
// If backend is empty, the worker stops ALL backends.
// The node stays registered and can receive another InstallBackend later.
func (a *RemoteUnloaderAdapter) StopBackend(nodeID, backend string) error {
return a.stopBackend(nodeID, backend, false)
}

// stopBackend asks a worker to stop a backend and waits for it to say what it
// did.
//
// This was a bare Publish, which returned nil as soon as the local publish
// succeeded and so reported success for a stop that killed nothing or failed
// outright. The unload endpoint answered HTTP 200 while the backend kept
// running and holding its VRAM, and the endpoint's own "backend stop failed"
// branch was unreachable.
//
// A silent worker is NOT treated as a failure. A worker that predates
// BackendStopReply still receives the request and still stops the backend; it
// only lacks the reply. Failing here would break every stop against a fleet
// that has not been upgraded yet, so a timeout degrades to the previous
// assumption and says so in the log.
//
// Only silence degrades. A transport error — the connection is closed, nothing
// subscribes to this node's subject at all — is reported, because under the
// old Publish it was reported too, and callers such as UnloadRemoteModel rely
// on that to skip the registry cleanup for a node they could not reach.
func (a *RemoteUnloaderAdapter) stopBackend(nodeID, backend string, force bool) error {
subject := messaging.SubjectNodeBackendStop(nodeID)
if backend == "" && !force {
return a.nats.Publish(subject, nil)
req := messaging.BackendStopRequest{Backend: backend, Force: force}

reply, err := messaging.RequestJSON[messaging.BackendStopRequest, messaging.BackendStopReply](
a.nats, subject, req, backendStopAckTimeout)
if err != nil {
if errors.Is(err, nats.ErrTimeout) {
xlog.Warn("Worker did not acknowledge backend.stop; assuming an older worker delivered it",
"nodeID", nodeID, "backend", backend, "force", force)
return nil
}
return fmt.Errorf("backend.stop on node %s: %w", nodeID, err)
}
if !reply.Success {
return fmt.Errorf("backend.stop on node %s: %s", nodeID, reply.Error)
}
return a.nats.Publish(subject, messaging.BackendStopRequest{Backend: backend, Force: force})
// An empty list from a worker that enumerates what it stopped is the answer
// to "was anything actually running under that name" — and the answer is
// no. That is not an error: stopping a backend that is not running leaves
// the caller in the state it asked for. It is worth saying out loud,
// because a caller that expected to reclaim VRAM did not.
if reply.ReportsStoppedProcesses && len(reply.StoppedProcessKeys) == 0 {
xlog.Warn("backend.stop matched no running process on the worker",
"nodeID", nodeID, "backend", backend)
return nil
}
xlog.Info("Worker stopped backend processes",
"nodeID", nodeID, "backend", backend, "stopped", reply.StoppedProcessKeys)
return nil
}

// DeleteBackend tells a worker node to delete a backend (stop + remove files).
Expand Down
94 changes: 80 additions & 14 deletions core/services/nodes/unloader_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -124,6 +124,12 @@ type fakeSubscription struct{}

func (f *fakeSubscription) Unsubscribe() error { return nil }

func mustJSON(v any) []byte {
data, err := json.Marshal(v)
Expect(err).ToNot(HaveOccurred())
return data
}

// --- Tests ---

var _ = Describe("RemoteUnloaderAdapter", func() {
Expand All @@ -136,6 +142,14 @@ var _ = Describe("RemoteUnloaderAdapter", func() {
BeforeEach(func() {
locator = &fakeModelLocator{}
mc = &fakeMessagingClient{}
// backend.stop is request-reply, so the default fake must answer the
// way a current worker does. Specs that care about the reply override
// requestReply themselves.
mc.requestReply = mustJSON(messaging.BackendStopReply{
Success: true,
StoppedProcessKeys: []string{"llama#0"},
ReportsStoppedProcesses: true,
})
adapter = NewRemoteUnloaderAdapter(locator, mc, 3*time.Minute, 15*time.Minute)
})

Expand Down Expand Up @@ -179,7 +193,7 @@ var _ = Describe("RemoteUnloaderAdapter", func() {
// tests/e2e/distributed/node_lifecycle_test.go — keep them in step.
locator.nodes = nil
Expect(adapter.UnloadRemoteModel("my-model")).To(Succeed())
Expect(mc.published).To(BeEmpty())
Expect(mc.requestCalls).To(BeEmpty())
})

It("broadcasts to all nodes with model", func() {
Expand All @@ -189,10 +203,10 @@ var _ = Describe("RemoteUnloaderAdapter", func() {
}
Expect(adapter.UnloadRemoteModel("llama")).To(Succeed())

// Should have published a StopBackend for each node.
Expect(mc.published).To(HaveLen(2))
Expect(mc.published[0].Subject).To(Equal(messaging.SubjectNodeBackendStop("node-1")))
Expect(mc.published[1].Subject).To(Equal(messaging.SubjectNodeBackendStop("node-2")))
// Should have asked each node to stop the backend.
Expect(mc.requestCalls).To(HaveLen(2))
Expect(mc.requestCalls[0].Subject).To(Equal(messaging.SubjectNodeBackendStop("node-1")))
Expect(mc.requestCalls[1].Subject).To(Equal(messaging.SubjectNodeBackendStop("node-2")))

// Should have removed the model from each node in the registry.
Expect(locator.removedPairs).To(HaveLen(2))
Expand All @@ -205,7 +219,7 @@ var _ = Describe("RemoteUnloaderAdapter", func() {
{ID: "node-fail", Name: "worker-fail"},
{ID: "node-ok", Name: "worker-ok"},
}
// Use a messaging client that fails the first Publish call only.
// Use a messaging client that fails the first Request call only.
failOnce := &failOnceMessagingClient{inner: mc, failOn: 0}
adapter = NewRemoteUnloaderAdapter(locator, failOnce, 3*time.Minute, 15*time.Minute)

Expand All @@ -223,25 +237,70 @@ var _ = Describe("RemoteUnloaderAdapter", func() {
Expect(adapter.UnloadRemoteModelContext(context.Background(), "llama", true)).To(Succeed())

var payload messaging.BackendStopRequest
Expect(json.Unmarshal(mc.published[0].Data, &payload)).To(Succeed())
Expect(json.Unmarshal(mc.requestCalls[0].Data, &payload)).To(Succeed())
Expect(payload).To(Equal(messaging.BackendStopRequest{Backend: "llama", Force: true}))
})
})

Describe("StopBackend", func() {
It("with empty backend publishes nil payload", func() {
It("with empty backend asks the worker to stop everything", func() {
Expect(adapter.StopBackend("node-1", "")).To(Succeed())
Expect(mc.published).To(HaveLen(1))
Expect(mc.published[0].Subject).To(Equal(messaging.SubjectNodeBackendStop("node-1")))
Expect(mc.published[0].Data).To(BeNil())
Expect(mc.requestCalls).To(HaveLen(1))
Expect(mc.requestCalls[0].Subject).To(Equal(messaging.SubjectNodeBackendStop("node-1")))

// An empty Backend is the wire signal for "stop all"; the worker's
// decodeBackendStopRequest reads it the same way it read the bare
// nil payload this replaced.
var payload messaging.BackendStopRequest
Expect(json.Unmarshal(mc.requestCalls[0].Data, &payload)).To(Succeed())
Expect(payload.Backend).To(BeEmpty())
})

It("with backend name publishes JSON", func() {
// The bug this reply exists for: the worker could not stop what was
// asked, and the caller was told everything was fine.
It("reports a stop the worker could not carry out", func() {
mc.requestReply = mustJSON(messaging.BackendStopReply{
Success: false,
Error: "llama#0: process refused to die",
ReportsStoppedProcesses: true,
})
err := adapter.StopBackend("node-1", "llama-backend")
Expect(err).To(HaveOccurred())
Expect(err.Error()).To(ContainSubstring("process refused to die"))
})

// Nothing running under that name is the state the caller asked for, so
// it stays a success — eviction and cleanup paths stop models that are
// already gone all the time.
It("succeeds when the worker matched no running process", func() {
mc.requestReply = mustJSON(messaging.BackendStopReply{
Success: true,
ReportsStoppedProcesses: true,
})
Expect(adapter.StopBackend("node-1", "llama-backend")).To(Succeed())
Expect(mc.published).To(HaveLen(1))
})

// A worker built before BackendStopReply performs the stop and never
// answers. Failing here would break every stop on a fleet mid-upgrade.
It("assumes delivery when an older worker never answers", func() {
mc.requestErr = nats.ErrTimeout
Expect(adapter.StopBackend("node-1", "llama-backend")).To(Succeed())
})

// A closed connection is not an old worker, and callers depend on
// hearing about it: UnloadRemoteModel skips the registry cleanup for a
// node it could not reach.
It("reports a transport failure rather than assuming delivery", func() {
mc.requestErr = nats.ErrConnectionClosed
Expect(adapter.StopBackend("node-1", "llama-backend")).To(HaveOccurred())
})

It("with backend name sends JSON", func() {
Expect(adapter.StopBackend("node-1", "llama-backend")).To(Succeed())
Expect(mc.requestCalls).To(HaveLen(1))

var payload messaging.BackendStopRequest
Expect(json.Unmarshal(mc.published[0].Data, &payload)).To(Succeed())
Expect(json.Unmarshal(mc.requestCalls[0].Data, &payload)).To(Succeed())
Expect(payload.Backend).To(Equal("llama-backend"))
Expect(payload.Force).To(BeFalse())
})
Expand Down Expand Up @@ -335,6 +394,13 @@ func (f *failOnceMessagingClient) SubscribeReply(subject string, handler func(da
}

func (f *failOnceMessagingClient) Request(subject string, data []byte, timeout time.Duration) ([]byte, error) {
f.mu.Lock()
idx := f.callIdx
f.callIdx++
f.mu.Unlock()
if idx == f.failOn {
return nil, fmt.Errorf("simulated failure")
}
return f.inner.Request(subject, data, timeout)
}

Expand Down
40 changes: 36 additions & 4 deletions core/services/worker/lifecycle.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import (
"maps"
"net"
"slices"
"strings"
"syscall"

"github.com/mudler/LocalAI/core/gallery"
Expand All @@ -27,7 +28,7 @@ func (s *backendSupervisor) subscribeLifecycleEvents() error {
if _, err := s.nats.SubscribeReply(messaging.SubjectNodeBackendUpgrade(s.nodeID), s.handleBackendUpgrade); err != nil {
return fmt.Errorf("subscribing to backend upgrade events: %w", err)
}
if _, err := s.nats.Subscribe(messaging.SubjectNodeBackendStop(s.nodeID), s.handleBackendStop); err != nil {
if _, err := s.nats.SubscribeReply(messaging.SubjectNodeBackendStop(s.nodeID), s.handleBackendStop); err != nil {
return fmt.Errorf("subscribing to backend stop events: %w", err)
}
if _, err := s.nats.SubscribeReply(messaging.SubjectNodeBackendDelete(s.nodeID), s.handleBackendDelete); err != nil {
Expand Down Expand Up @@ -154,27 +155,58 @@ func (s *backendSupervisor) handleBackendUpgrade(data []byte, reply func([]byte)
}

// handleBackendStop is the NATS callback for backend.stop — stop a specific
// backend process (fire-and-forget, no reply expected).
func (s *backendSupervisor) handleBackendStop(data []byte) {
// backend process and report what it terminated.
//
// The reply is what lets the controller tell a stop that worked from one that
// matched nothing or failed. Callers that publish without a reply subject (an
// older controller) still work: SubscribeReply drops the response.
func (s *backendSupervisor) handleBackendStop(data []byte, reply func([]byte)) {
req, stopAll, err := decodeBackendStopRequest(data)
if err != nil {
xlog.Error("Ignoring malformed NATS backend.stop event", "error", err)
replyJSON(reply, messaging.BackendStopReply{
Error: fmt.Sprintf("invalid request: %v", err),
ReportsStoppedProcesses: true,
})
return
}
if stopAll {
xlog.Info("Received NATS backend.stop event (all)", "force", req.Force)
s.stopAllBackends(req.Force)
stopped := s.stopAllBackends(req.Force)
replyJSON(reply, messaging.BackendStopReply{
Success: true,
StoppedProcessKeys: stopped,
ReportsStoppedProcesses: true,
})
return
}
xlog.Info("Received NATS backend.stop event", "backend", req.Backend, "force", req.Force)
// The identifier may be a backend name, a model name, or an exact
// modelID#replica key depending on the publisher; resolveStopTargets
// handles all three. stopBackend alone resolves only the model meanings.
var stopped []string
var failures []string
for _, key := range s.resolveStopTargets(req.Backend) {
if err := s.stopBackendExact(key, req.Force); err != nil {
xlog.Error("Failed to stop backend process", "backend", req.Backend, "processKey", key, "error", err)
failures = append(failures, fmt.Sprintf("%s: %v", key, err))
continue
}
stopped = append(stopped, key)
}
// Resolving to nothing is reported as success with an empty list, not as a
// failure: stopping a backend that is not running is the state the caller
// asked for. The empty list is what tells the caller nothing matched, and
// ReportsStoppedProcesses is what makes that emptiness trustworthy.
res := messaging.BackendStopReply{
Success: len(failures) == 0,
StoppedProcessKeys: stopped,
ReportsStoppedProcesses: true,
}
if len(failures) > 0 {
res.Error = strings.Join(failures, "; ")
}
replyJSON(reply, res)
}

func decodeBackendStopRequest(data []byte) (messaging.BackendStopRequest, bool, error) {
Expand Down
7 changes: 5 additions & 2 deletions core/services/worker/supervisor.go
Original file line number Diff line number Diff line change
Expand Up @@ -955,15 +955,18 @@ func (s *backendSupervisor) finishBackendStop(key string, bp *backendProcess, st
return nil
}

// stopAllBackends stops all running backend processes.
func (s *backendSupervisor) stopAllBackends(force bool) {
// stopAllBackends stops all running backend processes and returns the process
// keys it attempted, so a caller answering a backend.stop request can report
// what it acted on.
func (s *backendSupervisor) stopAllBackends(force bool) []string {
s.mu.Lock()
backends := slices.Collect(maps.Keys(s.processes))
s.mu.Unlock()

for _, b := range backends {
s.stopBackend(b, force)
}
return backends
}

// isRunning returns whether at least one backend process matching the given
Expand Down
15 changes: 15 additions & 0 deletions docs/content/features/distributed-mode.md
Original file line number Diff line number Diff line change
Expand Up @@ -560,6 +560,21 @@ The responses from `GET /api/node/:id/models` and `GET /api/nodes/:id/models` in

`model.unload` releases model memory inside a running backend. It does not replace the exact process stop that configuration cleanup requires. The `backend.stop` operation remains an administrative backend operation.

#### `backend.stop` is acknowledged

`backend.stop` is request-reply. The worker answers with what it terminated, so the controller can tell a stop that worked from one that matched nothing or failed outright.

This matters for `POST /api/nodes/:id/models/unload`, which stops the backend after unloading the model. The stop used to be fire-and-forget, so the endpoint answered `200` as soon as the message left the frontend — including when the backend was still running and still holding its VRAM. It now returns an error when the worker reports that the stop failed.

Two outcomes are deliberately **not** errors:

- **Nothing matched.** The worker reports an empty stopped-process list, logged as `backend.stop matched no running process`. Stopping a backend that is not running leaves the caller in the state it asked for, and eviction and cleanup paths stop already-gone models routinely.
- **No answer.** A worker built before this reply performs the stop and never responds. The controller waits 15 seconds, logs `Worker did not acknowledge backend.stop`, and assumes delivery, so a fleet mid-upgrade keeps working. A transport failure is reported rather than assumed.

{{% notice note %}}
On a mixed fleet, every stop against a worker that predates the reply costs the full 15-second wait before falling back. Upgrading the workers removes the delay.
{{% /notice %}}

### Per-node VRAM budget

Each worker advertises its detected VRAM, and the SmartRouter uses that number when picking a node with enough free memory. You can cap the VRAM a node offers for placement so it never gets scheduled beyond a chosen limit, leaving headroom for other workloads on that machine.
Expand Down
Loading