JCC-CSScheduler/schedulerMiddleware/internal/manager/jobmgr/job/state2/job_create.go

package state2

import (
	"fmt"
	"gitlink.org.cn/cloudream/common/pkgs/logger"
	sch "gitlink.org.cn/cloudream/common/sdks/pcmscheduler"
	schsdk "gitlink.org.cn/cloudream/common/sdks/scheduler"
	cdssdk "gitlink.org.cn/cloudream/common/sdks/storage"
	schglb "gitlink.org.cn/cloudream/scheduler/common/globals"
	jobmod "gitlink.org.cn/cloudream/scheduler/common/models/job"
	"gitlink.org.cn/cloudream/scheduler/schedulerMiddleware/internal/manager/jobmgr"
	"gitlink.org.cn/cloudream/scheduler/schedulerMiddleware/internal/manager/jobmgr/job/state"
)

type PCMJobCreate struct {
	jobInfo         *schsdk.PCMJobInfo
	blockChainToken string
	userID          cdssdk.UserID
}

func NewPCMJobCreate(userID cdssdk.UserID, info *schsdk.PCMJobInfo, blockChainToken string) *PCMJobCreate {
	return &PCMJobCreate{
		userID:          userID,
		jobInfo:         info,
		blockChainToken: blockChainToken,
	}
}

type UploadedData struct {
}

type CodeDistribute struct {
}

func (s *PCMJobCreate) Run(rtx jobmgr.JobStateRunContext, jo *jobmgr.Job) {
	scheduleData, err := s.do(rtx)
	if err != nil {
		logger.Error(err.Error())
		rtx.Mgr.ChangeState(jo, state.FailureComplete(err))
	} else {
		rtx.Mgr.ChangeState(jo, NewDataSchedule(s.userID, scheduleData.TaskID, scheduleData.ScheduleDatas, s.blockChainToken))
	}
}

// 根据数据分布情况和资源需求创建任务
func (s *PCMJobCreate) do(rtx jobmgr.JobStateRunContext) (*sch.CreateJobResp, error) {
	schCli, err := schglb.PCMSchePool.Acquire()
	if err != nil {
		return nil, fmt.Errorf("new scheduler client: %w", err)
	}
	defer schglb.PCMSchePool.Release(schCli)

	// 获取所有packageID
	var packages []cdssdk.PackageID
	var bindingIDs []int64
	packages, bindingIDs = collectDataID(s.jobInfo.Files.Code, packages, bindingIDs)
	packages, bindingIDs = collectDataID(s.jobInfo.Files.Dataset, packages, bindingIDs)
	packages, bindingIDs = collectDataID(s.jobInfo.Files.Image, packages, bindingIDs)
	packages, bindingIDs = collectDataID(s.jobInfo.Files.Model, packages, bindingIDs)

	if len(packages) == 0 && len(bindingIDs) == 0 {
		return nil, fmt.Errorf("no packageID")
	}

	// 从数据库中读取数据信息
	uploadDatas, err := rtx.Mgr.DB.UploadData().GetByPackageID(rtx.Mgr.DB.DefCtx(), packages, bindingIDs)
	if err != nil {
		return nil, err
	}
	if len(uploadDatas) == 0 {
		return nil, fmt.Errorf("no upload data")
	}

	// 获取数据分布信息
	var dataDistribute sch.DataDistribute
	for _, data := range uploadDatas {
		var clusters []sch.DataDetail
		for _, cluster := range data.UploadedCluster {
			clusters = append(clusters, sch.DataDetail{
				ClusterID: cluster.ClusterID,
				JsonData:  data.JsonData,
			})
		}

		switch data.DataType {
		case sch.CODE:
			dataDistribute.Code = append(dataDistribute.Code, sch.CodeDistribute{
				Clusters:  clusters,
				PackageID: data.PackageID,
				DataName:  "test",
			})
		case sch.DATASET:
			dataDistribute.Dataset = append(dataDistribute.Dataset, sch.DatasetDistribute{
				Clusters:  clusters,
				PackageID: data.PackageID,
				DataName:  "test",
			})
		case sch.MODEL:
			dataDistribute.Model = append(dataDistribute.Model, sch.ModelDistribute{
				Clusters:  clusters,
				PackageID: data.PackageID,
				DataName:  "test",
			})
		case sch.IMAGE:
			dataDistribute.Image = append(dataDistribute.Image, sch.ImageDistribute{
				Clusters:  clusters,
				PackageID: data.PackageID,
				DataName:  "test",
			})
		}
	}

	req := sch.CreateJobReq{
		Name:           s.jobInfo.Name,
		Description:    s.jobInfo.Description,
		DataDistribute: dataDistribute,
		JobResources:   s.jobInfo.JobResources,
	}

	resp, err := schCli.CreateJob(req)
	if err != nil {
		return nil, fmt.Errorf("create task: %w", err)
	}

	return resp, nil
}

func collectDataID(fileInfo schsdk.JobFileInfo, packageIDs []cdssdk.PackageID, bindingIDs []int64) ([]cdssdk.PackageID, []int64) {
	switch info := fileInfo.(type) {
	case *schsdk.PackageJobFileInfo:
		packageIDs = append(packageIDs, info.PackageID)
	case *schsdk.BindingJobFileInfo:
		bindingIDs = append(bindingIDs, info.BindingID)
	case *schsdk.ImageJobFileInfo:
		packageIDs = append(packageIDs, cdssdk.PackageID(info.ImageID))
	}
	return packageIDs, bindingIDs
}

func (s *PCMJobCreate) Dump(rtx jobmgr.JobStateRunContext, jo *jobmgr.Job) jobmod.JobStateDump {
	return &jobmod.NormalJobReadyToExecuteDump{}
}