pcm-octopus/internal/logic/getgeneralinfologic.go

package logic

import (
	"context"
	"encoding/json"
	"github.com/go-redis/redis"
	"github.com/zeromicro/go-zero/core/logx"
	"gitlink.org.cn/JointCloud/pcm-octopus/internal/common"
	"gitlink.org.cn/JointCloud/pcm-octopus/internal/config"
	"gitlink.org.cn/JointCloud/pcm-octopus/internal/svc"
	"gitlink.org.cn/JointCloud/pcm-octopus/octopus"
	"log"
	"strconv"
	"strings"
	"time"
)

const (
	GI                 = "Gi"
	OctopusGeneralInfo = "octopusGeneralInfo"
	Comma              = ","
)

type GetGeneralInfoLogic struct {
	ctx    context.Context
	svcCtx *svc.ServiceContext
	logx.Logger
}

func NewGetGeneralInfoLogic(ctx context.Context, svcCtx *svc.ServiceContext) *GetGeneralInfoLogic {
	return &GetGeneralInfoLogic{
		ctx:    ctx,
		svcCtx: svcCtx,
		Logger: logx.WithContext(ctx),
	}
}

func (l *GetGeneralInfoLogic) GetGeneralInfo(in *octopus.ResourceReq) (*octopus.GiResp, error) {
	var resp octopus.GiResp
	var octopusCpuCores int32
	var octopusMemoryInGi int32

	redisClient := redis.NewClient(&redis.Options{
		Addr:     l.svcCtx.Config.RedisConf.Host,
		Password: l.svcCtx.Config.RedisConf.Pass,
	})

	defer redisClient.Close()
	_, err := redisClient.Ping().Result()

	if err != nil {
		log.Println("redis连接失败", err)
		octopusCpuCores, octopusMemoryInGi = getGeneralInfoFromOctopus()
		octopusGeneralInfo := strconv.FormatInt(int64(octopusCpuCores), 10) + Comma + strconv.FormatInt(int64(octopusMemoryInGi), 10)
		redisClient.Set(OctopusGeneralInfo, octopusGeneralInfo, 168*time.Hour)
	} else {
		res, err := redisClient.Get(OctopusGeneralInfo).Result()
		if err == redis.Nil {
			log.Println("redis key未找到或已过期，重新请求")
			octopusCpuCores, octopusMemoryInGi = getGeneralInfoFromOctopus()
			octopusGeneralInfo := strconv.FormatInt(int64(octopusCpuCores), 10) + Comma + strconv.FormatInt(int64(octopusMemoryInGi), 10)
			redisClient.Set(OctopusGeneralInfo, octopusGeneralInfo, 168*time.Hour)
		} else {
			strs := strings.Split(res, Comma)
			cpu, _ := strconv.ParseInt(strs[0], 10, 32)
			memory, _ := strconv.ParseInt(strs[1], 10, 32)
			octopusCpuCores = int32(cpu)
			octopusMemoryInGi = int32(memory)
		}
	}

	resp.CpuCoreNum = octopusCpuCores
	resp.MemoryInGib = octopusMemoryInGi
	return &resp, nil
}

func getGeneralInfoFromOctopus() (int32, int32) {
	octopusConfig := config.Cfg

	urlMap := map[string]string{
		common.Hanwuji:   octopusConfig.OctopusConfig.HanwujiUrl + octopusConfig.OctopusConfig.OctopusResouceSpec,
		common.Suiyuan:   octopusConfig.OctopusConfig.SuiyuanUrl + octopusConfig.OctopusConfig.OctopusResouceSpec,
		common.Sailingsi: octopusConfig.OctopusConfig.SailingsiUrl + octopusConfig.OctopusConfig.OctopusResouceSpec,
	}

	var cpuCoreNum int32
	var memoryInGib int32
	for k, v := range urlMap {
		token := common.GetToken(k)
		if token == "" {
			continue
		}
		body, err := common.OctopusHttpClient("GET", v, nil, token)
		if err != nil {
			continue
		}
		//获取训练资源算力
		switch k {
		case common.Hanwuji:
			resourceSpec := common.HanwujiResp{}
			err := json.Unmarshal(body, &resourceSpec)
			if err != nil {
				log.Println("Hanwuji json转换失败 : ", err)
				continue
			}
			if !resourceSpec.Success {
				log.Println("Hanwuji 获取训练资源失败 : ", resourceSpec.Error)
				continue
			}
			for _, spec := range resourceSpec.Payload.MapResourceSpecIdList.Train.ResourceSpecs {
				cpuInfo, err := strconv.ParseInt(spec.ResourceQuantity.Cpu, 10, 32)
				if err != nil {
					continue
				}
				cpuCoreNum += int32(cpuInfo)

				memoryStr := strings.Replace(spec.ResourceQuantity.Memory, GI, "", -1)
				memoryInfo, err := strconv.ParseInt(memoryStr, 10, 32)
				if err != nil {
					continue
				}
				memoryInGib += int32(memoryInfo)
			}
		case common.Suiyuan:
			resourceSpec := common.SuiyuanResp{}
			err := json.Unmarshal(body, &resourceSpec)
			if err != nil {
				log.Println("Suiyuan json转换失败 : ", err)
				continue
			}
			if !resourceSpec.Success {
				log.Println("Suiyuan 获取训练资源失败 : ", resourceSpec.Error)
				continue
			}
			for _, spec := range resourceSpec.Payload.MapResourceSpecIdList.Train.ResourceSpecs {
				cpuInfo, err := strconv.ParseInt(spec.ResourceQuantity.Cpu, 10, 32)
				if err != nil {
					continue
				}
				cpuCoreNum += int32(cpuInfo)

				memoryStr := strings.Replace(spec.ResourceQuantity.Memory, GI, "", -1)
				memoryInfo, err := strconv.ParseInt(memoryStr, 10, 32)
				if err != nil {
					continue
				}
				memoryInGib += int32(memoryInfo)
			}
		}

	}

	return cpuCoreNum, memoryInGib
}