feat(p2p): route federation with shared clusterrouting policy (load + VRAM)

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

core/p2p/federated.go

@@ -4,7 +4,10 @@ import (
 	"fmt"
 	"math/rand/v2"
 	"sync"
+	"time"
 
+	"github.com/mudler/LocalAI/core/schema"
+	"github.com/mudler/LocalAI/pkg/clusterrouting"
 	"github.com/mudler/xlog"
 )
 
@@ -74,28 +77,46 @@ func (fs *FederatedServer) syncTableStatus() {
 	}
 }
 
-func (fs *FederatedServer) SelectLeastUsedServer() string {
-	fs.syncTableStatus()
+// buildFederatedCandidates maps the currently-online federated peers into the
+// shared routing policy's candidate form. InFlight comes from the per-peer
+// request counter (lower means fewer requests routed there); AvailableVRAM
+// comes from the gossiped NodeData. LastUsed is intentionally left zero:
+// federation has no per-peer last-used clock, and the request counter already
+// spreads load, so the VRAM tier deterministically breaks in-flight ties.
+func buildFederatedCandidates(nodes []schema.NodeData, requestTable map[string]int, now time.Time) []clusterrouting.ReplicaCandidate {
+	candidates := make([]clusterrouting.ReplicaCandidate, 0, len(nodes))
+	for _, nd := range nodes {
+		if !nd.IsOnlineAt(now) {
+			continue
+		}
+		candidates = append(candidates, clusterrouting.ReplicaCandidate{
+			NodeID:        nd.ID,
+			InFlight:      requestTable[nd.ID],
+			AvailableVRAM: nd.AvailableVRAM,
+		})
+	}
+	return candidates
+}
 
+// SelectBestServer picks the online federated peer to serve the next request
+// using the shared cluster-routing policy (least in-flight, then most free
+// VRAM). Returns "" when no peer is online.
+func (fs *FederatedServer) SelectBestServer() string {
+	fs.syncTableStatus()
+	// Snapshot the node set before taking fs.Lock so the fs critical section
+	// only guards requestTable. GetAvailableNodes takes its own global mutex;
+	// calling it outside fs.Lock avoids a fs.Mutex -> node.mu lock ordering.
+	nodes := GetAvailableNodes(fs.service)
 	fs.Lock()
 	defer fs.Unlock()
-
-	xlog.Debug("SelectLeastUsedServer()", "request_table", fs.requestTable)
-
-	// cycle over requestTable and find the entry with the lower number
-	// if there are multiple entries with the same number, select one randomly
-	// if there are no entries, return an empty string
-	var min int
-	var minKey string
-	for k, v := range fs.requestTable {
-		if min == 0 || v < min {
-			min = v
-			minKey = k
-		}
+	candidates := buildFederatedCandidates(nodes, fs.requestTable, time.Now())
+	best := clusterrouting.PickBestReplica(candidates)
+	if best == nil {
+		xlog.Debug("No online federated peers to select", "request_table", fs.requestTable)
+		return ""
 	}
-	xlog.Debug("Selected tunnel", "tunnel", minKey, "requests_served", min, "request_table", fs.requestTable)
-
-	return minKey
+	xlog.Debug("Selected federated peer", "peer", best.NodeID, "request_table", fs.requestTable)
+	return best.NodeID
 }
 
 func (fs *FederatedServer) RecordRequest(nodeID string) {

core/p2p/federated_server.go

@@ -70,9 +70,9 @@ func (fs *FederatedServer) proxy(ctx context.Context, node *node.Node) error {
 				} else if fs.loadBalanced {
 					xlog.Debug("Load balancing request")
 
-					workerID = fs.SelectLeastUsedServer()
+					workerID = fs.SelectBestServer()
 					if workerID == "" {
-						xlog.Debug("Least used server not found, selecting random")
+						xlog.Debug("Best server not found, selecting random")
 						workerID = fs.RandomServer()
 					}
 				} else {

core/p2p/federated_test.go

@@ -0,0 +1,63 @@
+package p2p
+
+import (
+	"time"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+
+	"github.com/mudler/LocalAI/core/schema"
+	"github.com/mudler/LocalAI/pkg/clusterrouting"
+)
+
+var _ = Describe("buildFederatedCandidates", func() {
+	ref := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC)
+	onlineSeen := ref.Add(-10 * time.Second)
+	offlineSeen := ref.Add(-2 * time.Minute)
+
+	It("excludes offline nodes", func() {
+		nodes := []schema.NodeData{
+			{ID: "online", LastSeen: onlineSeen},
+			{ID: "offline", LastSeen: offlineSeen},
+		}
+		cands := buildFederatedCandidates(nodes, map[string]int{}, ref)
+		Expect(cands).To(HaveLen(1))
+		Expect(cands[0].NodeID).To(Equal("online"))
+	})
+
+	It("maps the request counter to InFlight and defaults missing entries to zero", func() {
+		nodes := []schema.NodeData{
+			{ID: "a", LastSeen: onlineSeen},
+			{ID: "b", LastSeen: onlineSeen},
+		}
+		cands := buildFederatedCandidates(nodes, map[string]int{"a": 4}, ref)
+		byID := map[string]int{}
+		for _, c := range cands {
+			byID[c.NodeID] = c.InFlight
+		}
+		Expect(byID["a"]).To(Equal(4))
+		Expect(byID["b"]).To(Equal(0))
+	})
+
+	It("carries gossiped AvailableVRAM into the candidate", func() {
+		nodes := []schema.NodeData{
+			{ID: "gpu", LastSeen: onlineSeen, AvailableVRAM: 24_000_000_000},
+		}
+		cands := buildFederatedCandidates(nodes, map[string]int{}, ref)
+		Expect(cands[0].AvailableVRAM).To(Equal(uint64(24_000_000_000)))
+	})
+
+	It("produces candidates the shared policy ranks by least in-flight then most VRAM", func() {
+		// busy-big has the most VRAM but is busy, so it must lose. Among the two
+		// idle peers, the one with more free VRAM wins (VRAM breaks the tie).
+		nodes := []schema.NodeData{
+			{ID: "busy-big", LastSeen: onlineSeen, AvailableVRAM: 80_000_000_000},
+			{ID: "idle-small", LastSeen: onlineSeen, AvailableVRAM: 8_000_000_000},
+			{ID: "idle-big", LastSeen: onlineSeen, AvailableVRAM: 24_000_000_000},
+		}
+		cands := buildFederatedCandidates(nodes, map[string]int{"busy-big": 3}, ref)
+		best := clusterrouting.PickBestReplica(cands)
+		Expect(best).ToNot(BeNil())
+		Expect(best.NodeID).To(Equal("idle-big"))
+	})
+})

core/p2p/p2p_suite_test.go

@@ -0,0 +1,13 @@
+package p2p
+
+import (
+	"testing"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+func TestP2P(t *testing.T) {
+	RegisterFailHandler(Fail)
+	RunSpecs(t, "P2P Suite")
+}