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

package common

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

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

// 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

	// 主任务
	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,
			},
		},
		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,
				},
			},
		},
	}

	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{
		TaskInfo: models.InferenceTaskInfo{
			UserID: authData.JsmUserInfo.Data.UserID,
			JobSetInfo: models.InferenceJobSetInfo{
				Jobs: []models.InferenceJob{job},
			},
		},
		ResultInfo: models.InferenceResultInfo{
			LocalJobID: models.MainTaskID,
			JobSetID:   "", // 这个会在提交后由服务器返回
		},
	}
}