From e49f456ea388db2d1a597b956eff53d96224ccb0 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Mon, 2 Nov 2026 20:03:51 +0000 Subject: [PATCH] fix(p2p): bind resource limits at CLI boundaries Pass the resource-limit choice through application and node configuration. Keep only the exact string true enabled, including worker and explorer entry points. Leave other legacy environment reads unchanged. Assisted-by: nib:gpt-6.1-sol Signed-off-by: Ettore Di Giacinto Assisted-by: Claude Code:claude-sonnet-5-5 --- core/application/p2p.go | 4 ++-- core/cli/explorer.go | 5 ++-- core/cli/p2p_worker_hint_test.go | 23 +++++++++++++++++++ core/cli/run.go | 3 ++- core/cli/worker/worker_p2p.go | 5 ++-- core/cli/worker/worker_p2p_mlx.go | 3 ++- core/config/application_config.go | 5 ++++ core/explorer/discovery.go | 7 ++++-- core/p2p/node_config_internal_test.go | 20 ++++------------ core/p2p/p2p.go | 22 +++++++++++------- .../features/distributed_inferencing.md | 2 +- 11 files changed, 65 insertions(+), 34 deletions(-) diff --git a/core/application/p2p.go b/core/application/p2p.go index d39409a944..3c06bacaf1 100644 --- a/core/application/p2p.go +++ b/core/application/p2p.go @@ -70,7 +70,7 @@ func (a *Application) StartP2P() error { // 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), p2p.NodeConfig{EnableLimits: a.applicationConfig.P2PEnableLimits}) if err != nil { return err } @@ -88,7 +88,7 @@ func (a *Application) StartP2P() error { // If a node wasn't created previously, create it if n == nil { - node, err := p2p.NewNode(a.applicationConfig.P2PToken) + node, err := p2p.NewNode(a.applicationConfig.P2PToken, p2p.NodeConfig{EnableLimits: a.applicationConfig.P2PEnableLimits}) if err != nil { return err } diff --git a/core/cli/explorer.go b/core/cli/explorer.go index d520dac212..1997effaeb 100644 --- a/core/cli/explorer.go +++ b/core/cli/explorer.go @@ -12,6 +12,7 @@ import ( ) type ExplorerCMD struct { + P2PEnableLimits string `env:"LOCALAI_P2P_ENABLE_LIMITS" help:"Set to true to enable P2P resource limits" group:"p2p"` Address string `env:"LOCALAI_ADDRESS,ADDRESS" default:":8080" help:"Bind address for the API server" group:"api"` PoolDatabase string `env:"LOCALAI_POOL_DATABASE,POOL_DATABASE" default:"explorer.json" help:"Path to the pool database" group:"api"` ConnectionTimeout string `env:"LOCALAI_CONNECTION_TIMEOUT,CONNECTION_TIMEOUT" default:"2m" help:"Connection timeout for the explorer" group:"api"` @@ -34,12 +35,12 @@ func (e *ExplorerCMD) Run(ctx *cliContext.Context) error { } if e.WithSync { - ds := explorer.NewDiscoveryServer(db, dur, e.ConnectionErrorThreshold) + ds := explorer.NewDiscoveryServer(db, dur, e.ConnectionErrorThreshold, e.P2PEnableLimits == "true") go ds.Start(context.Background(), true) } if e.OnlySync { - ds := explorer.NewDiscoveryServer(db, dur, e.ConnectionErrorThreshold) + ds := explorer.NewDiscoveryServer(db, dur, e.ConnectionErrorThreshold, e.P2PEnableLimits == "true") ctx := context.Background() return ds.Start(ctx, false) diff --git a/core/cli/p2p_worker_hint_test.go b/core/cli/p2p_worker_hint_test.go index 39c7a269db..f28194f2a6 100644 --- a/core/cli/p2p_worker_hint_test.go +++ b/core/cli/p2p_worker_hint_test.go @@ -7,6 +7,7 @@ import ( . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" + "github.com/mudler/LocalAI/core/cli/worker" "github.com/mudler/LocalAI/core/config" ) @@ -34,3 +35,25 @@ var _ = Describe("p2p worker hint", func() { Expect(ctx.Command()).To(Equal("p2p-worker p2p-llama-cpp-rpc")) }) }) + +var _ = Describe("P2P limits CLI boundary", func() { + It("preserves exact true semantics at every entry point", func() { + for _, value := range []string{"", "true", "false", "1", "yes"} { + GinkgoT().Setenv("LOCALAI_P2P_ENABLE_LIMITS", value) + run := RunCMD{} + rpc := worker.P2P{} + mlx := worker.P2PMLX{} + explorer := ExplorerCMD{} + for _, entry := range []any{&run, &rpc, &mlx, &explorer} { + parser, err := kong.New(entry, kong.Vars{"basepath": ".", "generatedcontentpath": ".", "uploadpath": ".", "galleries": config.DefaultGalleriesJSON, "backends": config.DefaultBackendGalleriesJSON, "version": "test"}) + Expect(err).NotTo(HaveOccurred()) + _, err = parser.Parse(nil) + Expect(err).NotTo(HaveOccurred()) + } + Expect(run.P2PEnableLimits == "true").To(Equal(value == "true")) + Expect(rpc.P2PEnableLimits == "true").To(Equal(value == "true")) + Expect(mlx.P2PEnableLimits == "true").To(Equal(value == "true")) + Expect(explorer.P2PEnableLimits == "true").To(Equal(value == "true")) + } + }) +}) diff --git a/core/cli/run.go b/core/cli/run.go index 4ce1c8c4a4..a3d692bc00 100644 --- a/core/cli/run.go +++ b/core/cli/run.go @@ -40,6 +40,7 @@ import ( const workerJoinCommand = "local-ai p2p-worker p2p-llama-cpp-rpc" type RunCMD struct { + P2PEnableLimits string `env:"LOCALAI_P2P_ENABLE_LIMITS" help:"Set to true to enable P2P resource limits" group:"p2p"` AddonHostPolicy string `name:"addon-host-policy" env:"LOCALAI_ADDON_HOST_POLICY" help:"Absolute path to administrator host operation policy JSON, read once at startup. Empty disables environment policy" group:"addons"` AddonHostOps string `name:"addon-host-ops" env:"LOCALAI_ADDON_HOST_OPS" help:"Add-ons that may call host operations, as a comma list of addon-id=operation-prefix. An add-on that is not listed gets no host operation channel. Empty gives none" group:"addons"` AddonsEnabled bool `name:"addons-enabled" env:"LOCALAI_ADDONS_ENABLED" default:"true" help:"Run the add-on owner, which installs and supervises add-ons. Turn it off with --addons-enabled=false" group:"addons"` @@ -642,7 +643,7 @@ func (r *RunCMD) Run(ctx *cliContext.Context) (result error) { xlog.Info("To use the token, you can run the following command in another node or terminal:") fmt.Printf("export TOKEN=\"%s\"\n%s\n", token, workerJoinCommand) } - opts = append(opts, config.WithP2PToken(token)) + opts = append(opts, config.WithP2PToken(token), config.WithP2PEnableLimits(r.P2PEnableLimits == "true")) } if r.Federated { diff --git a/core/cli/worker/worker_p2p.go b/core/cli/worker/worker_p2p.go index 78dbd201eb..16493394d8 100644 --- a/core/cli/worker/worker_p2p.go +++ b/core/cli/worker/worker_p2p.go @@ -17,6 +17,7 @@ import ( ) type P2P struct { + P2PEnableLimits string `env:"LOCALAI_P2P_ENABLE_LIMITS" help:"Set to true to enable P2P resource limits" group:"p2p"` WorkerFlags `embed:""` Token string `env:"LOCALAI_TOKEN,LOCALAI_P2P_TOKEN,TOKEN" help:"P2P token to use"` NoRunner bool `env:"LOCALAI_NO_RUNNER,NO_RUNNER" help:"Do not start the llama-cpp-rpc-server"` @@ -63,7 +64,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), p2p.NodeConfig{EnableLimits: r.P2PEnableLimits == "true"}) if err != nil { return err } @@ -105,7 +106,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), p2p.NodeConfig{EnableLimits: r.P2PEnableLimits == "true"}) if err != nil { return err } diff --git a/core/cli/worker/worker_p2p_mlx.go b/core/cli/worker/worker_p2p_mlx.go index e7afc24317..0285e9cb8a 100644 --- a/core/cli/worker/worker_p2p_mlx.go +++ b/core/cli/worker/worker_p2p_mlx.go @@ -15,6 +15,7 @@ import ( ) type P2PMLX struct { + P2PEnableLimits string `env:"LOCALAI_P2P_ENABLE_LIMITS" help:"Set to true to enable P2P resource limits" group:"p2p"` WorkerFlags `embed:""` Token string `env:"LOCALAI_TOKEN,LOCALAI_P2P_TOKEN,TOKEN" help:"P2P token to use"` Peer2PeerNetworkID string `env:"LOCALAI_P2P_NETWORK_ID,P2P_NETWORK_ID" help:"Network ID for P2P mode" group:"p2p"` @@ -82,7 +83,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), p2p.NodeConfig{EnableLimits: r.P2PEnableLimits == "true"}) if err != nil { return err } diff --git a/core/config/application_config.go b/core/config/application_config.go index eb0576ebab..2d01e6b9e4 100644 --- a/core/config/application_config.go +++ b/core/config/application_config.go @@ -167,6 +167,7 @@ type ApplicationConfig struct { PreloadModelsFromPath string CORSAllowOrigins string ApiKeys []string + P2PEnableLimits bool P2PToken string P2PNetworkID string Federated bool @@ -1657,3 +1658,7 @@ func (o *ApplicationConfig) EffectiveChatsMaxIterations() int { } return o.ChatsMaxIterations } + +func WithP2PEnableLimits(enabled bool) AppOption { + return func(o *ApplicationConfig) { o.P2PEnableLimits = enabled } +} diff --git a/core/explorer/discovery.go b/core/explorer/discovery.go index a67395e738..f599d26a0c 100644 --- a/core/explorer/discovery.go +++ b/core/explorer/discovery.go @@ -16,6 +16,7 @@ import ( type DiscoveryServer struct { sync.Mutex + limits bool database *Database connectionTime time.Duration errorThreshold int @@ -32,14 +33,16 @@ func redactToken(t string) string { // NewDiscoveryServer creates a new DiscoveryServer with the given Database. // it keeps the db state in sync with the network state -func NewDiscoveryServer(db *Database, dur time.Duration, failureThreshold int) *DiscoveryServer { +func NewDiscoveryServer(db *Database, dur time.Duration, failureThreshold int, limits ...bool) *DiscoveryServer { if dur == 0 { dur = 50 * time.Second } if failureThreshold == 0 { failureThreshold = 3 } + enabled := len(limits) > 0 && limits[0] return &DiscoveryServer{ + limits: enabled, database: db, connectionTime: dur, errorThreshold: failureThreshold, @@ -64,7 +67,7 @@ func (s *DiscoveryServer) runBackground() { // Get the number of nodes // save it in the current state (mutex) // do not do in parallel - n, err := p2p.NewNode(token) + n, err := p2p.NewNode(token, p2p.NodeConfig{EnableLimits: s.limits}) if err != nil { xlog.Error("Failed to create node", "error", err) s.failedToken(token) diff --git a/core/p2p/node_config_internal_test.go b/core/p2p/node_config_internal_test.go index ffc6e6bc3b..aba63e75e8 100644 --- a/core/p2p/node_config_internal_test.go +++ b/core/p2p/node_config_internal_test.go @@ -6,24 +6,14 @@ import ( ) var _ = Describe("newNodeConfig", func() { - It("leaves resource limits disabled when LOCALAI_P2P_ENABLE_LIMITS is unset", func() { - GinkgoT().Setenv("LOCALAI_P2P_ENABLE_LIMITS", "") - + It("leaves resource limits disabled by default and when explicitly disabled", func() { Expect(newNodeConfig("token").Limit.Enable).To(BeFalse()) + + Expect(newNodeConfig("token", NodeConfig{EnableLimits: false}).Limit.Enable).To(BeFalse()) }) - It("enables resource limits when LOCALAI_P2P_ENABLE_LIMITS is \"true\"", func() { - GinkgoT().Setenv("LOCALAI_P2P_ENABLE_LIMITS", "true") - - Expect(newNodeConfig("token").Limit.Enable).To(BeTrue()) - }) - - It("leaves resource limits disabled for any other value", func() { - for _, v := range []string{"false", "1", "yes"} { - GinkgoT().Setenv("LOCALAI_P2P_ENABLE_LIMITS", v) - - Expect(newNodeConfig("token").Limit.Enable).To(BeFalse(), "value %q", v) - } + It("enables resource limits when explicitly requested", func() { + Expect(newNodeConfig("token", NodeConfig{EnableLimits: true}).Limit.Enable).To(BeTrue()) }) It("keeps the network token", func() { diff --git a/core/p2p/p2p.go b/core/p2p/p2p.go index 28525c40a3..62a7f7c4a0 100644 --- a/core/p2p/p2p.go +++ b/core/p2p/p2p.go @@ -311,13 +311,13 @@ 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, settings ...NodeConfig) (*node.Node, error) { if servicesID == "" { servicesID = defaultServicesID } llger := logger.New(log.LevelFatal) - nodeOpts, err := newNodeOpts(token) + nodeOpts, err := newNodeOpts(token, settings...) if err != nil { return nil, err } @@ -359,8 +359,8 @@ func ExposeService(ctx context.Context, host, port, token, servicesID string) (* return n, err } -func NewNode(token string) (*node.Node, error) { - nodeOpts, err := newNodeOpts(token) +func NewNode(token string, settings ...NodeConfig) (*node.Node, error) { + nodeOpts, err := newNodeOpts(token, settings...) if err != nil { return nil, err } @@ -373,10 +373,10 @@ func NewNode(token string) (*node.Node, error) { return n, nil } -func newNodeOpts(token string) ([]node.Option, error) { +func newNodeOpts(token string, settings ...NodeConfig) ([]node.Option, error) { llger := logger.New(log.LevelFatal) - c := newNodeConfig(token) + c := newNodeConfig(token, settings...) nodeOpts, _, err := c.ToOpts(llger) if err != nil { @@ -388,16 +388,22 @@ func newNodeOpts(token string) ([]node.Option, error) { return nodeOpts, nil } +// NodeConfig carries P2P settings supplied by the caller. +type NodeConfig struct{ EnableLimits bool } + // newNodeConfig builds the edgevpn node configuration from the LOCALAI_P2P_* // environment variables. -func newNodeConfig(token string) config.Config { +func newNodeConfig(token string, settings ...NodeConfig) config.Config { defaultInterval := 10 * time.Second // TODO: move this up, expose more config options when creating a node noDHT := os.Getenv("LOCALAI_P2P_DISABLE_DHT") == "true" // Resource limits are opt-in: the libp2p resource manager is enabled only // when this is set to "true", as documented. - enableLimits := os.Getenv("LOCALAI_P2P_ENABLE_LIMITS") == "true" + enableLimits := false + if len(settings) > 0 { + enableLimits = settings[0].EnableLimits + } var listenMaddrs []string var bootstrapPeers []string diff --git a/docs/content/features/distributed_inferencing.md b/docs/content/features/distributed_inferencing.md index 54fde1c2c4..12c4521a22 100644 --- a/docs/content/features/distributed_inferencing.md +++ b/docs/content/features/distributed_inferencing.md @@ -141,7 +141,7 @@ There are options that can be tweaked or parameters that can be set using enviro | **LOCALAI_FEDERATED** | Set to "true" to enable federated mode | | **FEDERATED_SERVER** | Set to "true" to enable federated server | | **LOCALAI_P2P_DISABLE_DHT** | Set to "true" to disable DHT and enable p2p layer to be local only (mDNS) | -| **LOCALAI_P2P_ENABLE_LIMITS** | Set to "true" to enable connection limits and resources management (useful when running with poor connectivity or want to limit resources consumption) | +| **LOCALAI_P2P_ENABLE_LIMITS** | Set to "true" (or pass `--p2p-enable-limits=true` to `run`, a `p2p-worker` command, or `explorer`) to enable connection limits and resources management (useful when running with poor connectivity or want to limit resources consumption) | | **LOCALAI_P2P_LISTEN_MADDRS** | Set to comma separated list of multiaddresses to override default libp2p 0.0.0.0 multiaddresses | | **LOCALAI_P2P_DHT_ANNOUNCE_MADDRS** | Set to comma separated list of multiaddresses to override announcing of listen multiaddresses (useful when external address:port is remapped) | | **LOCALAI_P2P_BOOTSTRAP_PEERS_MADDRS** | Set to comma separated list of multiaddresses to specify custom DHT bootstrap nodes |