diff --git a/core/cli/federated.go b/core/cli/federated.go index 92c84c61..8e5cfdcc 100644 --- a/core/cli/federated.go +++ b/core/cli/federated.go @@ -2,10 +2,12 @@ package cli import ( "context" + "fmt" cliContext "github.com/mudler/LocalAI/core/cli/context" cliP2P "github.com/mudler/LocalAI/core/cli/p2p" "github.com/mudler/LocalAI/core/p2p" + "github.com/rs/zerolog/log" ) type FederatedCLI struct { @@ -18,6 +20,24 @@ type FederatedCLI struct { 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) return fs.Start(context.Background(), f.P2PCommonFlags) diff --git a/core/cli/run.go b/core/cli/run.go index 5e39db05..d058f013 100644 --- a/core/cli/run.go +++ b/core/cli/run.go @@ -138,11 +138,6 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error { } opts = append(opts, config.WithP2PToken(token)) - if r.Federated { - p2pCfg.PeerGuard.Autocleanup = true - p2pCfg.PeerGuard.PeerGate = true - } - p2pCfg.NetworkToken = token } diff --git a/core/p2p/federated_server.go b/core/p2p/federated_server.go index de696abb..081bd134 100644 --- a/core/p2p/federated_server.go +++ b/core/p2p/federated_server.go @@ -15,9 +15,9 @@ import ( "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.NetworkToken = f.p2ptoken + p2pCfg.NetworkToken = fs.p2ptoken p2pCfg.PeerGuard.Autocleanup = 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) } - 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) }, false); err != nil { return err } - return f.proxy(ctx, n) + return fs.proxy(ctx, n) } func (fs *FederatedServer) proxy(ctx context.Context, node *node.Node) error {