forked from JointCloud/pcm-coordinator
111 lines
3.9 KiB
Go
111 lines
3.9 KiB
Go
package service
|
|
|
|
import (
|
|
"github.com/zeromicro/go-zero/core/logx"
|
|
"github.com/zeromicro/go-zero/zrpc"
|
|
hpcacclient "gitlink.org.cn/JointCloud/pcm-ac/hpcacclient"
|
|
"gitlink.org.cn/JointCloud/pcm-coordinator/api/internal/config"
|
|
"gitlink.org.cn/JointCloud/pcm-coordinator/api/internal/scheduler/database"
|
|
"gitlink.org.cn/JointCloud/pcm-coordinator/api/internal/scheduler/schedulers/option"
|
|
"gitlink.org.cn/JointCloud/pcm-coordinator/api/internal/scheduler/service/collector"
|
|
"gitlink.org.cn/JointCloud/pcm-coordinator/api/internal/scheduler/service/executor"
|
|
"gitlink.org.cn/JointCloud/pcm-coordinator/api/internal/scheduler/service/inference"
|
|
"gitlink.org.cn/JointCloud/pcm-coordinator/api/internal/scheduler/strategy"
|
|
"gitlink.org.cn/JointCloud/pcm-coordinator/api/internal/storeLink"
|
|
"gitlink.org.cn/JointCloud/pcm-coordinator/api/internal/types"
|
|
"gitlink.org.cn/JointCloud/pcm-modelarts/client/imagesservice"
|
|
"gitlink.org.cn/JointCloud/pcm-modelarts/client/modelartsservice"
|
|
"gitlink.org.cn/JointCloud/pcm-octopus/octopusclient"
|
|
"net/http"
|
|
"strconv"
|
|
"sync"
|
|
)
|
|
|
|
const (
|
|
OCTOPUS = "octopus"
|
|
MODELARTS = "modelarts"
|
|
SHUGUANGAI = "shuguangAi"
|
|
)
|
|
|
|
type AiService struct {
|
|
AiExecutorAdapterMap map[string]map[string]executor.AiExecutor
|
|
AiCollectorAdapterMap map[string]map[string]collector.AiCollector
|
|
Storage *database.AiStorage
|
|
mu sync.Mutex
|
|
}
|
|
|
|
func NewAiService(conf *config.Config, storages *database.AiStorage) (*AiService, error) {
|
|
var aiType = "1"
|
|
adapterIds, err := storages.GetAdapterIdsByType(aiType)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
aiService := &AiService{
|
|
AiExecutorAdapterMap: make(map[string]map[string]executor.AiExecutor),
|
|
AiCollectorAdapterMap: make(map[string]map[string]collector.AiCollector),
|
|
Storage: storages,
|
|
}
|
|
for _, id := range adapterIds {
|
|
clusters, err := storages.GetClustersByAdapterId(id)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if len(clusters.List) == 0 {
|
|
continue
|
|
}
|
|
exeClusterMap, colClusterMap := InitAiClusterMap(conf, clusters.List)
|
|
aiService.AiExecutorAdapterMap[id] = exeClusterMap
|
|
aiService.AiCollectorAdapterMap[id] = colClusterMap
|
|
}
|
|
|
|
return aiService, nil
|
|
}
|
|
|
|
func InitAiClusterMap(conf *config.Config, clusters []types.ClusterInfo) (map[string]executor.AiExecutor, map[string]collector.AiCollector) {
|
|
executorMap := make(map[string]executor.AiExecutor)
|
|
collectorMap := make(map[string]collector.AiCollector)
|
|
for _, c := range clusters {
|
|
switch c.Name {
|
|
case OCTOPUS:
|
|
id, _ := strconv.ParseInt(c.Id, 10, 64)
|
|
octopusRpc := octopusclient.NewOctopus(zrpc.MustNewClient(conf.OctopusRpcConf))
|
|
octopus := storeLink.NewOctopusLink(octopusRpc, c.Nickname, id)
|
|
collectorMap[c.Id] = octopus
|
|
executorMap[c.Id] = octopus
|
|
case MODELARTS:
|
|
id, _ := strconv.ParseInt(c.Id, 10, 64)
|
|
modelArtsRpc := modelartsservice.NewModelArtsService(zrpc.MustNewClient(conf.ModelArtsRpcConf))
|
|
modelArtsImgRpc := imagesservice.NewImagesService(zrpc.MustNewClient(conf.ModelArtsImgRpcConf))
|
|
modelarts := storeLink.NewModelArtsLink(modelArtsRpc, modelArtsImgRpc, c.Name, id, c.Nickname)
|
|
collectorMap[c.Id] = modelarts
|
|
executorMap[c.Id] = modelarts
|
|
case SHUGUANGAI:
|
|
id, _ := strconv.ParseInt(c.Id, 10, 64)
|
|
aCRpc := hpcacclient.NewHpcAC(zrpc.MustNewClient(conf.ACRpcConf))
|
|
sgai := storeLink.NewShuguangAi(aCRpc, c.Nickname, id)
|
|
collectorMap[c.Id] = sgai
|
|
executorMap[c.Id] = sgai
|
|
}
|
|
}
|
|
|
|
return executorMap, collectorMap
|
|
}
|
|
|
|
func (as *AiService) ImageInfer(opt *option.InferOption, id int64, adapterName string, clusters []*strategy.AssignedCluster, ts []*inference.ImageFile) {
|
|
|
|
r := http.Request{}
|
|
_, err := inference.Infer(opt, id, adapterName, clusters, ts, as.AiCollectorAdapterMap, as.Storage, r.Context())
|
|
if err != nil {
|
|
logx.Errorf(err.Error())
|
|
return
|
|
}
|
|
}
|
|
|
|
//func (a *AiService) AddCluster() error {
|
|
//
|
|
//}
|
|
//
|
|
//func (a *AiService) AddAdapter() error {
|
|
//
|
|
//}
|