pcm-coordinator/api/internal/scheduler/service/aiService.go

package service

import (
	"github.com/zeromicro/go-zero/zrpc"
	"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/service/collector"
	"gitlink.org.cn/JointCloud/pcm-coordinator/api/internal/scheduler/service/executor"
	"gitlink.org.cn/JointCloud/pcm-coordinator/api/internal/storeLink"
	"gitlink.org.cn/JointCloud/pcm-coordinator/api/internal/types"
	"gitlink.org.cn/JointCloud/pcm-octopus/octopusclient"
	"gitlink.org.cn/jcce-pcm/pcm-participant-modelarts/client/imagesservice"
	"gitlink.org.cn/jcce-pcm/pcm-participant-modelarts/client/modelartsservice"
	"strconv"
)

const (
	OCTOPUS    = "octopus"
	MODELARTS  = "modelarts"
	SHUGUANGAI = "shuguangAi"
)

type AiService struct {
	AiExecutorAdapterMap  map[string]map[string]executor.AiExecutor
	AiCollectorAdapterMap map[string]map[string]collector.AiCollector
}

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),
	}
	for _, id := range adapterIds {
		clusters, err := storages.GetClustersByAdapterId(id)
		if err != nil {
			return nil, err
		}
		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.Nickname, id)
			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
}