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

package service

import (
	"context"
	"gitlink.org.cn/jcce-pcm/pcm-coordinator/api/internal/scheduler/service/collector"
	"gitlink.org.cn/jcce-pcm/pcm-coordinator/api/internal/scheduler/service/executor"
	"gitlink.org.cn/jcce-pcm/pcm-coordinator/api/internal/storeLink"
	"gitlink.org.cn/jcce-pcm/pcm-coordinator/api/internal/svc"
)

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(ctx context.Context, svcCtx *svc.ServiceContext) (*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(ctx, svcCtx, k, 0)
			collectorMap[k] = octopus
			executorMap[k] = octopus
		case MODELARTS:
			modelarts := storeLink.NewModelArtsLink(ctx, svcCtx, k, 0)
			collectorMap[k] = modelarts
			executorMap[k] = modelarts
		case SHUGUANGAI:
			sgai := storeLink.NewShuguangAi(ctx, svcCtx, k, 0)
			collectorMap[k] = sgai
			executorMap[k] = sgai
		}
	}
	return &executorMap, &collectorMap
}