pcm-coordinator/internal/logic/schedule/schedulecreatetasklogic.go

package schedule

import (
	"context"
	"fmt"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/scheduler/common"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/scheduler/strategy"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/storeLink"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/svc"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/types"
	"gitlink.org.cn/JointCloud/pcm-coordinator/pkg/utils"
	"slices"
	"strconv"
	"strings"

	"github.com/zeromicro/go-zero/core/logx"
)

type ScheduleCreateTaskLogic struct {
	logx.Logger
	ctx           context.Context
	svcCtx        *svc.ServiceContext
	queryResource *QueryResourcesLogic
}

func NewScheduleCreateTaskLogic(ctx context.Context, svcCtx *svc.ServiceContext) *ScheduleCreateTaskLogic {
	return &ScheduleCreateTaskLogic{
		Logger:        logx.WithContext(ctx),
		ctx:           ctx,
		svcCtx:        svcCtx,
		queryResource: NewQueryResourcesLogic(ctx, svcCtx),
	}
}

func (l *ScheduleCreateTaskLogic) ScheduleCreateTask(req *types.CreateTaskReq) (resp *types.CreateTaskResp, err error) {
	resp = &types.CreateTaskResp{}
	if len(req.Resources) == 0 && len(req.Clusters) == 0 && len(req.StaticWeight) == 0 {
		return nil, fmt.Errorf("must specify at least one resource")
	}

	if len(req.Clusters) != 0 {
		schedatas, err := l.generateScheduleResult(req.DataDistributes, req.Clusters)
		if err != nil {
			return nil, err
		}
		taskId, err := l.createTask("SCHEDULE_TASK_"+utils.RandomString(10), "staticWeight")
		if err != nil {
			return nil, err
		}
		resp.ScheduleDatas = schedatas
		resp.TaskID = taskId
		return resp, nil
	}

	if len(req.StaticWeight) != 0 {
		strgy := strategy.NewStaticWeightStrategy(req.StaticWeight, 1)
		sche, err := strgy.Schedule()
		if err != nil {
			return nil, err
		}
		var clusters []string
		for _, c := range sche {
			if c.Replicas == 0 {
				continue
			}
			clusters = append(clusters, c.ClusterId)
		}
		schedatas, err := l.generateScheduleResult(req.DataDistributes, clusters)
		if err != nil {
			return nil, err
		}
		taskId, err := l.createTask("SCHEDULE_TASK_"+utils.RandomString(10), "staticWeight")
		if err != nil {
			return nil, err
		}
		resp.ScheduleDatas = schedatas
		resp.TaskID = taskId
		return resp, nil
	}

	if len(req.Resources) != 0 {
		_, err := l.queryResource.QueryResources(&types.QueryResourcesReq{AdapterId: req.AdapterId})
		if err != nil {
			return nil, err
		}
		return
	}
	return
}

func (l *ScheduleCreateTaskLogic) generateScheduleResult(distribute types.DataDistribute, clusters []string) ([]*types.ScheduleData, error) {
	var schedatas []*types.ScheduleData

	for _, d := range distribute.Dataset {
		data := &types.ScheduleData{
			DataType:   "dataset",
			PackageID:  d.PackageID,
			ClusterIDs: make([]int64, 0),
		}
		for _, cluster := range clusters {
			id, _ := strconv.ParseInt(cluster, 10, 64)
			if !slices.Contains(d.Clusters, id) {
				data.ClusterIDs = append(data.ClusterIDs, id)
			} else {
				continue
			}
		}
		if len(data.ClusterIDs) != 0 {
			schedatas = append(schedatas, data)
		}
	}

	for _, d := range distribute.Code {
		data := &types.ScheduleData{
			DataType:   "code",
			PackageID:  d.PackageID,
			ClusterIDs: make([]int64, 0),
		}
		for _, cluster := range clusters {
			id, _ := strconv.ParseInt(cluster, 10, 64)
			if !slices.Contains(d.Clusters, id) {
				data.ClusterIDs = append(data.ClusterIDs, id)
			} else {
				continue
			}
		}
		if len(data.ClusterIDs) != 0 {
			schedatas = append(schedatas, data)
		}
	}

	for _, d := range distribute.Image {
		data := &types.ScheduleData{
			DataType:   "image",
			PackageID:  d.PackageID,
			ClusterIDs: make([]int64, 0),
		}
		for _, cluster := range clusters {
			id, _ := strconv.ParseInt(cluster, 10, 64)
			if !slices.Contains(d.Clusters, id) {
				data.ClusterIDs = append(data.ClusterIDs, id)
			} else {
				continue
			}
		}
		if len(data.ClusterIDs) != 0 {
			schedatas = append(schedatas, data)
		}
	}

	for _, d := range distribute.Model {
		data := &types.ScheduleData{
			DataType:   "model",
			PackageID:  d.PackageID,
			ClusterIDs: make([]int64, 0),
		}
		for _, cluster := range clusters {
			id, _ := strconv.ParseInt(cluster, 10, 64)
			if !slices.Contains(d.Clusters, id) {
				data.ClusterIDs = append(data.ClusterIDs, id)
			} else {
				continue
			}
		}
		if len(data.ClusterIDs) != 0 {
			schedatas = append(schedatas, data)
		}
	}

	if len(schedatas) != 0 {
		err := l.updateStorageType(&schedatas)
		if err != nil {
			return nil, err
		}
	}

	return schedatas, nil
}

func (l *ScheduleCreateTaskLogic) createTask(taskName string, strategyName string) (int64, error) {
	strategyCode, err := l.svcCtx.Scheduler.AiStorages.GetStrategyCode(strategyName)
	if err != nil {
		return 0, err
	}
	//adapterName, err := l.svcCtx.Scheduler.AiStorages.GetAdapterNameById(rs[0].AdapterId)
	//if err != nil {
	//	return nil, err
	//}
	id, err := l.svcCtx.Scheduler.AiStorages.SaveTask(taskName, strategyCode, 0, "10")
	if err != nil {
		return 0, err
	}
	return id, nil
}

func (l *ScheduleCreateTaskLogic) updateStorageType(schedatas *[]*types.ScheduleData) error {

	for _, s := range *schedatas {
		var storageType string
		var sTypes []string
		for _, i := range s.ClusterIDs {
			id := strconv.FormatInt(i, 10)
			cluster, err := l.svcCtx.Scheduler.AiStorages.GetClustersById(id)
			if err != nil {
				return err
			}
			stype, ok := storeLink.StorageTypeMap[strings.Title(cluster.Name)]
			if ok {
				sTypes = append(sTypes, stype)
			}
		}
		sTypes = common.Unique(sTypes)
		for _, st := range sTypes {
			storageType += st + storeLink.COMMA
		}
		storageType = strings.TrimSuffix(storageType, storeLink.COMMA)
		s.StorageType = storageType
	}

	return nil
}