pcm-coordinator/internal/scheduler/service/hpc_service.go

package service

import (
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/config"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/scheduler/database"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/scheduler/service/collector"
	hpcservice "gitlink.org.cn/JointCloud/pcm-coordinator/internal/scheduler/service/hpc"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/types"
	"strconv"
	"sync"
)

const (
	Slurm_Arm = "slurm_arm"
)

type HpcService struct {
	HpcExecutorAdapterMap map[string]map[string]collector.HPCCollector
	Storage               *database.HpcStorage
	LocalCache            map[string]interface{}
	Conf                  *config.Config
	TaskSyncLock          sync.Mutex
}

func NewHpcService(conf *config.Config, storages *database.HpcStorage, localCache map[string]interface{}) (*HpcService, error) {
	var aiType = "2"
	adapterIds, err := storages.GetAdapterIdsByType(aiType)
	if err != nil {
		return nil, err
	}
	hpcService := &HpcService{
		HpcExecutorAdapterMap: make(map[string]map[string]collector.HPCCollector),
		Storage:               storages,
		LocalCache:            localCache,
		Conf:                  conf,
	}
	for _, id := range adapterIds {
		clusters, err := storages.GetClustersByAdapterId(id)
		if err != nil {
			return nil, err
		}
		if len(clusters.List) == 0 {
			continue
		}
		exeClusterMap := InitHpcClusterMap(conf, clusters.List)
		hpcService.HpcExecutorAdapterMap[id] = exeClusterMap
	}

	return hpcService, nil
}

func InitHpcClusterMap(conf *config.Config, clusters []types.ClusterInfo) map[string]collector.HPCCollector {
	executorMap := make(map[string]collector.HPCCollector)
	for _, c := range clusters {
		switch c.Name {
		case Slurm_Arm:
			id, _ := strconv.ParseInt(c.Id, 10, 64)
			slurm := hpcservice.NewHpc(c.Server, id, c.Nickname)
			executorMap[c.Id] = slurm
		}
	}
	return executorMap
}

func (as *HpcService) UpdateHpcClusterMaps(conf *config.Config, adapterId string, clusters []types.ClusterInfo) {
	for _, c := range clusters {
		_, ok := as.HpcExecutorAdapterMap[adapterId][c.Id]
		if !ok {
			switch c.Name {
			case Slurm_Arm:
				id, _ := strconv.ParseInt(c.Id, 10, 64)
				slurm := hpcservice.NewHpc(c.Server, id, c.Nickname)
				as.HpcExecutorAdapterMap[adapterId][c.Id] = slurm
			}
		} else {
			continue
		}
	}
}