dearwolf/LocalAI
1
0
Fork 0
mirror of https://github.com/mudler/LocalAI.git synced 2025-05-20 10:35:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 21

fix(p2p): allocate tunnels only when needed (#3259)

Browse source 
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
Ettore Di Giacinto 2024-08-17 15:03:55 +02:00 • committed by GitHub
parent 1ed5af1da8
commit 27b03a52f3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 23 additions and 20 deletions
Download patch file Download diff file Expand all files Collapse all files

4
core/cli/run.go
View file

@ -135,7 +135,7 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error {
os.Setenv("LLAMACPP_GRPC_SERVERS", tunnelEnvVar) os.Setenv("LLAMACPP_GRPC_SERVERS", tunnelEnvVar)
log.Debug().Msgf("setting LLAMACPP_GRPC_SERVERS to %s", tunnelEnvVar) log.Debug().Msgf("setting LLAMACPP_GRPC_SERVERS to %s", tunnelEnvVar)
}); err != nil { }, true); err != nil {
return err return err
} }
} }
@ -153,7 +153,7 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error {
return err return err
} }
if err := p2p.ServiceDiscoverer(context.Background(), node, token, p2p.NetworkID(r.Peer2PeerNetworkID, p2p.FederatedID), nil); err != nil { if err := p2p.ServiceDiscoverer(context.Background(), node, token, p2p.NetworkID(r.Peer2PeerNetworkID, p2p.FederatedID), nil, false); err != nil {
return err return err
} }
} }

2
core/p2p/federated_server.go
View file

@ -29,7 +29,7 @@ func (f *FederatedServer) Start(ctx context.Context) error {
if err := ServiceDiscoverer(ctx, n, f.p2ptoken, f.service, func(servicesID string, tunnel NodeData) { if err := ServiceDiscoverer(ctx, n, f.p2ptoken, f.service, func(servicesID string, tunnel NodeData) {
log.Debug().Msgf("Discovered node: %s", tunnel.ID) log.Debug().Msgf("Discovered node: %s", tunnel.ID)
}); err != nil { }, true); err != nil {
return err return err
} }

35
core/p2p/p2p.go
View file

@ -139,11 +139,11 @@ func allocateLocalService(ctx context.Context, node *node.Node, listenAddr, serv
// This is the main of the server (which keeps the env variable updated) // This is the main of the server (which keeps the env variable updated)
// This starts a goroutine that keeps LLAMACPP_GRPC_SERVERS updated with the discovered services // This starts a goroutine that keeps LLAMACPP_GRPC_SERVERS updated with the discovered services
func ServiceDiscoverer(ctx context.Context, n *node.Node, token, servicesID string, discoveryFunc func(serviceID string, node NodeData)) error { func ServiceDiscoverer(ctx context.Context, n *node.Node, token, servicesID string, discoveryFunc func(serviceID string, node NodeData), allocate bool) error {
if servicesID == "" { if servicesID == "" {
servicesID = defaultServicesID servicesID = defaultServicesID
} }
tunnels, err := discoveryTunnels(ctx, n, token, servicesID) tunnels, err := discoveryTunnels(ctx, n, token, servicesID, allocate)
if err != nil { if err != nil {
return err return err
} }
@ -170,7 +170,7 @@ func ServiceDiscoverer(ctx context.Context, n *node.Node, token, servicesID stri
return nil return nil
} }
func discoveryTunnels(ctx context.Context, n *node.Node, token, servicesID string) (chan NodeData, error) { func discoveryTunnels(ctx context.Context, n *node.Node, token, servicesID string, allocate bool) (chan NodeData, error) {
tunnels := make(chan NodeData) tunnels := make(chan NodeData)
err := n.Start(ctx) err := n.Start(ctx)
@ -209,7 +209,7 @@ func discoveryTunnels(ctx context.Context, n *node.Node, token, servicesID strin
zlog.Error().Msg("cannot unmarshal node data") zlog.Error().Msg("cannot unmarshal node data")
continue continue
} }
ensureService(ctx, n, nd, k) ensureService(ctx, n, nd, k, allocate)
muservice.Lock() muservice.Lock()
if _, ok := service[nd.Name]; ok { if _, ok := service[nd.Name]; ok {
tunnels <- service[nd.Name].NodeData tunnels <- service[nd.Name].NodeData
@ -231,7 +231,7 @@ type nodeServiceData struct {
var service = map[string]nodeServiceData{} var service = map[string]nodeServiceData{}
var muservice sync.Mutex var muservice sync.Mutex
func ensureService(ctx context.Context, n *node.Node, nd *NodeData, sserv string) { func ensureService(ctx context.Context, n *node.Node, nd *NodeData, sserv string, allocate bool) {
muservice.Lock() muservice.Lock()
defer muservice.Unlock() defer muservice.Unlock()
if ndService, found := service[nd.Name]; !found { if ndService, found := service[nd.Name]; !found {
@ -240,22 +240,25 @@ func ensureService(ctx context.Context, n *node.Node, nd *NodeData, sserv string
zlog.Debug().Msgf("Node %s is offline", nd.ID) zlog.Debug().Msgf("Node %s is offline", nd.ID)
return return
} }
newCtxm, cancel := context.WithCancel(ctx)
// Start the service
port, err := freeport.GetFreePort()
if err != nil {
zlog.Error().Err(err).Msgf("Could not allocate a free port for %s", nd.ID)
return
}
tunnelAddress := fmt.Sprintf("127.0.0.1:%d", port) newCtxm, cancel := context.WithCancel(ctx)
nd.TunnelAddress = tunnelAddress if allocate {
// Start the service
port, err := freeport.GetFreePort()
if err != nil {
zlog.Error().Err(err).Msgf("Could not allocate a free port for %s", nd.ID)
return
}
tunnelAddress := fmt.Sprintf("127.0.0.1:%d", port)
nd.TunnelAddress = tunnelAddress
go allocateLocalService(newCtxm, n, tunnelAddress, sserv)
zlog.Debug().Msgf("Starting service %s on %s", sserv, tunnelAddress)
}
service[nd.Name] = nodeServiceData{ service[nd.Name] = nodeServiceData{
NodeData: *nd, NodeData: *nd,
CancelFunc: cancel, CancelFunc: cancel,
} }
go allocateLocalService(newCtxm, n, tunnelAddress, sserv)
zlog.Debug().Msgf("Starting service %s on %s", sserv, tunnelAddress)
} else { } else {
// Check if the service is still alive // Check if the service is still alive
// if not cancel the context // if not cancel the context

2
core/p2p/p2p_disabled.go
View file

@ -18,7 +18,7 @@ func (f *FederatedServer) Start(ctx context.Context) error {
return fmt.Errorf("not implemented") return fmt.Errorf("not implemented")
} }
func ServiceDiscoverer(ctx context.Context, node *node.Node, token, servicesID string, fn func(string, NodeData)) error { func ServiceDiscoverer(ctx context.Context, node *node.Node, token, servicesID string, fn func(string, NodeData), allocate bool) error {
return fmt.Errorf("not implemented") return fmt.Errorf("not implemented")
} }