feat(p2p): gossip each peer's model set for model-aware federation

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

core/application/p2p.go

@@ -68,9 +68,21 @@ func (a *Application) StartP2P() error {
 			return err
 		}
 
+		// modelsFn reports the model names this instance currently serves so the
+		// federation proxy can route a request only to peers that have the
+		// requested model. It is re-evaluated on every announce tick.
+		modelsFn := func() []string {
+			cfgs := a.ModelConfigLoader().GetAllModelsConfigs()
+			names := make([]string, 0, len(cfgs))
+			for _, c := range cfgs {
+				names = append(names, c.Name)
+			}
+			return names
+		}
+
 		// Here a new node is created and started
 		// and a service is exposed by the node
-		node, err := p2p.ExposeService(ctx, "localhost", port, a.applicationConfig.P2PToken, p2p.NetworkID(networkID, p2p.FederatedID))
+		node, err := p2p.ExposeService(ctx, "localhost", port, a.applicationConfig.P2PToken, p2p.NetworkID(networkID, p2p.FederatedID), modelsFn)
 		if err != nil {
 			return err
 		}

core/cli/worker/worker_p2p.go

@@ -62,7 +62,7 @@ func (r *P2P) Run(ctx *cliContext.Context) error {
 			p = r.RunnerPort
 		}
 
-		_, err = p2p.ExposeService(c, address, p, r.Token, p2p.NetworkID(r.Peer2PeerNetworkID, p2p.LlamaCPPWorkerID))
+		_, err = p2p.ExposeService(c, address, p, r.Token, p2p.NetworkID(r.Peer2PeerNetworkID, p2p.LlamaCPPWorkerID), nil)
 		if err != nil {
 			return err
 		}
@@ -104,7 +104,7 @@ func (r *P2P) Run(ctx *cliContext.Context) error {
 			}
 		}()
 
-		_, err = p2p.ExposeService(c, address, fmt.Sprint(port), r.Token, p2p.NetworkID(r.Peer2PeerNetworkID, p2p.LlamaCPPWorkerID))
+		_, err = p2p.ExposeService(c, address, fmt.Sprint(port), r.Token, p2p.NetworkID(r.Peer2PeerNetworkID, p2p.LlamaCPPWorkerID), nil)
 		if err != nil {
 			return err
 		}

core/cli/worker/worker_p2p_mlx.go

@@ -81,7 +81,7 @@ func (r *P2PMLX) Run(ctx *cliContext.Context) error {
 		}
 	}()
 
-	_, err = p2p.ExposeService(c, address, fmt.Sprint(port), r.Token, p2p.NetworkID(r.Peer2PeerNetworkID, p2p.MLXWorkerID))
+	_, err = p2p.ExposeService(c, address, fmt.Sprint(port), r.Token, p2p.NetworkID(r.Peer2PeerNetworkID, p2p.MLXWorkerID), nil)
 	if err != nil {
 		return err
 	}

core/p2p/p2p.go

@@ -312,7 +312,7 @@ func ensureService(ctx context.Context, n *node.Node, nd *schema.NodeData, sserv
 }
 
 // This is the P2P worker main
-func ExposeService(ctx context.Context, host, port, token, servicesID string) (*node.Node, error) {
+func ExposeService(ctx context.Context, host, port, token, servicesID string, modelsFn func() []string) (*node.Node, error) {
 	if servicesID == "" {
 		servicesID = defaultServicesID
 	}
@@ -348,11 +348,16 @@ func ExposeService(ctx context.Context, host, port, token, servicesID string) (*
 		20*time.Second,
 		func() {
 			updatedMap := map[string]any{}
+			var models []string
+			if modelsFn != nil {
+				models = modelsFn()
+			}
 			updatedMap[name] = &schema.NodeData{
 				Name:          name,
 				LastSeen:      time.Now(),
 				ID:            nodeID(name),
 				AvailableVRAM: xsysinfo.GetGPUAggregateInfo().FreeVRAM,
+				Models:        models,
 			}
 			ledger.Add(servicesID, updatedMap)
 		},

core/schema/localai.go

@@ -194,6 +194,12 @@ type NodeData struct {
 	// Zero for CPU-only nodes and for peers on an older version that does not
 	// publish it; the routing policy treats zero as the lowest VRAM tier.
 	AvailableVRAM uint64
+	// Models is the set of model names this peer currently serves, gossiped so
+	// the federation proxy can route a request only to peers that have the
+	// requested model. Empty means "unknown" (an older peer, or one that has
+	// not loaded any model yet) and is treated as eligible for any model so a
+	// mixed-version swarm is not starved.
+	Models []string
 }
 
 func (d NodeData) IsOnline() bool {

core/schema/nodedata_test.go

@@ -31,4 +31,9 @@ var _ = Describe("NodeData", func() {
 		nd := schema.NodeData{AvailableVRAM: 8_000_000_000}
 		Expect(nd.AvailableVRAM).To(Equal(uint64(8_000_000_000)))
 	})
+
+	It("carries the advertised model set", func() {
+		nd := schema.NodeData{Models: []string{"llama-3", "qwen"}}
+		Expect(nd.Models).To(ConsistOf("llama-3", "qwen"))
+	})
 })