From a3054442e2f9fb937d44648a52cf4725ea2625c8 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Mon, 31 Aug 2026 21:56:27 +0000 Subject: [PATCH] feat(cluster): record frontend replicas in a shared instances table Replicas need to find each other to relay worker traffic, and nothing in the tree recorded a replica's address. The advertised address is discovered by opening a UDP socket toward PostgreSQL and reading back the local address, which yields the interface every replica demonstrably shares without asking an operator to configure one. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/services/cluster/cluster_suite_test.go | 13 ++ core/services/cluster/instance.go | 189 ++++++++++++++++++++ core/services/cluster/instance_test.go | 101 +++++++++++ core/services/nodes/registry.go | 3 +- 4 files changed, 305 insertions(+), 1 deletion(-) create mode 100644 core/services/cluster/cluster_suite_test.go create mode 100644 core/services/cluster/instance.go create mode 100644 core/services/cluster/instance_test.go diff --git a/core/services/cluster/cluster_suite_test.go b/core/services/cluster/cluster_suite_test.go new file mode 100644 index 000000000..d82148729 --- /dev/null +++ b/core/services/cluster/cluster_suite_test.go @@ -0,0 +1,13 @@ +package cluster_test + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestCluster(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Cluster Package Suite") +} diff --git a/core/services/cluster/instance.go b/core/services/cluster/instance.go new file mode 100644 index 000000000..7d4347204 --- /dev/null +++ b/core/services/cluster/instance.go @@ -0,0 +1,189 @@ +// Package cluster records the frontend replicas that make up one LocalAI +// deployment and, later, the links between them. It is deliberately free of +// dependencies on core/services/nodes: nodes migrates and consumes the models +// declared here, so an import in the other direction would be a cycle. +package cluster + +import ( + "context" + "errors" + "fmt" + "net" + "net/url" + "strconv" + "strings" + "time" + + "gorm.io/gorm" + "gorm.io/gorm/clause" +) + +// ErrInstanceNotFound reports that no row exists for the requested instance ID. +// Callers distinguish it from a transport failure to decide whether to +// re-register or to retry. +var ErrInstanceNotFound = errors.New("cluster: instance not found") + +// Instance is one live frontend replica, keyed by the ID that replica chose for +// itself. Column sizes mirror nodes.BackendNode so both tables agree on what an +// ID and a host:port look like. +type Instance struct { + ID string `gorm:"primaryKey;size:36" json:"id"` + AdvertisedAddr string `gorm:"size:255" json:"advertised_addr"` // host:port other replicas dial + Version string `gorm:"size:64" json:"version"` + LastSeen time.Time `gorm:"index" json:"last_seen"` +} + +// Registry reads and writes the instances table. +type Registry struct { + db *gorm.DB +} + +// NewRegistry returns a Registry over db. Migration is the caller's job; the +// nodes registry owns the AutoMigrate for every table in this deployment so +// that a single advisory lock covers them all. +func NewRegistry(db *gorm.DB) *Registry { + return &Registry{db: db} +} + +// Register records this replica's address, refreshing LastSeen. It upserts on +// the primary key rather than deleting and re-inserting, so a concurrent Live +// never observes a live replica as missing. +func (r *Registry) Register(ctx context.Context, id, addr, version string) error { + inst := Instance{ + ID: id, + AdvertisedAddr: addr, + Version: version, + LastSeen: time.Now(), + } + if err := r.db.WithContext(ctx).Clauses(clause.OnConflict{ + Columns: []clause.Column{{Name: "id"}}, + DoUpdates: clause.AssignmentColumns([]string{"advertised_addr", "version", "last_seen"}), + }).Create(&inst).Error; err != nil { + return fmt.Errorf("registering instance %q: %w", id, err) + } + return nil +} + +// Heartbeat refreshes LastSeen for an already-registered instance. An unknown +// ID is an error rather than an insert: a heartbeat carries no address, so +// inserting would publish a replica nobody can reach. +func (r *Registry) Heartbeat(ctx context.Context, id string) error { + // gorm reports no error when a Where matches nothing, so the miss has to be + // read off RowsAffected. + res := r.db.WithContext(ctx).Model(&Instance{}). + Where("id = ?", id). + Update("last_seen", time.Now()) + if res.Error != nil { + return fmt.Errorf("heartbeating instance %q: %w", id, res.Error) + } + if res.RowsAffected == 0 { + return fmt.Errorf("heartbeating instance %q: %w", id, ErrInstanceNotFound) + } + return nil +} + +// Live returns the instances whose LastSeen is newer than now-within. +func (r *Registry) Live(ctx context.Context, within time.Duration) ([]Instance, error) { + var out []Instance + if err := r.db.WithContext(ctx). + Where("last_seen > ?", time.Now().Add(-within)). + Order("id"). + Find(&out).Error; err != nil { + return nil, fmt.Errorf("listing live instances: %w", err) + } + return out, nil +} + +// Get returns one instance, or ErrInstanceNotFound if it is not registered. +func (r *Registry) Get(ctx context.Context, id string) (*Instance, error) { + var inst Instance + err := r.db.WithContext(ctx).Where("id = ?", id).First(&inst).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, fmt.Errorf("getting instance %q: %w", id, ErrInstanceNotFound) + } + if err != nil { + return nil, fmt.Errorf("getting instance %q: %w", id, err) + } + return &inst, nil +} + +// DiscoverAdvertisedAddr determines the address this replica should advertise +// to its peers, with no operator configuration. +// +// Every replica in a deployment reaches the same PostgreSQL server, so the +// local interface that routes to PostgreSQL is on a network all the replicas +// demonstrably share. Opening a UDP socket toward the database sends no packet; +// it only asks the kernel to pick a source address for that route, which is the +// address to advertise. The caller supplies the port, since the frontend's +// listening port has nothing to do with the database's. +// +// Failures are returned rather than papered over with a fallback such as +// 127.0.0.1, which would publish an address no peer can dial. +func DiscoverAdvertisedAddr(dsn string, port int) (string, error) { + host, dbPort, err := dsnHostPort(dsn) + if err != nil { + return "", err + } + conn, err := net.Dial("udp", net.JoinHostPort(host, dbPort)) + if err != nil { + return "", fmt.Errorf("resolving route to database host %q: %w", host, err) + } + // Nothing was ever sent on this socket, so a close failure carries no + // information about the address we just read. + defer func() { _ = conn.Close() }() + local, ok := conn.LocalAddr().(*net.UDPAddr) + if !ok || local.IP == nil || local.IP.IsUnspecified() { + return "", fmt.Errorf("no local address on the route to database host %q", host) + } + return net.JoinHostPort(local.IP.String(), strconv.Itoa(port)), nil +} + +// dsnHostPort extracts the host and port from either DSN form gorm's postgres +// driver accepts: a URL ("postgres://user:pass@host:5432/db") or libpq keyword +// pairs ("host=... port=..."). +func dsnHostPort(dsn string) (string, string, error) { + const defaultPort = "5432" + dsn = strings.TrimSpace(dsn) + if dsn == "" { + return "", "", errors.New("empty database DSN") + } + + if strings.HasPrefix(dsn, "postgres://") || strings.HasPrefix(dsn, "postgresql://") { + u, err := url.Parse(dsn) + if err != nil { + return "", "", fmt.Errorf("parsing database DSN: %w", err) + } + host := u.Hostname() + if host == "" { + return "", "", errors.New("database DSN has no host") + } + port := u.Port() + if port == "" { + port = defaultPort + } + return host, port, nil + } + + host, port := "", defaultPort + for _, field := range strings.Fields(dsn) { + key, value, found := strings.Cut(field, "=") + if !found { + continue + } + switch key { + case "host": + host = value + case "port": + port = value + } + } + if host == "" { + return "", "", errors.New("database DSN has no host") + } + // A Unix socket directory tells us nothing about which interface reaches + // the database, so there is no address to derive. + if strings.HasPrefix(host, "/") { + return "", "", fmt.Errorf("database DSN uses a unix socket (%q); no routable address to advertise", host) + } + return host, port, nil +} diff --git a/core/services/cluster/instance_test.go b/core/services/cluster/instance_test.go new file mode 100644 index 000000000..112cc4851 --- /dev/null +++ b/core/services/cluster/instance_test.go @@ -0,0 +1,101 @@ +package cluster_test + +import ( + "context" + "net" + "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("Instance registry", func() { + var ( + db *gorm.DB + reg *cluster.Registry + ctx context.Context + ) + + BeforeEach(func() { + db = testutil.SetupTestDB() + Expect(db.AutoMigrate(&cluster.Instance{})).To(Succeed()) + reg = cluster.NewRegistry(db) + ctx = context.Background() + }) + + It("registers an instance and reads it back", func() { + Expect(reg.Register(ctx, "inst-a", "10.0.0.1:8080", "v1")).To(Succeed()) + + got, err := reg.Get(ctx, "inst-a") + Expect(err).ToNot(HaveOccurred()) + Expect(got.AdvertisedAddr).To(Equal("10.0.0.1:8080")) + Expect(got.Version).To(Equal("v1")) + }) + + It("re-registering the same id updates the address instead of duplicating", func() { + Expect(reg.Register(ctx, "inst-a", "10.0.0.1:8080", "v1")).To(Succeed()) + Expect(reg.Register(ctx, "inst-a", "10.0.0.9:9090", "v2")).To(Succeed()) + + live, err := reg.Live(ctx, time.Hour) + Expect(err).ToNot(HaveOccurred()) + Expect(live).To(HaveLen(1)) + Expect(live[0].AdvertisedAddr).To(Equal("10.0.0.9:9090")) + }) + + It("reports a missing instance distinguishably", func() { + _, err := reg.Get(ctx, "nope") + Expect(err).To(MatchError(cluster.ErrInstanceNotFound)) + }) + + It("excludes instances whose heartbeat has aged out", func() { + Expect(reg.Register(ctx, "stale", "10.0.0.1:8080", "v1")).To(Succeed()) + // Age the row directly; sleeping in a spec is forbidden. + Expect(db.Model(&cluster.Instance{}).Where("id = ?", "stale"). + Update("last_seen", time.Now().Add(-10*time.Minute)).Error).To(Succeed()) + + live, err := reg.Live(ctx, time.Minute) + Expect(err).ToNot(HaveOccurred()) + Expect(live).To(BeEmpty()) + }) + + It("brings a stale instance back with a heartbeat", func() { + Expect(reg.Register(ctx, "revive", "10.0.0.1:8080", "v1")).To(Succeed()) + Expect(db.Model(&cluster.Instance{}).Where("id = ?", "revive"). + Update("last_seen", time.Now().Add(-10*time.Minute)).Error).To(Succeed()) + Expect(reg.Heartbeat(ctx, "revive")).To(Succeed()) + + live, err := reg.Live(ctx, time.Minute) + Expect(err).ToNot(HaveOccurred()) + Expect(live).To(HaveLen(1)) + }) + + It("heartbeating an unknown instance is an error, not a silent insert", func() { + Expect(reg.Heartbeat(ctx, "ghost")).To(MatchError(cluster.ErrInstanceNotFound)) + }) +}) + +var _ = Describe("Advertised address discovery", func() { + // The address itself depends on host networking and is deliberately not + // asserted. What is portable is the shape: whatever interface routes to the + // database, the port must be the one the caller asked for, not the + // database's. + It("combines a local interface with the caller's port", func() { + addr, err := cluster.DiscoverAdvertisedAddr("postgres://198.51.100.1:5432/testdb", 8080) + if err != nil { + Skip("no route to a database host on this machine: " + err.Error()) + } + host, port, splitErr := net.SplitHostPort(addr) + Expect(splitErr).ToNot(HaveOccurred()) + Expect(port).To(Equal("8080")) + Expect(net.ParseIP(host)).ToNot(BeNil()) + }) + + It("refuses a DSN it cannot derive an address from", func() { + _, err := cluster.DiscoverAdvertisedAddr("", 8080) + Expect(err).To(HaveOccurred()) + }) +}) diff --git a/core/services/nodes/registry.go b/core/services/nodes/registry.go index 4b4f7f1c8..291ccd2db 100644 --- a/core/services/nodes/registry.go +++ b/core/services/nodes/registry.go @@ -9,6 +9,7 @@ import ( "github.com/google/uuid" "github.com/mudler/LocalAI/core/services/advisorylock" + "github.com/mudler/LocalAI/core/services/cluster" "github.com/mudler/LocalAI/pkg/system" "github.com/mudler/LocalAI/pkg/vrambudget" "github.com/mudler/xlog" @@ -442,7 +443,7 @@ 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 { - return db.AutoMigrate(&BackendNode{}, &NodeModel{}, &NodeLabel{}, &ModelSchedulingConfig{}, &PendingBackendOp{}, &ModelLoadInfo{}, &ModelLoadJob{}, &ModelConfigState{}) + return db.AutoMigrate(&BackendNode{}, &NodeModel{}, &NodeLabel{}, &ModelSchedulingConfig{}, &PendingBackendOp{}, &ModelLoadInfo{}, &ModelLoadJob{}, &ModelConfigState{}, &cluster.Instance{}) }); err != nil { return nil, fmt.Errorf("migrating node tables: %w", err) }