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

package service

import (
	"fmt"
	"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 (
	PcmSlurm = "pcm-slurm"
)

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

// NewHpcService 创建并初始化HpcService实例
func NewHpcService(conf *config.Config, storages *database.HpcStorage, localCache map[string]interface{}) (*HpcService, error) {
	hpcService := &HpcService{
		HpcExecutorAdapterMap: make(map[string]collector.HPCCollector),
		Storage:               storages,
		LocalCache:            localCache,
		Conf:                  conf,
	}

	if err := hpcService.initAdapters(); err != nil {
		return nil, err
	}

	return hpcService, nil
}

// initAdapters 初始化所有适配器
func (s *HpcService) initAdapters() error {
	adapters, err := s.loadAdapters()
	if err != nil {
		return err
	}
	for _, adapter := range adapters {
		if err := s.processAdapter(*adapter); err != nil {
			return err
		}
	}

	return nil
}

// loadAdapters 从存储中加载适配器
func (s *HpcService) loadAdapters() ([]*types.AdapterInfo, error) {
	const aiType = "2"
	return s.Storage.GetAdaptersByType(aiType)
}

// processAdapter 处理单个适配器
func (s *HpcService) processAdapter(adapter types.AdapterInfo) error {
	if adapter.Id == "" {
		return nil
	}

	executor, err := s.createExecutor(adapter)
	if err != nil {
		return err
	}

	if executor != nil {
		s.HpcExecutorAdapterMap[adapter.Id] = executor
	}

	return nil
}

// createExecutor 根据适配器类型创建对应的执行器
func (s *HpcService) createExecutor(adapter types.AdapterInfo) (collector.HPCCollector, error) {
	switch adapter.Nickname {
	case PcmSlurm:
		return s.CreateSlurmExecutor(adapter)
	// 可以在这里添加其他类型的适配器
	default:
		return nil, nil // 或者返回错误，取决于业务需求
	}
}

// CreateSlurmExecutor 创建Slurm执行器
func (s *HpcService) CreateSlurmExecutor(adapter types.AdapterInfo) (collector.HPCCollector, error) {
	id, err := strconv.ParseInt(adapter.Id, 10, 64)
	if err != nil {
		return nil, fmt.Errorf("failed to parse adapter ID %s: %v", adapter.Id, err)
	}
	return hpcservice.NewHpc(adapter.Server, id, adapter.Nickname), nil
}