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

package schedule

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

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

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

func NewScheduleCreateTaskLogic(ctx context.Context, svcCtx *svc.ServiceContext) *ScheduleCreateTaskLogic {
	return &ScheduleCreateTaskLogic{
		Logger: logx.WithContext(ctx),
		ctx:    ctx,
		svcCtx: 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 := generateScheduleResult(req.DataDistributes, req.Clusters)
		taskId := l.createTask()
		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 := generateScheduleResult(req.DataDistributes, clusters)
		taskId := l.createTask()
		resp.ScheduleDatas = schedatas
		resp.TaskID = taskId
		return resp, nil
	}

	if len(req.Resources) != 0 {

	}
	return
}

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

	for _, cluster := range clusters {
		id, _ := strconv.ParseInt(cluster, 10, 64)
		for _, dd := range distribute.Dataset {
			if slices.Contains(dd.Clusters, id) {
				data := &types.ScheduleData{
					DataType:   "dataset",
					PackageID:  dd.PackageID,
					ClusterIDs: dd.Clusters,
				}
				schedatas = append(schedatas, data)
			}
		}

		for _, dd := range distribute.Code {
			if slices.Contains(dd.Clusters, id) {
				data := &types.ScheduleData{
					DataType:   "code",
					PackageID:  dd.PackageID,
					ClusterIDs: dd.Clusters,
				}
				schedatas = append(schedatas, data)
			}
		}

		for _, dd := range distribute.Image {
			if slices.Contains(dd.Clusters, id) {
				data := &types.ScheduleData{
					DataType:   "image",
					PackageID:  dd.PackageID,
					ClusterIDs: dd.Clusters,
				}
				schedatas = append(schedatas, data)
			}
		}

		for _, dd := range distribute.Model {
			if slices.Contains(dd.Clusters, id) {
				data := &types.ScheduleData{
					DataType:   "model",
					PackageID:  dd.PackageID,
					ClusterIDs: dd.Clusters,
				}
				schedatas = append(schedatas, data)
			}
		}
	}

	return schedatas
}

func (l *ScheduleCreateTaskLogic) createTask() int64 {
	return 123456789
}