File size: 1,200 Bytes
0f07ba7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
package p2p
import (
"slices"
"strings"
"sync"
"github.com/mudler/LocalAI/core/schema"
)
const (
defaultServicesID = "services"
WorkerID = "worker"
)
var mu sync.Mutex
var nodes = map[string]map[string]schema.NodeData{}
func GetAvailableNodes(serviceID string) []schema.NodeData {
if serviceID == "" {
serviceID = defaultServicesID
}
mu.Lock()
defer mu.Unlock()
var availableNodes = []schema.NodeData{}
for _, v := range nodes[serviceID] {
availableNodes = append(availableNodes, v)
}
slices.SortFunc(availableNodes, func(a, b schema.NodeData) int {
return strings.Compare(a.ID, b.ID)
})
return availableNodes
}
func GetNode(serviceID, nodeID string) (schema.NodeData, bool) {
if serviceID == "" {
serviceID = defaultServicesID
}
mu.Lock()
defer mu.Unlock()
if _, ok := nodes[serviceID]; !ok {
return schema.NodeData{}, false
}
nd, exists := nodes[serviceID][nodeID]
return nd, exists
}
func AddNode(serviceID string, node schema.NodeData) {
if serviceID == "" {
serviceID = defaultServicesID
}
mu.Lock()
defer mu.Unlock()
if nodes[serviceID] == nil {
nodes[serviceID] = map[string]schema.NodeData{}
}
nodes[serviceID][node.ID] = node
}
|