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 <mudler@localai.io>

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")
+}

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
+}

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())
+	})
+})

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)
 	}