Add token generation for federated server; remove additional P2P setup for federated instance mode

This commit is contained in:
mintyleaf 2025-03-29 16:54:29 +04:00
parent aa7171dd5d
commit 5c35029ae7
3 changed files with 24 additions and 9 deletions

View file

@ -2,10 +2,12 @@ package cli
import ( import (
"context" "context"
"fmt"
cliContext "github.com/mudler/LocalAI/core/cli/context" cliContext "github.com/mudler/LocalAI/core/cli/context"
cliP2P "github.com/mudler/LocalAI/core/cli/p2p" cliP2P "github.com/mudler/LocalAI/core/cli/p2p"
"github.com/mudler/LocalAI/core/p2p" "github.com/mudler/LocalAI/core/p2p"
"github.com/rs/zerolog/log"
) )
type FederatedCLI struct { type FederatedCLI struct {
@ -18,6 +20,24 @@ type FederatedCLI struct {
func (f *FederatedCLI) Run(ctx *cliContext.Context) error { func (f *FederatedCLI) Run(ctx *cliContext.Context) error {
if f.Peer2PeerToken == "" {
log.Info().Msg("No token provided, generating one")
connectionData, err := p2p.GenerateNewConnectionData(
f.Peer2PeerDHTInterval, f.Peer2PeerOTPInterval,
f.Peer2PeerPrivkey, f.Peer2PeerUsePeerguard,
)
if err != nil {
log.Warn().Msgf("Error generating token: %s", err.Error())
}
f.Peer2PeerToken = connectionData.Base64()
log.Info().Msg("Generated Token:")
fmt.Println(f.Peer2PeerToken)
log.Info().Msg("To use the token, you can run the following command in another node or terminal:")
fmt.Printf("export TOKEN=\"%s\"\nlocal-ai worker p2p-llama-cpp-rpc\n", f.Peer2PeerToken)
}
fs := p2p.NewFederatedServer(f.Address, p2p.NetworkID(f.Peer2PeerNetworkID, p2p.FederatedID), f.Peer2PeerToken, !f.RandomWorker, f.TargetWorker) fs := p2p.NewFederatedServer(f.Address, p2p.NetworkID(f.Peer2PeerNetworkID, p2p.FederatedID), f.Peer2PeerToken, !f.RandomWorker, f.TargetWorker)
return fs.Start(context.Background(), f.P2PCommonFlags) return fs.Start(context.Background(), f.P2PCommonFlags)

View file

@ -138,11 +138,6 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error {
} }
opts = append(opts, config.WithP2PToken(token)) opts = append(opts, config.WithP2PToken(token))
if r.Federated {
p2pCfg.PeerGuard.Autocleanup = true
p2pCfg.PeerGuard.PeerGate = true
}
p2pCfg.NetworkToken = token p2pCfg.NetworkToken = token
} }

View file

@ -15,9 +15,9 @@ import (
"github.com/rs/zerolog/log" "github.com/rs/zerolog/log"
) )
func (f *FederatedServer) Start(ctx context.Context, p2pCommonFlags cliP2P.P2PCommonFlags) error { func (fs *FederatedServer) Start(ctx context.Context, p2pCommonFlags cliP2P.P2PCommonFlags) error {
p2pCfg := NewP2PConfig(p2pCommonFlags) p2pCfg := NewP2PConfig(p2pCommonFlags)
p2pCfg.NetworkToken = f.p2ptoken p2pCfg.NetworkToken = fs.p2ptoken
p2pCfg.PeerGuard.Autocleanup = true p2pCfg.PeerGuard.Autocleanup = true
p2pCfg.PeerGuard.PeerGate = true p2pCfg.PeerGuard.PeerGate = true
@ -30,13 +30,13 @@ func (f *FederatedServer) Start(ctx context.Context, p2pCommonFlags cliP2P.P2PCo
return fmt.Errorf("creating a new node: %w", err) return fmt.Errorf("creating a new node: %w", err)
} }
if err := ServiceDiscoverer(ctx, n, f.service, func(servicesID string, tunnel NodeData) { if err := ServiceDiscoverer(ctx, n, fs.service, func(servicesID string, tunnel NodeData) {
log.Debug().Msgf("Discovered node: %s", tunnel.ID) log.Debug().Msgf("Discovered node: %s", tunnel.ID)
}, false); err != nil { }, false); err != nil {
return err return err
} }
return f.proxy(ctx, n) return fs.proxy(ctx, n)
} }
func (fs *FederatedServer) proxy(ctx context.Context, node *node.Node) error { func (fs *FederatedServer) proxy(ctx context.Context, node *node.Node) error {