diff --git a/core/p2p/federated.go b/core/p2p/federated.go index 7ed8cbab6..cd1793a57 100644 --- a/core/p2p/federated.go +++ b/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) { diff --git a/core/p2p/federated_server.go b/core/p2p/federated_server.go index 36d7eb45d..7ed9920dc 100644 --- a/core/p2p/federated_server.go +++ b/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 { diff --git a/core/p2p/federated_test.go b/core/p2p/federated_test.go new file mode 100644 index 000000000..eb9eea145 --- /dev/null +++ b/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")) + }) +}) diff --git a/core/p2p/p2p_suite_test.go b/core/p2p/p2p_suite_test.go new file mode 100644 index 000000000..f24f9bbc2 --- /dev/null +++ b/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") +}