diff --git a/core/application/p2p.go b/core/application/p2p.go index d39409a94..cb0ad6712 100644 --- a/core/application/p2p.go +++ b/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 } diff --git a/core/cli/worker/worker_p2p.go b/core/cli/worker/worker_p2p.go index c7ff254ea..223f4f908 100644 --- a/core/cli/worker/worker_p2p.go +++ b/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 } diff --git a/core/cli/worker/worker_p2p_mlx.go b/core/cli/worker/worker_p2p_mlx.go index 7edd1673d..77a95394e 100644 --- a/core/cli/worker/worker_p2p_mlx.go +++ b/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 } diff --git a/core/p2p/p2p.go b/core/p2p/p2p.go index a108147bf..4cf892c62 100644 --- a/core/p2p/p2p.go +++ b/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) }, diff --git a/core/schema/localai.go b/core/schema/localai.go index 890136bf8..0ab66aaa7 100644 --- a/core/schema/localai.go +++ b/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 { diff --git a/core/schema/nodedata_test.go b/core/schema/nodedata_test.go index 61a21ed21..3c38fc2af 100644 --- a/core/schema/nodedata_test.go +++ b/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")) + }) })