diff --git a/core/application/distributed.go b/core/application/distributed.go index b7dc0bf91..404b2112c 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -5,6 +5,8 @@ import ( "encoding/json" "fmt" "io" + "net" + "strconv" "strings" "sync" "time" @@ -12,12 +14,14 @@ import ( "github.com/google/uuid" "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/services/agents" + "github.com/mudler/LocalAI/core/services/cluster" "github.com/mudler/LocalAI/core/services/distributed" "github.com/mudler/LocalAI/core/services/jobs" "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/core/services/storage" + "github.com/mudler/LocalAI/internal" "github.com/mudler/LocalAI/pkg/distributedhdr" "github.com/mudler/LocalAI/pkg/sanitize" "github.com/mudler/xlog" @@ -43,6 +47,16 @@ type DistributedServices struct { Unloader *nodes.RemoteUnloaderAdapter ModelCleanup *nodes.ModelCleanupService + // Cluster is the replica-membership registry: which frontend replicas are + // alive, at which address, and which of them holds a given worker's tunnel. + Cluster *cluster.Registry + // Membership publishes this replica's row and reaps the dead. Nil when no + // peer-reachable address could be determined, which leaves this replica + // invisible to its peers but otherwise fully functional. + Membership *cluster.Membership + // PeerSessions owns the peer links other replicas dialled into this one. + PeerSessions *cluster.SessionStore + shutdownOnce sync.Once } @@ -53,6 +67,15 @@ func (ds *DistributedServices) Shutdown() { return } ds.shutdownOnce.Do(func() { + // Peer state first: a replica that is going away should stop claiming + // to be alive before it stops answering, so peers re-home rather than + // dial a process in teardown. + if ds.Membership != nil { + ds.Membership.Stop() + } + if ds.PeerSessions != nil { + ds.PeerSessions.CloseAll() + } if ds.Health != nil { ds.Health.Stop() } @@ -162,6 +185,29 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade } xlog.Info("Node registry initialized") + // Replica membership. NewNodeRegistry has just migrated the tables this + // reads, so it has to come after it. + clusterRegistry := cluster.NewRegistry(authDB) + // Accepted peer links are held with no stream handler: this replica has + // somewhere to put a link a peer dials, and refuses the streams on it, + // because nothing relays worker traffic yet. + peerSessions := cluster.NewSessionStore(nil) + var membership *cluster.Membership + if advertised, err := advertisedPeerAddr(cfg); err != nil { + // Not fatal. A replica that cannot publish an address still serves + // every request that reaches it directly; what it cannot do is have + // another replica relay to it. Failing startup here would take out + // every existing single-host deployment, whose route to a local + // database is loopback. + xlog.Warn("This replica will not be reachable by its peers: no advertised address", + "error", err, "knob", "LOCALAI_DISTRIBUTED_ADVERTISE_ADDR") + } else { + membership = cluster.NewMembership(clusterRegistry, cfg.Distributed.InstanceID, advertised, internal.PrintableVersion()) + if err := membership.Start(cfg.Context); err != nil { + return nil, fmt.Errorf("registering this replica in the cluster: %w", err) + } + } + // Let scheduling rules be keyed by a model alias. The registry resolves a // rule's name through the config loader to find the model it governs, so an // operator can pin placement to a stable name like "production" and have it @@ -450,9 +496,36 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade ModelAdapter: modelAdapter, Unloader: remoteUnloader, ModelCleanup: modelCleanup, + Cluster: clusterRegistry, + Membership: membership, + PeerSessions: peerSessions, }, nil } +// advertisedPeerAddr is the host:port peers dial to reach this replica. +// +// The operator's value wins outright. Otherwise it is derived from the port +// this process serves on and the local address that routes to PostgreSQL, which +// is only a peer-reachable answer when the database is on another host; +// DiscoverAdvertisedAddr refuses rather than guessing when it is not. +func advertisedPeerAddr(cfg *config.ApplicationConfig) (string, error) { + if cfg.Distributed.AdvertiseAddr != "" { + return cfg.Distributed.AdvertiseAddr, nil + } + if cfg.APIAddress == "" { + return "", fmt.Errorf("no API address to derive a peer port from") + } + _, port, err := net.SplitHostPort(cfg.APIAddress) + if err != nil { + return "", fmt.Errorf("reading the peer port out of API address %q: %w", cfg.APIAddress, err) + } + portNumber, err := strconv.Atoi(port) + if err != nil { + return "", fmt.Errorf("API address %q has a non-numeric port: %w", cfg.APIAddress, err) + } + return cluster.DiscoverAdvertisedAddr(cfg.Auth.DatabaseURL, portNumber) +} + func isPostgresURL(url string) bool { return strings.HasPrefix(url, "postgres://") || strings.HasPrefix(url, "postgresql://") } diff --git a/core/cli/run.go b/core/cli/run.go index 6b9b3e4dc..4408edd12 100644 --- a/core/cli/run.go +++ b/core/cli/run.go @@ -165,6 +165,7 @@ type RunCMD struct { Distributed bool `env:"LOCALAI_DISTRIBUTED" default:"false" help:"Enable distributed mode (requires PostgreSQL + NATS)" group:"distributed"` InstanceID string `env:"LOCALAI_INSTANCE_ID" help:"Unique instance ID for distributed mode (auto-generated UUID if empty)" group:"distributed"` NatsURL string `env:"LOCALAI_NATS_URL" help:"NATS server URL (e.g., nats://localhost:4222)" group:"distributed"` + DistributedAdvertiseAddr string `env:"LOCALAI_DISTRIBUTED_ADVERTISE_ADDR" help:"host:port other frontend replicas dial to reach this one (peer link). Empty = derived from the local address that routes to PostgreSQL, which only works when the database is on another host." group:"distributed"` StorageURL string `env:"LOCALAI_STORAGE_URL" help:"S3-compatible storage endpoint URL (e.g., http://minio:9000)" group:"distributed"` StorageBucket string `env:"LOCALAI_STORAGE_BUCKET" default:"localai" help:"S3 bucket name for object storage" group:"distributed"` StorageRegion string `env:"LOCALAI_STORAGE_REGION" default:"us-east-1" help:"S3 region" group:"distributed"` @@ -351,6 +352,9 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error { if r.InstanceID != "" { opts = append(opts, config.WithDistributedInstanceID(r.InstanceID)) } + if r.DistributedAdvertiseAddr != "" { + opts = append(opts, config.WithDistributedAdvertiseAddr(r.DistributedAdvertiseAddr)) + } if r.NatsURL != "" { opts = append(opts, config.WithNatsURL(r.NatsURL)) } diff --git a/core/config/distributed_config.go b/core/config/distributed_config.go index 5a48a84e9..bbb141592 100644 --- a/core/config/distributed_config.go +++ b/core/config/distributed_config.go @@ -13,8 +13,15 @@ import ( // DistributedConfig holds configuration for horizontal scaling mode. // When Enabled is true, PostgreSQL and NATS are required. type DistributedConfig struct { - Enabled bool // --distributed / LOCALAI_DISTRIBUTED - InstanceID string // --instance-id / LOCALAI_INSTANCE_ID (auto-generated UUID if empty) + Enabled bool // --distributed / LOCALAI_DISTRIBUTED + InstanceID string // --instance-id / LOCALAI_INSTANCE_ID (auto-generated UUID if empty) + // AdvertiseAddr is the host:port OTHER REPLICAS dial to reach this one, + // which is not the address this process binds: a replica behind a service + // or a NAT binds one and is reached at another. Empty means "work it out", + // by asking the kernel which local address routes to PostgreSQL; that + // answer is only usable when the database is remote, so a deployment with + // a local or sidecar database has to set this. + AdvertiseAddr string // LOCALAI_DISTRIBUTED_ADVERTISE_ADDR NatsURL string // --nats-url / LOCALAI_NATS_URL StorageURL string // --storage-url / LOCALAI_STORAGE_URL (S3 endpoint) RegistrationToken string // --registration-token / LOCALAI_REGISTRATION_TOKEN (required token for node registration) @@ -195,6 +202,14 @@ func WithDistributedInstanceID(id string) AppOption { } } +// WithDistributedAdvertiseAddr pins the host:port peers dial to reach this +// replica, overriding the route-based discovery. +func WithDistributedAdvertiseAddr(addr string) AppOption { + return func(o *ApplicationConfig) { + o.Distributed.AdvertiseAddr = addr + } +} + func WithNatsURL(url string) AppOption { return func(o *ApplicationConfig) { o.Distributed.NatsURL = url diff --git a/core/http/app.go b/core/http/app.go index 2e1453ac0..d6418165b 100644 --- a/core/http/app.go +++ b/core/http/app.go @@ -576,6 +576,14 @@ func API(application *application.Application) (*echo.Echo, error) { routes.RegisterNodeSelfServiceRoutes(e, registry, distCfg.RegistrationToken, distCfg.AutoApproveNodes, application.AuthDB(), application.ApplicationConfig().Auth.APIKeyHMACSecret, natsCfg) routes.RegisterNodeAdminRoutes(e, registry, remoteUnloader, application.GalleryService(), opcache, application.ApplicationConfig(), adminMiddleware, application.AuthDB(), application.ApplicationConfig().Auth.APIKeyHMACSecret, application.ApplicationConfig().Distributed.RegistrationToken, natsCfg) + // Replica-to-replica peer link. Registered only in distributed mode: in + // single-node mode there are no peers, and the route authenticates with the + // registration token, so publishing it unconditionally would put a + // multiplexer on every single-binary install. + if d := application.Distributed(); d != nil && d.PeerSessions != nil { + routes.RegisterClusterRoutes(e, distCfg.RegistrationToken, d.PeerSessions.Accept) + } + // Distributed SSE routes (job progress + agent events via NATS) if d := application.Distributed(); d != nil { if d.Dispatcher != nil { diff --git a/core/http/endpoints/cluster/peer.go b/core/http/endpoints/cluster/peer.go index 85985eee5..f0fb96ebc 100644 --- a/core/http/endpoints/cluster/peer.go +++ b/core/http/endpoints/cluster/peer.go @@ -18,18 +18,6 @@ import ( "github.com/mudler/xlog" ) -// RegisterClusterRoutes registers the peer link. onPeer receives every -// authenticated session; see PeerHandler for what it is expected to do with it. -// -// The route is core/services/cluster's own constant, so the handler and the -// dialler cannot be registered and dialled at different paths. That the path -// also falls under auth.ClusterPathPrefix, and so bypasses the session -// middleware, is asserted by a spec in this package: it is the only place that -// can see both, since core/services/cluster must not import core/http/auth. -func RegisterClusterRoutes(e *echo.Echo, token string, onPeer func(string, *yamux.Session)) { - e.GET(clustersvc.PeerPath, PeerHandler(token, onPeer)) -} - // PeerHandler upgrades an authenticated peer dial to a WebSocket, wraps it as // a yamux server session and hands it to onSession. // diff --git a/core/http/endpoints/cluster/peer_test.go b/core/http/endpoints/cluster/peer_test.go index a86ab8add..dc8cba3da 100644 --- a/core/http/endpoints/cluster/peer_test.go +++ b/core/http/endpoints/cluster/peer_test.go @@ -5,8 +5,9 @@ import ( "net/http/httptest" "strings" + "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/http/auth" - clusterep "github.com/mudler/LocalAI/core/http/endpoints/cluster" + "github.com/mudler/LocalAI/core/http/routes" clustersvc "github.com/mudler/LocalAI/core/services/cluster" "github.com/gorilla/websocket" @@ -16,6 +17,11 @@ import ( . "github.com/onsi/gomega" ) +// wsPeerURL is the peer route on a test server, named as peer-1. +func wsPeerURL(s *httptest.Server) string { + return "ws" + strings.TrimPrefix(s.URL, "http") + clustersvc.PeerPath + "?id=peer-1" +} + var _ = Describe("Peer link handler", func() { var ( srv *httptest.Server @@ -25,19 +31,15 @@ var _ = Describe("Peer link handler", func() { BeforeEach(func() { sessions = make(chan *yamux.Session, 1) e := echo.New() - clusterep.RegisterClusterRoutes(e, "peer-token", func(_ string, s *yamux.Session) { + routes.RegisterClusterRoutes(e, "peer-token", func(_ string, s *yamux.Session) { sessions <- s }) srv = httptest.NewServer(e) DeferCleanup(srv.Close) }) - wsURL := func(s *httptest.Server) string { - return "ws" + strings.TrimPrefix(s.URL, "http") + "/api/cluster/peer?id=peer-1" - } - It("rejects a connection with no token", func() { - _, resp, err := websocket.DefaultDialer.Dial(wsURL(srv), nil) + _, resp, err := websocket.DefaultDialer.Dial(wsPeerURL(srv), nil) Expect(err).To(HaveOccurred()) Expect(resp).ToNot(BeNil()) Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) @@ -46,7 +48,7 @@ var _ = Describe("Peer link handler", func() { It("rejects a connection with the wrong token", func() { h := http.Header{} h.Set("Authorization", "Bearer wrong") - _, resp, err := websocket.DefaultDialer.Dial(wsURL(srv), h) + _, resp, err := websocket.DefaultDialer.Dial(wsPeerURL(srv), h) Expect(err).To(HaveOccurred()) Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) }) @@ -54,7 +56,7 @@ var _ = Describe("Peer link handler", func() { It("accepts an authenticated peer and yields a usable yamux session", func() { h := http.Header{} h.Set("Authorization", "Bearer peer-token") - conn, _, err := websocket.DefaultDialer.Dial(wsURL(srv), h) + conn, _, err := websocket.DefaultDialer.Dial(wsPeerURL(srv), h) Expect(err).ToNot(HaveOccurred()) DeferCleanup(func() { _ = conn.Close() }) @@ -98,11 +100,11 @@ var _ = Describe("Peer link handler", func() { h.Set("Authorization", "Bearer peer-token") ids := make(chan string, 1) e := echo.New() - clusterep.RegisterClusterRoutes(e, "peer-token", func(id string, _ *yamux.Session) { ids <- id }) + routes.RegisterClusterRoutes(e, "peer-token", func(id string, _ *yamux.Session) { ids <- id }) s2 := httptest.NewServer(e) DeferCleanup(s2.Close) - conn, _, err := websocket.DefaultDialer.Dial(wsURL(s2), h) + conn, _, err := websocket.DefaultDialer.Dial(wsPeerURL(s2), h) Expect(err).ToNot(HaveOccurred()) DeferCleanup(func() { _ = conn.Close() }) @@ -115,12 +117,12 @@ var _ = Describe("Peer link handler", func() { // unauthenticated yamux multiplexer to anyone who can reach the port. e := echo.New() accepted := make(chan *yamux.Session, 1) - clusterep.RegisterClusterRoutes(e, "", func(_ string, sess *yamux.Session) { accepted <- sess }) + routes.RegisterClusterRoutes(e, "", func(_ string, sess *yamux.Session) { accepted <- sess }) s2 := httptest.NewServer(e) DeferCleanup(s2.Close) for _, header := range []http.Header{nil, {"Authorization": []string{"Bearer "}}, {"Authorization": []string{"Bearer anything"}}} { - _, resp, err := websocket.DefaultDialer.Dial(wsURL(s2), header) + _, resp, err := websocket.DefaultDialer.Dial(wsPeerURL(s2), header) Expect(err).To(HaveOccurred()) Expect(resp).ToNot(BeNil()) Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) @@ -132,7 +134,7 @@ var _ = Describe("Peer link handler", func() { // RFC 7235 makes the scheme case-insensitive. The token after it is not. h := http.Header{} h.Set("Authorization", "bearer peer-token") - conn, _, err := websocket.DefaultDialer.Dial(wsURL(srv), h) + conn, _, err := websocket.DefaultDialer.Dial(wsPeerURL(srv), h) Expect(err).ToNot(HaveOccurred()) DeferCleanup(func() { _ = conn.Close() }) Eventually(sessions, "5s").Should(Receive()) @@ -143,7 +145,7 @@ var _ = Describe("Peer link handler", func() { // without the handler's own recover the peer would keep a link nobody // ever accepts streams on. e := echo.New() - clusterep.RegisterClusterRoutes(e, "peer-token", func(_ string, _ *yamux.Session) { + routes.RegisterClusterRoutes(e, "peer-token", func(_ string, _ *yamux.Session) { panic("callback exploded") }) s2 := httptest.NewServer(e) @@ -157,7 +159,7 @@ var _ = Describe("Peer link handler", func() { h := http.Header{} h.Set("Authorization", "Bearer peer-token") - conn, _, err := websocket.DefaultDialer.Dial(wsURL(s2), h) + conn, _, err := websocket.DefaultDialer.Dial(wsPeerURL(s2), h) Expect(err).ToNot(HaveOccurred()) DeferCleanup(func() { _ = conn.Close() }) @@ -186,17 +188,69 @@ var _ = Describe("Peer link handler", func() { }) }) -var _ = Describe("Peer link auth prefix", func() { - It("keeps the peer route inside the alternative-authentication prefix", func() { - // The peer route authenticates with the cluster token, not the global - // session middleware, which only holds while the route sits under the - // prefix auth exempts. Moving either one alone 401s every peer dial. - // - // This lives here because it is the only package that can see both: - // core/services/cluster owns the route and must stay free of any - // core/http dependency, and core/http/auth owns the exemption. - Expect(strings.HasPrefix(clustersvc.PeerPath, auth.ClusterPathPrefix)).To(BeTrue(), - "peer route %q is no longer under the auth-exempt prefix %q", - clustersvc.PeerPath, auth.ClusterPathPrefix) +var _ = Describe("Peer link auth coverage", func() { + // These specs put the REAL global auth middleware in front of the REAL + // registrar and prove a peer dial reaches the handler anyway. The peer link + // authenticates with the cluster token, not a session, so it only works + // while its path sits under the prefix auth exempts; moving either one + // alone 401s every peer dial, and the two live in packages that must not + // import each other. + // + // The predicate that grants the exemption is unexported, so this asserts on + // its effect rather than on it: what a caller can observe is whether the + // request reaches the handler. + var ( + srv *httptest.Server + sessions chan *yamux.Session + ) + + BeforeEach(func() { + sessions = make(chan *yamux.Session, 1) + e := echo.New() + // A nil DB with one legacy API key is the cheapest configuration that + // turns the middleware ON without a database. With neither, Middleware + // short-circuits to next() and every assertion below would pass against + // a server that has no auth at all. + e.Use(auth.Middleware(nil, &config.ApplicationConfig{ApiKeys: []string{"an-api-key"}})) + routes.RegisterClusterRoutes(e, "peer-token", func(_ string, s *yamux.Session) { sessions <- s }) + // A route outside the cluster prefix, registered on the same server, is + // the control: it proves the middleware in front of both is live. + e.GET("/api/nodes", func(c echo.Context) error { return c.NoContent(http.StatusOK) }) + srv = httptest.NewServer(e) + DeferCleanup(srv.Close) + }) + + It("refuses an uncredentialed request to a route outside the cluster prefix", func() { + resp, err := http.Get(srv.URL + "/api/nodes") + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = resp.Body.Close() }() + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized), + "the global auth middleware is not actually guarding this server, so the peer-route assertions below would prove nothing") + }) + + It("lets a peer dial reach the handler, which is the only thing that can authenticate it", func() { + // The cluster token is not one of the API keys the middleware knows, so + // a 400 from the handler's own missing-id check can only mean the + // request was let through unauthenticated by the middleware. + req, err := http.NewRequestWithContext(GinkgoT().Context(), http.MethodGet, srv.URL+clustersvc.PeerPath, nil) + Expect(err).ToNot(HaveOccurred()) + req.Header.Set("Authorization", "Bearer peer-token") + + resp, err := http.DefaultClient.Do(req) + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = resp.Body.Close() }() + Expect(resp.StatusCode).To(Equal(http.StatusBadRequest), + "a peer dial must reach the handler; 401 here means the peer route left the auth-exempt prefix %q", auth.ClusterPathPrefix) + }) + + It("completes a full peer handshake through the guarded server", func() { + // The status-code assertion above cannot see the upgrade, and the + // upgrade is what a peer actually does. + h := http.Header{} + h.Set("Authorization", "Bearer peer-token") + conn, _, err := websocket.DefaultDialer.Dial(wsPeerURL(srv), h) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + Eventually(sessions, "5s").Should(Receive()) }) }) diff --git a/core/http/routes/cluster.go b/core/http/routes/cluster.go new file mode 100644 index 000000000..3e462db1c --- /dev/null +++ b/core/http/routes/cluster.go @@ -0,0 +1,25 @@ +package routes + +import ( + clusterep "github.com/mudler/LocalAI/core/http/endpoints/cluster" + clustersvc "github.com/mudler/LocalAI/core/services/cluster" + + "github.com/labstack/echo/v4" + "github.com/libp2p/go-yamux/v5" +) + +// RegisterClusterRoutes registers the replica-to-replica peer link. onPeer +// receives every authenticated session; see clusterep.PeerHandler for what it +// is expected to do with it. +// +// The path is core/services/cluster's own constant, so the handler and the +// dialler cannot be registered and dialled at different paths. That the path +// also falls under auth.ClusterPathPrefix, and so bypasses the global session +// middleware, is asserted by driving a request through that middleware in +// core/http/endpoints/cluster/peer_test.go. +// +// The route carries no auth middleware: it authenticates itself against the +// cluster token, because a peer replica has no session and no user. +func RegisterClusterRoutes(e *echo.Echo, token string, onPeer func(string, *yamux.Session)) { + e.GET(clustersvc.PeerPath, clusterep.PeerHandler(token, onPeer)) +} diff --git a/core/services/cluster/instance_test.go b/core/services/cluster/instance_test.go index 8e74e930e..d8075a7a3 100644 --- a/core/services/cluster/instance_test.go +++ b/core/services/cluster/instance_test.go @@ -22,9 +22,9 @@ var _ = Describe("Instance registry", func() { BeforeEach(func() { db = testutil.SetupTestDB() - Expect(db.AutoMigrate(&cluster.Instance{})).To(Succeed()) - reg = cluster.NewRegistry(db) ctx = context.Background() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + reg = cluster.NewRegistry(db) }) It("registers an instance and reads it back", func() { diff --git a/core/services/cluster/membership.go b/core/services/cluster/membership.go new file mode 100644 index 000000000..3b1182b17 --- /dev/null +++ b/core/services/cluster/membership.go @@ -0,0 +1,175 @@ +// SPDX-License-Identifier: MIT + +package cluster + +import ( + "context" + "errors" + "fmt" + "sync" + "time" + + "github.com/mudler/xlog" + "gorm.io/gorm" +) + +const ( + // InstanceHeartbeat is how often a replica refreshes its own row. + InstanceHeartbeat = 5 * time.Second + + // InstanceLiveness is how long a replica may go without a heartbeat before + // its peers treat it as gone: six consecutive misses. + // + // The window is generous on purpose. Declaring a replica dead deletes the + // connection rows it owned, and a worker whose row is deleted while its + // owner is merely slow has to be re-homed for nothing. The cost of waiting + // is bounded and symmetric: traffic for that worker is retried, not lost. + InstanceLiveness = 30 * time.Second +) + +// Membership publishes this replica's address and keeps the instances table +// free of replicas that have stopped answering. +// +// It is the only writer of this replica's row and the only sweeper of anyone +// else's, which is what keeps one fact on one clock: whether a replica is +// alive is answered by its last_seen and by nothing else. +type Membership struct { + reg *Registry + id string + addr string + version string + + interval time.Duration + liveness time.Duration + + stop chan struct{} + done chan struct{} + stopOnce sync.Once +} + +// NewMembership returns the membership loop for one replica. The address is +// what peers will dial, so it must be reachable from another host, not the +// address this process binds. +func NewMembership(reg *Registry, id, addr, version string) *Membership { + return &Membership{ + reg: reg, + id: id, + addr: addr, + version: version, + interval: InstanceHeartbeat, + liveness: InstanceLiveness, + stop: make(chan struct{}), + done: make(chan struct{}), + } +} + +// Start registers this replica and begins heartbeating and sweeping. The first +// registration is synchronous and its failure is returned: a replica whose +// address never reaches the table is invisible to its peers, and starting +// anyway would hide that behind a background log line. +func (m *Membership) Start(ctx context.Context) error { + if err := m.reg.Register(ctx, m.id, m.addr, m.version); err != nil { + return err + } + xlog.Info("Cluster instance registered", "id", m.id, "addr", m.addr) + go m.loop(ctx) + return nil +} + +// Stop ends the loop and waits for it. Safe to call more than once. +func (m *Membership) Stop() { + if m == nil { + return + } + m.stopOnce.Do(func() { close(m.stop) }) + <-m.done +} + +func (m *Membership) loop(ctx context.Context) { + defer close(m.done) + + ticker := time.NewTicker(m.interval) + defer ticker.Stop() + + for { + select { + case <-m.stop: + return + case <-ctx.Done(): + return + case <-ticker.C: + m.tick(ctx) + } + } +} + +// tick refreshes this replica's row and sweeps the dead. +// +// Every replica sweeps, rather than one elected sweeper. The deletes are +// idempotent and cheap, and an elected sweeper is one more thing that has to be +// alive for the cluster to notice that something is not. +func (m *Membership) tick(ctx context.Context) { + err := m.reg.Heartbeat(ctx, m.id) + if errors.Is(err, ErrInstanceNotFound) { + // Another replica swept this row while this process was stalled long + // enough to look dead. Re-register rather than heartbeat: a heartbeat + // carries no address, so the row has to be rebuilt from scratch. + xlog.Warn("Cluster instance row was reaped, re-registering", "id", m.id) + if err := m.reg.Register(ctx, m.id, m.addr, m.version); err != nil { + xlog.Error("Re-registering cluster instance failed", "id", m.id, "error", err) + } + } else if err != nil { + xlog.Warn("Cluster instance heartbeat failed", "id", m.id, "error", err) + } + + instances, connections, err := m.reg.ReapStale(ctx, m.id, m.liveness) + if err != nil { + xlog.Warn("Reaping stale cluster instances failed", "error", err) + return + } + if instances > 0 || connections > 0 { + xlog.Info("Reaped cluster state left by dead replicas", "instances", instances, "connections", connections) + } +} + +// ReapStale deletes the replicas that have not heartbeated within the liveness +// window, and the connection rows whose owner is no longer among the survivors. +// +// The two deletes are one sweeper on purpose. A connection row is only ever +// orphaned by its owner dying, so the moment that is decided is the moment to +// clean up after it; a second sweeper with its own schedule would either lag +// this one or race it, and would need its own answer to "is that replica +// alive", which is the one fact this table already owns. +// +// self is never reaped. This process may fail to heartbeat for longer than the +// window (a long stall, a database blip) and still be serving: deleting its own +// row would then delete the connections of workers that are, at that moment, +// connected to it. The stall is recovered by the re-register in tick instead. +// +// PostgreSQL only, like Live: distributed mode requires it, and the interval +// arithmetic is measured on the database's clock because liveness is compared +// across replicas. +func (r *Registry) ReapStale(ctx context.Context, self string, within time.Duration) (instances int64, connections int64, err error) { + err = r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + res := tx.Where("id <> ? AND last_seen <= now() - make_interval(secs => ?)", self, within.Seconds()). + Delete(&Instance{}) + if res.Error != nil { + return fmt.Errorf("deleting stale instances: %w", res.Error) + } + instances = res.RowsAffected + + // Whatever survived the delete above is the live set, so this needs no + // second liveness rule and cannot disagree with the first one. + res = tx.Where("owner_instance_id NOT IN (SELECT id FROM instances)"). + Delete(&NodeConnection{}) + if res.Error != nil { + return fmt.Errorf("deleting orphaned node connections: %w", res.Error) + } + connections = res.RowsAffected + return nil + }) + if err != nil { + return 0, 0, fmt.Errorf("reaping stale cluster state: %w", err) + } + return instances, connections, nil +} diff --git a/core/services/cluster/membership_test.go b/core/services/cluster/membership_test.go new file mode 100644 index 000000000..a00a08c0b --- /dev/null +++ b/core/services/cluster/membership_test.go @@ -0,0 +1,116 @@ +package cluster_test + +import ( + "context" + "time" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/testutil" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" +) + +var _ = Describe("Reaping dead replicas", func() { + var ( + db *gorm.DB + reg *cluster.Registry + ctx context.Context + ) + + // age pushes a replica's heartbeat into the past. Sleeping in a spec is + // forbidden, and the liveness window is measured in tens of seconds. + age := func(id string, by time.Duration) { + GinkgoHelper() + Expect(db.Model(&cluster.Instance{}).Where("id = ?", id). + Update("last_seen", gorm.Expr("now() - make_interval(secs => ?)", by.Seconds())).Error).To(Succeed()) + } + + BeforeEach(func() { + db = testutil.SetupTestDB() + ctx = context.Background() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + reg = cluster.NewRegistry(db) + }) + + It("deletes a replica that stopped heartbeating, and the connections it owned", func() { + Expect(reg.Register(ctx, "live", "10.0.0.1:8080", "v1")).To(Succeed()) + Expect(reg.Register(ctx, "dead", "10.0.0.2:8080", "v1")).To(Succeed()) + _, err := reg.Claim(ctx, "w1", "dead") + Expect(err).ToNot(HaveOccurred()) + age("dead", time.Hour) + + instances, connections, err := reg.ReapStale(ctx, "live", time.Minute) + Expect(err).ToNot(HaveOccurred()) + Expect(instances).To(Equal(int64(1))) + Expect(connections).To(Equal(int64(1)), + "a worker whose owner no longer exists is recorded as connected to nothing") + + _, _, err = reg.Owner(ctx, "w1") + Expect(err).To(MatchError(cluster.ErrNoConnection)) + }) + + It("leaves the connections of a live replica alone", func() { + Expect(reg.Register(ctx, "live", "10.0.0.1:8080", "v1")).To(Succeed()) + Expect(reg.Register(ctx, "other", "10.0.0.2:8080", "v1")).To(Succeed()) + epoch, err := reg.Claim(ctx, "w1", "other") + Expect(err).ToNot(HaveOccurred()) + + _, connections, err := reg.ReapStale(ctx, "live", time.Minute) + Expect(err).ToNot(HaveOccurred()) + Expect(connections).To(BeZero()) + + owner, stored, err := reg.Owner(ctx, "w1") + Expect(err).ToNot(HaveOccurred()) + Expect(owner).To(Equal("other")) + Expect(stored).To(Equal(epoch)) + }) + + It("never reaps the sweeper itself, however stale its own row looks", func() { + // A replica whose heartbeat stalled longer than the window is still + // serving the workers connected to it. Reaping its own row would delete + // their connection rows in the same pass, re-homing workers that never + // went anywhere. + Expect(reg.Register(ctx, "me", "10.0.0.1:8080", "v1")).To(Succeed()) + _, err := reg.Claim(ctx, "w1", "me") + Expect(err).ToNot(HaveOccurred()) + age("me", time.Hour) + + instances, connections, err := reg.ReapStale(ctx, "me", time.Minute) + Expect(err).ToNot(HaveOccurred()) + Expect(instances).To(BeZero()) + Expect(connections).To(BeZero()) + + owner, _, err := reg.Owner(ctx, "w1") + Expect(err).ToNot(HaveOccurred()) + Expect(owner).To(Equal("me")) + }) + + It("keeps this replica's row alive and reaps the dead while it runs", func() { + Expect(reg.Register(ctx, "dead", "10.0.0.2:8080", "v1")).To(Succeed()) + age("dead", time.Hour) + + membership := cluster.NewMembership(reg, "me", "10.0.0.1:8080", "v1") + Expect(membership.Start(ctx)).To(Succeed()) + DeferCleanup(membership.Stop) + + // Rows, not live rows: an aged-out replica drops out of Live + // immediately, and what the sweeper adds is deleting it. Asserting on + // Live here would pass with no sweeper at all. + rows := func() int64 { + var n int64 + if err := db.Model(&cluster.Instance{}).Count(&n).Error; err != nil { + return -1 + } + return n + } + Expect(rows()).To(Equal(int64(2)), "the stale row is still in the table until a sweep deletes it") + + Eventually(rows, 3*cluster.InstanceHeartbeat, time.Second).Should(Equal(int64(1))) + live, err := reg.Live(ctx, time.Minute) + Expect(err).ToNot(HaveOccurred()) + Expect(live).To(HaveLen(1)) + Expect(live[0].ID).To(Equal("me"), "the sweeper deleted the wrong row") + }) +}) diff --git a/core/services/cluster/ownership.go b/core/services/cluster/ownership.go index a2e2ba9cb..3b6b6cb62 100644 --- a/core/services/cluster/ownership.go +++ b/core/services/cluster/ownership.go @@ -59,17 +59,34 @@ type NodeConnection struct { ConnectedAt time.Time `gorm:"not null" json:"connected_at"` } -// EnsureEpochSequence creates the sequence Claim draws epochs from. It lives +// Migrate creates every table and sequence this package owns. It is the one +// call a caller has to remember: gorm's AutoMigrate models tables and columns +// but has no notion of a sequence, and the connection fence draws its epochs +// from one, so a caller that knew only about AutoMigrate would leave a schema +// that looks complete and cannot claim. Safe to call repeatedly. +// +// It does not take the migration advisory lock itself. The caller holds it +// across every table in the deployment, and taking a second one here would +// either nest inside that one or, worse, be the reason someone stops holding +// the outer one. +func Migrate(ctx context.Context, db *gorm.DB) error { + if err := db.WithContext(ctx).AutoMigrate(&Instance{}, &NodeConnection{}); err != nil { + return fmt.Errorf("migrating cluster tables: %w", err) + } + return ensureEpochSequence(ctx, db) +} + +// ensureEpochSequence creates the sequence Claim draws epochs from. It lives // here, beside the model that needs it, because gorm's AutoMigrate models // tables and columns but has no notion of a sequence; the caller that owns the -// migration advisory lock calls it so that concurrently starting replicas do -// not race on the DDL. It is safe to call repeatedly. +// migration advisory lock calls Migrate so that concurrently starting replicas +// do not race on the DDL. It is safe to call repeatedly. // // The sequence is not attached as a column DEFAULT on purpose: AutoMigrate // compares the struct's declared default against the one PostgreSQL reports // (`nextval('...'::regclass)`), and a mismatch there makes every startup ALTER // the column. Naming the sequence in the statement keeps the schema stable. -func EnsureEpochSequence(ctx context.Context, db *gorm.DB) error { +func ensureEpochSequence(ctx context.Context, db *gorm.DB) error { // CREATE SEQUENCE is PostgreSQL-only, and the same migration path runs // against SQLite in single-binary mode. Nothing there can claim a // connection (Claim refuses the dialect outright), so there is nothing to diff --git a/core/services/cluster/ownership_test.go b/core/services/cluster/ownership_test.go index ba634ebcd..5744e07cd 100644 --- a/core/services/cluster/ownership_test.go +++ b/core/services/cluster/ownership_test.go @@ -67,8 +67,7 @@ var _ = Describe("Connection ownership", func() { BeforeEach(func() { db = testutil.SetupTestDB() ctx = context.Background() - Expect(db.AutoMigrate(&cluster.NodeConnection{})).To(Succeed()) - Expect(cluster.EnsureEpochSequence(ctx, db)).To(Succeed()) + Expect(cluster.Migrate(ctx, db)).To(Succeed()) reg = cluster.NewRegistry(db) }) @@ -243,13 +242,11 @@ var _ = Describe("Connection ownership on a non-PostgreSQL dialect", func() { It("migrates, because the single-binary path shares this schema", func() { // A PostgreSQL-only column DEFAULT here breaks AutoMigrate for every // SQLite caller of nodes.NewNodeRegistry, which is how this regressed. - Expect(db.AutoMigrate(&cluster.NodeConnection{})).To(Succeed()) - Expect(cluster.EnsureEpochSequence(ctx, db)).To(Succeed()) + Expect(cluster.Migrate(ctx, db)).To(Succeed()) }) It("refuses to claim, rather than pretending to fence", func() { - Expect(db.AutoMigrate(&cluster.NodeConnection{})).To(Succeed()) - Expect(cluster.EnsureEpochSequence(ctx, db)).To(Succeed()) + Expect(cluster.Migrate(ctx, db)).To(Succeed()) _, err := cluster.NewRegistry(db).Claim(ctx, "w1", "inst-a") Expect(err).To(HaveOccurred()) diff --git a/core/services/cluster/peerlink_test.go b/core/services/cluster/peerlink_test.go index cfaa4dc69..692f2ee6c 100644 --- a/core/services/cluster/peerlink_test.go +++ b/core/services/cluster/peerlink_test.go @@ -20,6 +20,16 @@ import ( "gorm.io/gorm" ) +// servePeerRoute mounts the peer handler on the route both sides agree on. +// +// It deliberately does not call routes.RegisterClusterRoutes: that registrar +// lives in core/http/routes, which imports half the server, and these specs are +// about the handler and the dialler rather than about the route table. The path +// comes from the same constant the registrar uses, so the two cannot drift. +func servePeerRoute(e *echo.Echo, token string, onPeer func(string, *yamux.Session)) { + e.GET(cluster.PeerPath, clusterep.PeerHandler(token, onPeer)) +} + var _ = Describe("Peer pool", func() { var ( db *gorm.DB @@ -33,7 +43,7 @@ var _ = Describe("Peer pool", func() { // startPeer stands up a real peer server and registers it under peerID. startPeer := func(peerID string) *httptest.Server { e := echo.New() - clusterep.RegisterClusterRoutes(e, "peer-token", func(_ string, s *yamux.Session) { + servePeerRoute(e, "peer-token", func(_ string, s *yamux.Session) { accepted <- s }) ts := httptest.NewServer(e) @@ -45,7 +55,7 @@ var _ = Describe("Peer pool", func() { BeforeEach(func() { ctx = context.Background() db = testutil.SetupTestDB() - Expect(db.AutoMigrate(&cluster.Instance{})).To(Succeed()) + Expect(cluster.Migrate(ctx, db)).To(Succeed()) reg = cluster.NewRegistry(db) accepted = make(chan *yamux.Session, 4) pool = cluster.NewPeerPool("self", "peer-token", reg) @@ -88,7 +98,7 @@ var _ = Describe("Peer pool", func() { // link anonymous and indistinguishable from every other. ids := make(chan string, 1) e := echo.New() - clusterep.RegisterClusterRoutes(e, "peer-token", func(id string, _ *yamux.Session) { ids <- id }) + servePeerRoute(e, "peer-token", func(id string, _ *yamux.Session) { ids <- id }) ts := httptest.NewServer(e) DeferCleanup(ts.Close) Expect(reg.Register(ctx, "peer-named", strings.TrimPrefix(ts.URL, "http://"), "test")).To(Succeed()) @@ -131,7 +141,7 @@ var _ = Describe("Peer pool", func() { // row. Reporting absence here would evict every worker behind a peer // that was merely rolled out with a stale secret. e := echo.New() - clusterep.RegisterClusterRoutes(e, "a-different-token", func(_ string, s *yamux.Session) { accepted <- s }) + servePeerRoute(e, "a-different-token", func(_ string, s *yamux.Session) { accepted <- s }) ts := httptest.NewServer(e) DeferCleanup(ts.Close) Expect(reg.Register(ctx, "peer-strict", strings.TrimPrefix(ts.URL, "http://"), "test")).To(Succeed()) diff --git a/core/services/cluster/sessions.go b/core/services/cluster/sessions.go new file mode 100644 index 000000000..7ee5fa332 --- /dev/null +++ b/core/services/cluster/sessions.go @@ -0,0 +1,140 @@ +// SPDX-License-Identifier: MIT + +package cluster + +import ( + "net" + "sync" + + "github.com/libp2p/go-yamux/v5" + "github.com/mudler/xlog" +) + +// SessionStore holds the peer links this replica has ACCEPTED, which is the +// mirror image of PeerPool: the pool owns the sessions this replica dialled, +// this owns the ones its peers dialled into it. +// +// Something has to own an accepted session. The HTTP handler cannot: it returns +// as soon as the upgrade is done, and the hijacked connection outlives it. And +// something has to accept the streams that arrive on it, because yamux only +// acknowledges a stream once the far side accepts it, so a session nobody +// accepts on does not fail a peer's Open, it hangs it. +type SessionStore struct { + // onStream handles one accepted stream and owns closing it. A nil handler + // closes the stream immediately, which is what a replica with no relay + // installed should do: refuse promptly rather than leave a peer parked. + onStream func(peerID string, stream net.Conn) + + mu sync.Mutex + sessions map[string]*yamux.Session + closed bool +} + +// NewSessionStore returns a store whose accepted streams are handled by +// onStream. Pass nil to refuse every stream, closing it at once. +func NewSessionStore(onStream func(peerID string, stream net.Conn)) *SessionStore { + return &SessionStore{onStream: onStream, sessions: map[string]*yamux.Session{}} +} + +// Accept takes ownership of a session a peer dialled in. It is the callback +// shape RegisterClusterRoutes wants, and it returns promptly: the serving loop +// runs on its own goroutine, because the handler's return is what completes the +// hijack. +func (s *SessionStore) Accept(peerID string, sess *yamux.Session) { + if sess == nil { + return + } + + s.mu.Lock() + if s.closed { + s.mu.Unlock() + // Shutdown raced the dial. Leaving the session open would keep the peer + // believing it has a live link into a process that is going away. + _ = sess.Close() + return + } + previous := s.sessions[peerID] + s.sessions[peerID] = sess + s.mu.Unlock() + + // A peer that dials again has lost its previous link, whether or not this + // side has noticed. Keeping both would leave a session nothing can ever be + // routed to, since the map holds one per peer. + if previous != nil { + xlog.Debug("cluster peer re-dialled, dropping its previous link", "peer", peerID) + _ = previous.Close() + } + + go s.serve(peerID, sess) +} + +// Get returns the session this replica accepted from peerID. The second result +// is false when no link from that peer is held, which a caller must not read as +// the peer being absent: it may be about to dial, or dialling this replica may +// simply not be its job. +func (s *SessionStore) Get(peerID string) (*yamux.Session, bool) { + s.mu.Lock() + defer s.mu.Unlock() + sess, ok := s.sessions[peerID] + return sess, ok +} + +// serve accepts streams until the session dies, then forgets it. +func (s *SessionStore) serve(peerID string, sess *yamux.Session) { + defer func() { + s.forget(peerID, sess) + _ = sess.Close() + }() + + for { + stream, err := sess.AcceptStream() + if err != nil { + // A peer link ending is ordinary: a rolling update closes every + // session it holds. The error is the session's, not one stream's, + // so there is nothing to recover to. + xlog.Debug("cluster peer link ended", "peer", peerID, "error", err) + return + } + if s.onStream == nil { + // No relay installed. Closing is deliberate and is not the same as + // ignoring: a stream nobody answers parks the peer's request until + // its own deadline, and reports nothing about why. + xlog.Debug("cluster peer stream refused: no relay installed", "peer", peerID) + _ = stream.Close() + continue + } + // One goroutine per stream: the handler relays a whole request, and + // serving them from the accept loop would let one request stall every + // other stream on the link. + go s.onStream(peerID, stream) + } +} + +// forget drops the entry only if it still names this session. A peer that +// re-dialled has already replaced it, and deleting blindly would evict the live +// link when the old one finally noticed it was dead. +func (s *SessionStore) forget(peerID string, sess *yamux.Session) { + s.mu.Lock() + defer s.mu.Unlock() + if s.sessions[peerID] == sess { + delete(s.sessions, peerID) + } +} + +// CloseAll drops every held session. An Accept after it closes the session +// rather than storing it, so a dial racing shutdown cannot leak a link. +func (s *SessionStore) CloseAll() { + s.mu.Lock() + if s.closed { + s.mu.Unlock() + return + } + s.closed = true + held := s.sessions + s.sessions = map[string]*yamux.Session{} + s.mu.Unlock() + + for _, sess := range held { + _ = sess.Close() + } +} diff --git a/core/services/cluster/sessions_test.go b/core/services/cluster/sessions_test.go new file mode 100644 index 000000000..f3fb44ca8 --- /dev/null +++ b/core/services/cluster/sessions_test.go @@ -0,0 +1,122 @@ +package cluster_test + +import ( + "net" + "time" + + "github.com/mudler/LocalAI/core/services/cluster" + + "github.com/libp2p/go-yamux/v5" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// yamuxPair returns a client and a server session over an in-memory pipe. It +// stands in for a dialled peer link: everything the store does with a session +// is transport-agnostic, and the WebSocket half is covered where it is used. +func yamuxPair() (client *yamux.Session, server *yamux.Session) { + GinkgoHelper() + a, b := net.Pipe() + var err error + server, err = yamux.Server(a, nil, nil) + Expect(err).ToNot(HaveOccurred()) + client, err = yamux.Client(b, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { + _ = client.Close() + _ = server.Close() + }) + return client, server +} + +var _ = Describe("Accepted peer sessions", func() { + It("accepts and refuses a stream rather than leaving the peer parked", func() { + // yamux only acknowledges a stream once the far side accepts it, so a + // store that held the session without accepting would not fail a peer's + // Open, it would hang it, and every relayed request behind it. + store := cluster.NewSessionStore(nil) + DeferCleanup(store.CloseAll) + client, server := yamuxPair() + store.Accept("peer-1", server) + + stream, err := client.OpenStream(GinkgoT().Context()) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = stream.Close() }) + + Expect(stream.SetReadDeadline(time.Now().Add(10 * time.Second))).To(Succeed()) + _, err = stream.Read(make([]byte, 1)) + Expect(err).To(HaveOccurred(), "a refused stream must end, not hang") + }) + + It("hands a stream to the relay when one is installed", func() { + streams := make(chan net.Conn, 1) + store := cluster.NewSessionStore(func(_ string, stream net.Conn) { streams <- stream }) + DeferCleanup(store.CloseAll) + client, server := yamuxPair() + store.Accept("peer-1", server) + + stream, err := client.OpenStream(GinkgoT().Context()) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = stream.Close() }) + + var relayed net.Conn + Eventually(streams, "10s").Should(Receive(&relayed)) + go func() { + defer GinkgoRecover() + _, _ = stream.Write([]byte("hello")) + }() + buf := make([]byte, 5) + Expect(relayed.SetReadDeadline(time.Now().Add(10 * time.Second))).To(Succeed()) + _, err = relayed.Read(buf) + Expect(err).ToNot(HaveOccurred()) + Expect(string(buf)).To(Equal("hello")) + }) + + It("replaces a peer's link when it dials again, and closes the one it lost", func() { + // A peer only re-dials because its previous link is gone from where it + // stands. Keeping both would leave a session nothing can be routed to, + // since the store holds one per peer. + store := cluster.NewSessionStore(nil) + DeferCleanup(store.CloseAll) + _, first := yamuxPair() + _, second := yamuxPair() + + store.Accept("peer-1", first) + store.Accept("peer-1", second) + + held, ok := store.Get("peer-1") + Expect(ok).To(BeTrue()) + Expect(held).To(BeIdenticalTo(second)) + Eventually(first.IsClosed, "10s").Should(BeTrue()) + Expect(second.IsClosed()).To(BeFalse(), "the link the peer is actually using was dropped") + }) + + It("forgets a session that ended, without evicting the one that replaced it", func() { + store := cluster.NewSessionStore(nil) + DeferCleanup(store.CloseAll) + client, server := yamuxPair() + store.Accept("peer-1", server) + + Expect(client.Close()).To(Succeed()) + Eventually(func() bool { + _, ok := store.Get("peer-1") + return ok + }, "10s").Should(BeFalse()) + }) + + It("closes every held link on shutdown, and refuses to store one afterwards", func() { + store := cluster.NewSessionStore(nil) + _, server := yamuxPair() + store.Accept("peer-1", server) + + store.CloseAll() + Eventually(server.IsClosed, "10s").Should(BeTrue()) + + _, late := yamuxPair() + store.Accept("peer-late", late) + _, ok := store.Get("peer-late") + Expect(ok).To(BeFalse()) + Eventually(late.IsClosed, "10s").Should(BeTrue(), + "a dial racing shutdown must not be left believing it holds a live link") + }) +}) diff --git a/core/services/cluster/wsconn_test.go b/core/services/cluster/wsconn_test.go index 305ce99b3..d7213fd87 100644 --- a/core/services/cluster/wsconn_test.go +++ b/core/services/cluster/wsconn_test.go @@ -11,7 +11,6 @@ import ( "strings" "time" - clusterep "github.com/mudler/LocalAI/core/http/endpoints/cluster" "github.com/mudler/LocalAI/core/services/cluster" "github.com/gorilla/websocket" @@ -239,7 +238,7 @@ var _ = Describe("Peer link payloads", func() { It("carries a payload far larger than one yamux frame end to end", func() { sessions := make(chan *yamux.Session, 1) e := echo.New() - clusterep.RegisterClusterRoutes(e, "peer-token", func(_ string, s *yamux.Session) { sessions <- s }) + servePeerRoute(e, "peer-token", func(_ string, s *yamux.Session) { sessions <- s }) srv := httptest.NewServer(e) DeferCleanup(srv.Close) diff --git a/core/services/nodes/registry.go b/core/services/nodes/registry.go index 72b957bcc..e147a9955 100644 --- a/core/services/nodes/registry.go +++ b/core/services/nodes/registry.go @@ -443,14 +443,14 @@ func (r *NodeRegistry) nodeModelNames(ctx context.Context, db *gorm.DB, nodeID s // when multiple instances (frontend + workers) start at the same time. func NewNodeRegistry(db *gorm.DB) (*NodeRegistry, error) { if err := advisorylock.WithLockCtx(context.Background(), db, advisorylock.KeySchemaMigrate, func() error { - if err := db.AutoMigrate(&BackendNode{}, &NodeModel{}, &NodeLabel{}, &ModelSchedulingConfig{}, &PendingBackendOp{}, &ModelLoadInfo{}, &ModelLoadJob{}, &ModelConfigState{}, &cluster.Instance{}, &cluster.NodeConnection{}); err != nil { + if err := db.AutoMigrate(&BackendNode{}, &NodeModel{}, &NodeLabel{}, &ModelSchedulingConfig{}, &PendingBackendOp{}, &ModelLoadInfo{}, &ModelLoadJob{}, &ModelConfigState{}); err != nil { return err } - // AutoMigrate models tables and columns but has no notion of a - // sequence, and the connection-ownership fence draws its epochs from - // one. It runs under this same lock so concurrently starting replicas - // do not race on the DDL. - return cluster.EnsureEpochSequence(context.Background(), db) + // The cluster package owns its own tables AND the sequence its + // ownership fence draws epochs from, which AutoMigrate cannot express. + // It runs under this same lock so concurrently starting replicas do not + // race on the DDL. + return cluster.Migrate(context.Background(), db) }); err != nil { return nil, fmt.Errorf("migrating node tables: %w", err) } diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index 0231c2dc4..3a6f9768f 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -64,6 +64,7 @@ The frontend is a standard LocalAI instance with distributed mode enabled. These | `--distributed` | `LOCALAI_DISTRIBUTED` | `false` | Enable distributed mode | | `--instance-id` | `LOCALAI_INSTANCE_ID` | auto UUID | Unique instance ID for this frontend | | `--nats-url` | `LOCALAI_NATS_URL` | *(required)* | NATS server URL (e.g., `nats://localhost:4222`) | +| `--distributed-advertise-addr` | `LOCALAI_DISTRIBUTED_ADVERTISE_ADDR` | *(derived)* | `host:port` the **other frontend replicas** dial to reach this one. See [Replica peer links](#replica-peer-links). | | `--registration-token` | `LOCALAI_REGISTRATION_TOKEN` | *(empty)* | Token that workers must provide to register | | `--registration-require-auth` | `LOCALAI_REGISTRATION_REQUIRE_AUTH` | `false` | Fail startup when distributed mode is enabled but the registration token is empty (node endpoints and worker file-transfer would otherwise be unauthenticated) | | `--distributed-require-auth` | `LOCALAI_DISTRIBUTED_REQUIRE_AUTH` | `false` | **Umbrella switch.** Implies both `--nats-require-auth` and `--registration-require-auth` - one knob to lock down the NATS bus *and* the registration/file-transfer layer. Set this in production instead of the two granular flags. | @@ -78,6 +79,27 @@ The frontend is a standard LocalAI instance with distributed mode enabled. These | *(env only)* | `LOCALAI_MODEL_LOAD_WAIT` | `60s` | How long an inference request waits for a model that is still cold-loading onto a worker before it is answered with `503`, a `Retry-After` header and live staging progress. The request is served the moment the model becomes ready, so a model already most of the way staged needs no client retry. Set to `0` to wait as long as the load takes — only safe when no ingress or load balancer with an idle timeout sits in front. See [Requests for a model that is still loading](#requests-for-a-model-that-is-still-loading). | | `--expose-node-header` | `LOCALAI_EXPOSE_NODE_HEADER` | `false` | When enabled, inference responses carry an `X-LocalAI-Node` header with the ID of the worker node that served the request. Coverage spans the OpenAI-compatible endpoints (chat completions, completions, embeddings, audio transcriptions, audio speech / TTS, image generations, image inpainting), the Jina rerank endpoint (`/v1/rerank`), the VAD endpoints (`/v1/vad`, `/vad`), and the Anthropic Messages (`/v1/messages`) and Ollama (`/api/chat`, `/api/generate`, `/api/embed`) shims. Useful for debugging, observability and load-balancer attribution. Off by default: the node ID reveals internal cluster topology and should not be exposed on a public endpoint. Best-effort: under heavy concurrency for the same model across multiple replicas, the header may reflect a recent routing decision rather than this exact request's. Acceptable for observability and debugging. | +### Replica peer links + +Frontend replicas record themselves in an `instances` table and open direct links to each other, so that a request arriving at one replica can be served by state another replica holds. Each replica publishes one address for this, and every other replica dials it: it is the address **peers** use, which is not necessarily the address the process binds. A replica behind a Kubernetes Service, a load balancer or a NAT binds one and is reached at another. + +When `LOCALAI_DISTRIBUTED_ADVERTISE_ADDR` is unset, the address is derived: LocalAI asks the kernel which local address routes to PostgreSQL, and pairs it with the port it serves on. Every replica reaches the same database, so that address is on a network they demonstrably share. + +That only holds while the database is on **another host**. If PostgreSQL runs on the same host or pod (compose, single-node, a sidecar), the route to it is loopback, and advertising a loopback address would send every peer to itself. LocalAI refuses to guess in that case and logs: + +``` +This replica will not be reachable by its peers: no advertised address +``` + +The replica keeps serving every request that reaches it directly; what it cannot do is have another replica reach it. Set the address explicitly to fix it: + +```yaml +environment: + LOCALAI_DISTRIBUTED_ADVERTISE_ADDR: "10.0.1.7:8080" # or the pod IP, service DNS name, etc. +``` + +The peer link is served at `/api/cluster/peer` and authenticates with `LOCALAI_REGISTRATION_TOKEN`, the same shared secret workers register with. Replicas that disagree about it cannot link. A replica that stops heartbeating for 30 seconds is dropped from the table by the others, along with the worker-connection rows it owned. + ### The model load deadline scales with the checkpoint The `LoadModel` deadline starts *after* the backend is installed and the model files are staged, so it covers only the worker backend's own checkpoint read and pipeline init. That work is proportional to the bytes on disk, which makes any fixed deadline a model-size cliff rather than a timeout: a 70 GB video checkpoint on a Jetson Thor worker failed reproducibly against the old fixed 5m default (`rpc error: code = DeadlineExceeded` after 953.5s of wall clock, roughly 11m of which was backend install and staging), and simply raising the constant would only move the cliff to the next larger model while making a genuinely wedged *small* model hang for the whole inflated duration. diff --git a/tests/e2e/distributed/cluster/cluster.go b/tests/e2e/distributed/cluster/cluster.go index 1783d71e1..521882efe 100644 --- a/tests/e2e/distributed/cluster/cluster.go +++ b/tests/e2e/distributed/cluster/cluster.go @@ -209,6 +209,14 @@ func (c *Cluster) startFrontend(i int, port int) (*Process, error) { // Pinning makes the cross-replica session a property of the harness. "LOCALAI_AUTH_HMAC_SECRET="+testHMACSecret, "LOCALAI_REGISTRATION_TOKEN="+c.opts.RegistrationToken, + // Every replica here shares one host, so the address a peer dials is + // this process's own loopback address. It has to be said explicitly: + // the automatic discovery asks which local address routes to + // PostgreSQL, and this suite's PostgreSQL is a container published on + // 127.0.0.1, so the discovery refuses (correctly) rather than + // advertising a loopback address that would mean "yourself" on a + // multi-host deployment. + fmt.Sprintf("LOCALAI_DISTRIBUTED_ADVERTISE_ADDR=127.0.0.1:%d", port), "LOCALAI_AUTO_APPROVE_NODES=true", "DEBUG=true", ) @@ -334,6 +342,14 @@ func (c *Cluster) FrontendURL(i int) string { return fmt.Sprintf("http://127.0.0.1:%d", c.frontends[i].Port) } +// RegistrationToken is the shared secret this cluster was started with. It +// authenticates worker registration AND the replica-to-replica peer link, so a +// spec acting as a peer needs it rather than a second literal that can drift +// from Options. +func (c *Cluster) RegistrationToken() string { + return c.opts.RegistrationToken +} + // WorkerName is the node name worker i registered under. func (c *Cluster) WorkerName(i int) string { return c.workers[i].Name diff --git a/tests/e2e/distributed/cluster_baseline_test.go b/tests/e2e/distributed/cluster_baseline_test.go index 6cbaaa0ea..63995b419 100644 --- a/tests/e2e/distributed/cluster_baseline_test.go +++ b/tests/e2e/distributed/cluster_baseline_test.go @@ -122,6 +122,14 @@ func mockBackendBinary() string { // existing caller keeps the plain two-argument form and the default shape. func startCluster(frontends, workers int, customise ...func(*cluster.Options)) *cluster.Cluster { GinkgoHelper() + c, _ := startClusterOnFreshDB(frontends, workers, customise...) + return c +} + +// startClusterOnFreshDB is startCluster plus the DSN of the database it was +// given, for a spec that has to read a table no endpoint exposes. +func startClusterOnFreshDB(frontends, workers int, customise ...func(*cluster.Options)) (*cluster.Cluster, string) { + GinkgoHelper() // Resolved before SetupInfra so a missing binary skips without having paid // for a database that the skip would then leave to DeferCleanup. @@ -163,7 +171,7 @@ func startCluster(frontends, workers int, customise ...func(*cluster.Options)) * } c.Stop() }) - return c + return c, infra.PGURL } // rosterProbe polls one frontend's node roster. diff --git a/tests/e2e/distributed/cluster_peerlink_test.go b/tests/e2e/distributed/cluster_peerlink_test.go new file mode 100644 index 000000000..8055ba297 --- /dev/null +++ b/tests/e2e/distributed/cluster_peerlink_test.go @@ -0,0 +1,263 @@ +package distributed_test + +import ( + "context" + "fmt" + "net" + "strings" + "time" + + clustersvc "github.com/mudler/LocalAI/core/services/cluster" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/driver/postgres" + "gorm.io/gorm" + gormlogger "gorm.io/gorm/logger" +) + +const ( + // instanceRosterTimeout bounds the wait for a replica's row to appear. + // Registration is synchronous in startup, so this only has to cover the gap + // between /readyz answering and this spec's first query. + instanceRosterTimeout = "30s" + instanceRosterPoll = "500ms" + + // deadReplicaTimeout bounds the wait for a survivor to reap a replica that + // was killed: the liveness window plus a sweep interval plus slack. It is + // deliberately derived from the constants rather than a round number, so + // tightening the window shortens the spec instead of leaving it passing for + // the wrong reason. + deadReplicaTimeout = clustersvc.InstanceLiveness + 4*clustersvc.InstanceHeartbeat + + // peerDialTimeout bounds one peer dial. Every replica here is a local + // process, so a dial that needs longer has failed, not slowed. + peerDialTimeout = 20 * time.Second +) + +// openClusterDB connects to the database the cluster was given, so a spec can +// read the tables the peer link keeps. Nothing serves them over HTTP: they are +// replica-to-replica state, not an admin surface, and inventing an endpoint to +// observe them would be a bigger change than the thing under test. +func openClusterDB(dsn string) *gorm.DB { + GinkgoHelper() + db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{Logger: gormlogger.Discard}) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { closeDB(db) }) + return db +} + +// hostPortOf strips the scheme off a frontend URL, giving the form the +// instances table stores. +func hostPortOf(url string) string { + return strings.TrimPrefix(strings.TrimPrefix(url, "http://"), "https://") +} + +// instanceRoster reads the live replica rows, keeping the last error so a +// failing Eventually can name it. +type instanceRoster struct { + registry *clustersvc.Registry + ctx context.Context + + lastErr error + lastSaw []clustersvc.Instance +} + +func newInstanceRoster(db *gorm.DB) *instanceRoster { + return &instanceRoster{registry: clustersvc.NewRegistry(db), ctx: context.Background()} +} + +// addresses returns the advertised address of every live replica, or nil on a +// query error so Eventually keeps trying. +func (r *instanceRoster) addresses() []string { + live, err := r.registry.Live(r.ctx, clustersvc.InstanceLiveness) + if err != nil { + r.lastErr = err + return nil + } + r.lastErr = nil + r.lastSaw = live + addrs := []string{} + for _, instance := range live { + addrs = append(addrs, instance.AdvertisedAddr) + } + return addrs +} + +// idAt returns the id of the live replica advertising addr, or "" if no such +// row is present yet. +func (r *instanceRoster) idAt(addr string) string { + for _, instance := range r.lastSaw { + if instance.AdvertisedAddr == addr { + return instance.ID + } + } + return "" +} + +func (r *instanceRoster) describe() string { + if r.lastErr != nil { + return fmt.Sprintf("the last read of the instances table failed: %v", r.lastErr) + } + return fmt.Sprintf("the instances table held %d live replica(s): %+v", len(r.lastSaw), r.lastSaw) +} + +// awaitReplicas waits for every frontend of c to publish its address and +// returns the roster, positioned on that reading. +func awaitReplicas(roster *instanceRoster, addrs ...string) { + GinkgoHelper() + Eventually(roster.addresses, instanceRosterTimeout, instanceRosterPoll). + Should(ConsistOf(addrs), roster.describe) +} + +var _ = Describe("Cluster peer link", Label("Distributed"), Label("Cluster"), func() { + It("publishes an address for every replica that peers can actually dial", func() { + // A wrong implementation registers nothing (the whole of phase 1 had no + // call site until this spec), registers one row for two replicas, or + // records an address nothing can connect to: the bind address of a + // replica behind a service, or the loopback address the route to a + // co-located database would suggest. + c, dsn := startClusterOnFreshDB(2, 0) + + roster := newInstanceRoster(openClusterDB(dsn)) + awaitReplicas(roster, hostPortOf(c.FrontendURL(0)), hostPortOf(c.FrontendURL(1))) + + // "Routable" is not a property of the string. Connect to each address, + // which is the only check that would have caught a replica publishing + // the port it was configured with rather than the one it serves on. + for _, instance := range roster.lastSaw { + conn, err := net.DialTimeout("tcp", instance.AdvertisedAddr, peerDialTimeout) + Expect(err).ToNot(HaveOccurred(), + "replica %s advertises %q, which nothing can connect to", instance.ID, instance.AdvertisedAddr) + Expect(conn.Close()).To(Succeed()) + } + }) + + It("carries a peer stream between two replicas, and refuses one without the cluster token", func() { + // A wrong implementation fails here on WebSocket framing, which is the + // likeliest defect in the peer link: the adapter has to turn + // message-oriented WebSocket frames into the undelimited byte stream + // yamux drives. It also fails if the route was never registered on the + // real server, or if the global session middleware answers it: a peer + // carries no session and no user, only the cluster token. + // + // The stream is opened with the production dialler, resolving the peer + // through the production registry, over a real socket to a real + // process. This spec plays the sibling replica, because phase 1 has + // nothing that makes a frontend dial one on its own. + c, dsn := startClusterOnFreshDB(2, 0) + + roster := newInstanceRoster(openClusterDB(dsn)) + awaitReplicas(roster, hostPortOf(c.FrontendURL(0)), hostPortOf(c.FrontendURL(1))) + + peerID := roster.idAt(hostPortOf(c.FrontendURL(1))) + Expect(peerID).ToNot(BeEmpty()) + + ctx, cancel := context.WithTimeout(context.Background(), peerDialTimeout) + defer cancel() + + pool := clustersvc.NewPeerPool("e2e-peer", c.RegistrationToken(), roster.registry) + DeferCleanup(pool.Close) + + // OpenStream is only acknowledged once the far side accepts, so this + // returning at all proves the frontend is accepting streams on the + // session it took, in addition to proving the handshake. + stream, err := pool.Open(ctx, peerID) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = stream.Close() }) + + // Phase 1 installs no relay, so the accepted stream is refused at once. + // What matters here is that the refusal arrives: a replica that held + // the session without accepting on it would leave this read parked + // until the deadline, which is the failure mode a live cluster would + // experience as every relayed request hanging. + Expect(stream.SetReadDeadline(time.Now().Add(peerDialTimeout))).To(Succeed()) + _, err = stream.Read(make([]byte, 1)) + Expect(err).To(HaveOccurred(), + "the peer accepted the stream and then neither answered nor closed it") + Expect(err).ToNot(MatchError(context.DeadlineExceeded)) + + // The same dial with the wrong credentials must be refused, otherwise + // the success above says nothing about authentication. + impostor := clustersvc.NewPeerPool("e2e-peer", "not-the-cluster-token", roster.registry) + DeferCleanup(impostor.Close) + _, err = impostor.Open(ctx, peerID) + Expect(err).To(MatchError(clustersvc.ErrPeerUnreachable)) + Expect(err).ToNot(MatchError(clustersvc.ErrInstanceNotFound), + "a peer refusing credentials is a live peer; reading it as absence is how a replica evicts healthy workers") + }) + + It("reports a killed replica as unreachable, reaps what it owned, and evicts no worker", func() { + // This is the absence rule, pinned before phase 2 can depend on it. A + // wrong implementation lets a peer that will not answer surface as node + // absence, and a caller entitled to act on absence then reclaims what + // the peer was running: a network hiccup between two healthy replicas + // evicts healthy workers. + // + // It also pins the reaper: the connection rows a dead replica owned are + // swept by the same sweeper that decides the replica is dead, so the + // two can never disagree about who is alive. + c, dsn := startClusterOnFreshDB(2, 1) + + client, err := c.AdminSession(0) + Expect(err).ToNot(HaveOccurred()) + + // The worker registers with frontend 0, so frontend 1 is the replica + // that can die without taking the worker's registrar with it. + registrar, err := c.WorkerRegistrar(0) + Expect(err).ToNot(HaveOccurred()) + Expect(registrar).To(Equal(0), "this spec kills frontend 1 and needs the worker to have registered elsewhere") + + probe := newRosterProbe(c, client, 0) + Eventually(probe.healthyNames, nodeRosterTimeout, nodeRosterPoll). + Should(ContainElement(c.WorkerName(0)), probe.describe) + workerID := probe.idOf(c.WorkerName(0)) + Expect(workerID).ToNot(BeEmpty()) + + roster := newInstanceRoster(openClusterDB(dsn)) + awaitReplicas(roster, hostPortOf(c.FrontendURL(0)), hostPortOf(c.FrontendURL(1))) + survivorID := roster.idAt(hostPortOf(c.FrontendURL(0))) + doomedID := roster.idAt(hostPortOf(c.FrontendURL(1))) + Expect(survivorID).ToNot(BeEmpty()) + Expect(doomedID).ToNot(BeEmpty()) + + // Give frontend 1 the worker's tunnel. Phase 2 makes the worker do this + // by dialling; here the claim is written directly, because the point + // under test is what happens to the claim when its owner dies. + ctx := context.Background() + epoch, err := roster.registry.Claim(ctx, workerID, doomedID) + Expect(err).ToNot(HaveOccurred()) + Expect(epoch).ToNot(BeZero()) + + Expect(c.KillFrontend(1)).To(Succeed()) + Eventually(func() bool { return c.FrontendAlive(1) }, "20s", "500ms").Should(BeFalse()) + + // The row is still there for the whole liveness window, so this is the + // case that matters: the peer is KNOWN and will not answer. + dialCtx, cancel := context.WithTimeout(ctx, peerDialTimeout) + defer cancel() + pool := clustersvc.NewPeerPool("e2e-peer", c.RegistrationToken(), roster.registry) + DeferCleanup(pool.Close) + _, err = pool.Open(dialCtx, doomedID) + Expect(err).To(MatchError(clustersvc.ErrPeerUnreachable)) + Expect(err).ToNot(MatchError(clustersvc.ErrInstanceNotFound), + "a dead replica whose row is still present is unreachable, not absent") + + // The survivor sweeps the dead replica and, in the same pass, the claim + // it left behind. + Eventually(roster.addresses, deadReplicaTimeout, instanceRosterPoll). + Should(ConsistOf(hostPortOf(c.FrontendURL(0))), roster.describe) + ownerErr := func() error { + _, _, err := roster.registry.Owner(ctx, workerID) + return err + } + Eventually(ownerErr, deadReplicaTimeout, instanceRosterPoll). + Should(MatchError(clustersvc.ErrNoConnection), + "the claim held by a replica that no longer exists was never reaped") + + // And the worker itself is untouched throughout. Nothing about a peer + // dying may reach the node roster. + Consistently(probe.healthyNames, "6s", "1s"). + Should(ContainElement(c.WorkerName(0)), probe.describe) + }) +})