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

package service

import (
	"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/jcce-pcm/pcm-ac/hpcacclient"
	"gitlink.org.cn/jcce-pcm/pcm-participant-modelarts/client/imagesservice"
	"gitlink.org.cn/jcce-pcm/pcm-participant-modelarts/client/modelartsservice"
	"gitlink.org.cn/jcce-pcm/pcm-participant-octopus/octopusclient"
)

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

var (
	AiTypeMap = map[string]string{
		"hanwuji": OCTOPUS,
		//"suiyan":                OCTOPUS,
		//"sailingsi":             OCTOPUS,
		//"modelarts-CloudBrain2": MODELARTS,
		"shuguangAi": SHUGUANGAI,
	}
)

func InitAiClusterMap(octopusRpc octopusclient.Octopus, modelArtsRpc modelartsservice.ModelArtsService, modelArtsImgRpc imagesservice.ImagesService, aCRpc hpcacclient.HpcAC) (*map[string]executor.AiExecutor, *map[string]collector.AiCollector) {
	executorMap := make(map[string]executor.AiExecutor)
	collectorMap := make(map[string]collector.AiCollector)
	for k, v := range AiTypeMap {
		switch v {
		case OCTOPUS:
			octopus := storeLink.NewOctopusLink(octopusRpc, k, 0)
			collectorMap[k] = octopus
			executorMap[k] = octopus
		case MODELARTS:
			modelarts := storeLink.NewModelArtsLink(modelArtsRpc, modelArtsImgRpc, k, 0)
			collectorMap[k] = modelarts
			executorMap[k] = modelarts
		case SHUGUANGAI:
			sgai := storeLink.NewShuguangAi(aCRpc, k, 0)
			collectorMap[k] = sgai
			executorMap[k] = sgai
		}
	}
	return &executorMap, &collectorMap
}