chore: allow to disable gallery endpoints, improve p2p connection handling (#3256)

* Add more debug messages

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* feat: allow to disable gallery endpoints

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* improve p2p messaging

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* improve error handling

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* Make sure to close the listening socket when context is exhausted

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

---------

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
Ettore Di Giacinto 2024-08-17 08:28:52 +02:00 committed by GitHub
parent d6b3fbb4ad
commit 7278bf3de8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 44 additions and 35 deletions

View file

@ -80,7 +80,7 @@ func (fs *FederatedServer) SelectLeastUsedServer() string {
fs.Lock()
defer fs.Unlock()
log.Debug().Any("request_table", fs.requestTable).Msgf("Current request table")
log.Debug().Any("request_table", fs.requestTable).Msgf("SelectLeastUsedServer()")
// cycle over requestTable and find the entry with the lower number
// if there are multiple entries with the same number, select one randomly
@ -93,7 +93,7 @@ func (fs *FederatedServer) SelectLeastUsedServer() string {
minKey = k
}
}
log.Debug().Any("requests_served", min).Msgf("Selected tunnel %s", minKey)
log.Debug().Any("requests_served", min).Any("request_table", fs.requestTable).Msgf("Selected tunnel %s", minKey)
return minKey
}
@ -104,7 +104,7 @@ func (fs *FederatedServer) RecordRequest(nodeID string) {
// increment the counter for the nodeID in the requestTable
fs.requestTable[nodeID]++
log.Debug().Any("request_table", fs.requestTable).Msgf("Current request table")
log.Debug().Any("request_table", fs.requestTable).Any("request", nodeID).Msgf("Recording request")
}
func (fs *FederatedServer) ensureRecordExist(nodeID string) {
@ -114,5 +114,5 @@ func (fs *FederatedServer) ensureRecordExist(nodeID string) {
fs.requestTable[nodeID] = 0
}
log.Debug().Any("request_table", fs.requestTable).Msgf("Current request table")
log.Debug().Any("request_table", fs.requestTable).Any("request", nodeID).Msgf("Ensure record exists")
}

View file

@ -46,7 +46,10 @@ func (fs *FederatedServer) proxy(ctx context.Context, node *node.Node) error {
return err
}
// ll.Info("Binding local port on", srcaddr)
go func() {
<-ctx.Done()
l.Close()
}()
ledger, _ := node.Ledger()
// Announce ourselves so nodes accepts our connection
@ -54,17 +57,12 @@ func (fs *FederatedServer) proxy(ctx context.Context, node *node.Node) error {
ctx,
10*time.Second,
func() {
// Retrieve current ID for ip in the blockchain
//_, found := ledger.GetKey(protocol.UsersLedgerKey, node.Host().ID().String())
// If mismatch, update the blockchain
//if !found {
updatedMap := map[string]interface{}{}
updatedMap[node.Host().ID().String()] = &types.User{
PeerID: node.Host().ID().String(),
Timestamp: time.Now().String(),
}
ledger.Add(protocol.UsersLedgerKey, updatedMap)
// }
},
)

View file

@ -51,6 +51,11 @@ func allocateLocalService(ctx context.Context, node *node.Node, listenAddr, serv
zlog.Error().Err(err).Msg("Error listening")
return err
}
go func() {
<-ctx.Done()
l.Close()
}()
// ll.Info("Binding local port on", srcaddr)
ledger, _ := node.Ledger()
@ -60,17 +65,12 @@ func allocateLocalService(ctx context.Context, node *node.Node, listenAddr, serv
ctx,
10*time.Second,
func() {
// Retrieve current ID for ip in the blockchain
//_, found := ledger.GetKey(protocol.UsersLedgerKey, node.Host().ID().String())
// If mismatch, update the blockchain
//if !found {
updatedMap := map[string]interface{}{}
updatedMap[node.Host().ID().String()] = &types.User{
PeerID: node.Host().ID().String(),
Timestamp: time.Now().String(),
}
ledger.Add(protocol.UsersLedgerKey, updatedMap)
// }
},
)
@ -197,14 +197,13 @@ func discoveryTunnels(ctx context.Context, n *node.Node, token, servicesID strin
return
default:
time.Sleep(5 * time.Second)
zlog.Debug().Msg("Searching for workers")
data := ledger.LastBlock().Storage[servicesID]
zlog.Debug().Any("data", ledger.LastBlock().Storage).Msg("Ledger data")
for k, v := range data {
zlog.Info().Msgf("Found worker %s", k)
zlog.Debug().Msgf("New worker found in the ledger data '%s'", k)
nd := &NodeData{}
if err := v.Unmarshal(nd); err != nil {
zlog.Error().Msg("cannot unmarshal node data")
@ -245,8 +244,10 @@ func ensureService(ctx context.Context, n *node.Node, nd *NodeData, sserv string
// Start the service
port, err := freeport.GetFreePort()
if err != nil {
fmt.Print(err)
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
service[nd.Name] = nodeServiceData{
@ -310,10 +311,6 @@ func ExposeService(ctx context.Context, host, port, token, servicesID string) er
ctx,
20*time.Second,
func() {
// Retrieve current ID for ip in the blockchain
//_, found := ledger.GetKey("services_localai", name)
// If mismatch, update the blockchain
//if !found {
updatedMap := map[string]interface{}{}
updatedMap[name] = &NodeData{
Name: name,
@ -321,7 +318,6 @@ func ExposeService(ctx context.Context, host, port, token, servicesID string) er
ID: nodeID(name),
}
ledger.Add(servicesID, updatedMap)
// }
},
)
@ -354,7 +350,10 @@ func newNodeOpts(token string) ([]node.Option, error) {
if loglevel == "" {
loglevel = "info"
}
libp2ploglevel := os.Getenv("LOCALAI_LIBP2P_LOGLEVEL")
if libp2ploglevel == "" {
libp2ploglevel = "info"
}
c := config.Config{
Limit: config.ResourceLimit{
Enable: noLimits,
@ -363,7 +362,7 @@ func newNodeOpts(token string) ([]node.Option, error) {
NetworkToken: token,
LowProfile: false,
LogLevel: loglevel,
Libp2pLogLevel: "fatal",
Libp2pLogLevel: libp2ploglevel,
Ledger: config.Ledger{
SyncInterval: defaultInterval,
AnnounceInterval: defaultInterval,