diff --git a/core/services/nodes/unloader.go b/core/services/nodes/unloader.go index 3b1cd15b8..460be8acf 100644 --- a/core/services/nodes/unloader.go +++ b/core/services/nodes/unloader.go @@ -372,17 +372,36 @@ func (a *RemoteUnloaderAdapter) ListBackends(nodeID string) (*messaging.BackendL // scheduler could pick a node that could not be given work and the request // failed with "no responders available". // -// It reuses the models.running subject rather than a dedicated ping subject on -// purpose: a new subject would go unanswered by any worker that has not been -// upgraded yet, and this check would then report every one of them as dead. -// The worker answers out of its in-memory process table, so a live node -// replies immediately, and NATS reports no-responders without waiting out the -// timeout. +// The subject asked has to be one every worker in the fleet subscribes to, or +// this check condemns the workers that do not. models.running was the obvious +// choice and the wrong one: it arrived in 4.6, so a 4.5 worker that is alive +// and serving never answers it, and a model pinned to that node could never be +// scheduled. backend.list has been part of the worker protocol far longer, so +// it is the safer question to ask. +// +// A worker that answers anything is alive. Only when every subject reports no +// responders is the node treated as absent, so adding a newer subject here can +// never condemn an older worker. func (a *RemoteUnloaderAdapter) PingNode(nodeID string) error { - subject := messaging.SubjectNodeModelsRunning(nodeID) - _, err := messaging.RequestJSON[messaging.ModelsRunningRequest, messaging.ModelsRunningReply]( - a.nats, subject, messaging.ModelsRunningRequest{}, 5*time.Second) - return err + subjects := []string{ + messaging.SubjectNodeBackendList(nodeID), + messaging.SubjectNodeModelsRunning(nodeID), + } + var lastErr error + for _, subject := range subjects { + _, err := messaging.RequestJSON[messaging.BackendListRequest, messaging.BackendListReply]( + a.nats, subject, messaging.BackendListRequest{}, 5*time.Second) + if err == nil { + return nil + } + if !errors.Is(err, nats.ErrNoResponders) { + // Reached someone, or failed for a reason that is not absence. + // Either way the node is not proven gone. + return nil + } + lastErr = err + } + return lastErr } // ListRunningModels asks a worker node which model backend processes it diff --git a/core/services/nodes/unloader_ping_test.go b/core/services/nodes/unloader_ping_test.go new file mode 100644 index 000000000..a9b3a5889 --- /dev/null +++ b/core/services/nodes/unloader_ping_test.go @@ -0,0 +1,51 @@ +package nodes + +import ( + "errors" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "github.com/nats-io/nats.go" + + "github.com/mudler/LocalAI/core/services/messaging" +) + +// The scheduler's liveness probe asks a worker a question over NATS and treats +// "no responders" as proof the worker is gone. That is only sound if every +// worker in the fleet subscribes to the subject asked. +// +// It originally asked models.running, which arrived in 4.6. A 4.5 worker is +// perfectly alive and serving, answers backend.list, and never subscribes to +// models.running, so the probe condemned it on every scheduling attempt. A +// model pinned to such a node could then never be placed at all. +var _ = Describe("Node liveness probe subject", func() { + var ( + mc *scriptedMessagingClient + adapter *RemoteUnloaderAdapter + ) + + const nodeID = "11111111-2222-3333-4444-555555555555" + + BeforeEach(func() { + mc = newScriptedMessagingClient() + adapter = NewRemoteUnloaderAdapter(nil, mc, 3*time.Minute, 15*time.Minute) + }) + + It("treats a worker that answers backend.list as alive", func() { + // A worker old enough to predate models.running: it answers the + // long-standing backend.list subject and nothing else. + mc.scriptReply(messaging.SubjectNodeBackendList(nodeID), messaging.BackendListReply{}) + mc.scriptNoResponders(messaging.SubjectNodeModelsRunning(nodeID)) + + Expect(errors.Is(adapter.PingNode(nodeID), nats.ErrNoResponders)).To(BeFalse(), + "a worker answering backend.list is alive regardless of newer subjects") + }) + + It("still reports a worker that answers nothing as absent", func() { + mc.scriptNoResponders(messaging.SubjectNodeBackendList(nodeID)) + mc.scriptNoResponders(messaging.SubjectNodeModelsRunning(nodeID)) + + Expect(errors.Is(adapter.PingNode(nodeID), nats.ErrNoResponders)).To(BeTrue()) + }) +})