remote-task-excutor-cli/pkg/service/common/task_builder.go

package common

import (
	"context"
	"fmt"
	"remote-task-excutor-cli/pkg/models"
	"remote-task-excutor-cli/pkg/service/preparation"
)

// TaskType 任务类型常量
const (
	TrainingTaskType  = "AI"            // 训练任务类型
	InferenceTaskType = "PCM_Inference" // 推理任务类型
)

func formatResource(resource *models.ResourceConfig) models.ResourceConfig {
	accerlerateItem := models.ResourcesItem{
		Type:   resource.Type,
		Name:   resource.Name,
		Number: resource.AvailableCount,
	}

	for i, _ := range resource.Resources {
		resource.Resources[i].Number = resource.Resources[i].AvailableValue
		resource.Resources[i].AvailableValue = 0
	}
	resource.Resources = append(resource.Resources, accerlerateItem)
	return *resource
}

// BuildSubmitTaskRequest 构建提交任务请求的通用函数
func BuildSubmitTaskRequest(ctx context.Context, authData *models.AuthData, config *models.RunConfig,
	clusterID string, bindResultSet *models.BindResultSet, taskType string) models.SubtaskRequest {

	// 根据任务类型决定是否包含数据返回任务
	var jobs []models.Job
	//var resource models.ResourceConfig
	//fmt.Println("before")
	//if taskType == InferenceTaskType {
	//	resource = formatResource(&config.Resource)
	//} else {
	//	resource = config.Resource
	//}
	// 主任务
	fmt.Println("build train req resource is %+v", config.Resource)
	mainJob := models.Job{
		LocalJobID:  models.MainTaskID,
		Name:        "remote-task" + preparation.GenerateUniqueID(),
		Description: "材料平台远程任务",
		Type:        taskType,
		Files: &models.JobFiles{
			Dataset: models.FileBinding{
				Type:      preparation.TaskBindingType,
				BindingID: bindResultSet.BindDatasetID,
			},
			Model: models.FileBinding{
				Type:      preparation.TaskBindingType,
				BindingID: bindResultSet.BindModelID,
			},
			Image: models.ImageBinding{
				Type:    "Image",
				ImageID: config.Image,
			},
			SubModel: models.FileBinding{
				Type:      preparation.TaskBindingType,
				BindingID: bindResultSet.BindSubModelID,
			},
		},
		JobResources: &models.JobResources{
			ScheduleStrategy: "dataLocality",
			Clusters: []models.Cluster{
				{
					ClusterID: clusterID,
					Runtime: models.Runtime{
						Envs:   config.RunArgs,
						Params: config.RunArgs,
					},
					Code: models.CodeInfo{
						Type:      preparation.TaskBindingType,
						BindingID: bindResultSet.BindCodeID,
					},
					Resources: config.Resource.Resources,
				},
			},
		},
	}

	if taskType == InferenceTaskType {
		mainJob.OnlyCreate = true
	}

	jobs = append(jobs, mainJob)

	// 只有训练任务才需要数据返回任务
	if taskType == TrainingTaskType {
		dataReturnJob := models.Job{
			LocalJobID: models.DataReturnTaskID,
			Type:       models.DataReturnType,
			TargetJob: []models.TargetJob{
				{
					TargetJobID: models.MainTaskID,
					InputParams: models.InputParams{
						PackageName: "Name",
						ClusterID:   "ClusterID",
						Output:      "Output",
					},
				},
			},
		}
		jobs = append(jobs, dataReturnJob)
	}

	return models.SubtaskRequest{
		UserID: authData.JsmUserInfo.Data.UserID,
		JobSetInfo: models.JobSetInfo{
			Jobs: jobs,
		},
	}
}

//// BuildInferenceSubmitTaskRequest 构建推理任务提交请求
//func BuildInferenceSubmitTaskRequest(ctx context.Context, authData *models.AuthData,
//	codeBindingID, modelBindingID, subModelBindingID int, imageID int, clusterID string,
//	resourceType string, resources []models.ResourcesItem, command string, description string) models.InferenceSubmitTaskResponseData {
//
//	// 构建文件信息
//	files := models.InferenceJobFiles{
//		Model: models.FileBinding{
//			Type:      "Binding",
//			BindingID: modelBindingID,
//		},
//		Image: models.ImageBinding{
//			Type:    "Image",
//			ImageID: imageID,
//		},
//	}
//
//	// 如果有增量模型，添加到文件信息中
//	if subModelBindingID > 0 {
//		files.SubModel = &models.FileBinding{
//			Type:      "Binding",
//			BindingID: subModelBindingID,
//		}
//	}
//
//	// 构建任务
//	job := models.InferenceJob{
//		LocalJobID:  models.MainTaskID,
//		Name:        "inference-task" + preparation.GenerateUniqueID(),
//		Description: description,
//		Type:        "PCM_Inference",
//		Files:       &files,
//		JobResources: &models.InferenceJobResources{
//			ScheduleStrategy: "dataLocality",
//			Clusters: []models.InferenceCluster{
//				{
//					ClusterID: clusterID,
//					Runtime: models.Runtime{
//						Envs:   make(map[string]string),
//						Params: make(map[string]string),
//					},
//					Code: models.CodeInfo{
//						Type:      "Binding",
//						BindingID: codeBindingID,
//					},
//					Resources: resources,
//				},
//			},
//		},
//	}
//
//	return models.InferenceSubmitTaskResponseData{}
//}