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.Executor, *map[string]collector.ResourceCollector) {
	executorMap := make(map[string]executor.Executor)
	collectorMap := make(map[string]collector.ResourceCollector)
	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
}