remote-task-excutor-cli/pkg/service/task/task_service.go

package task

import (
	"context"
	"encoding/json"
	"fmt"
	"github.com/mholt/archiver/v3"
	"os"
	"remote-task-excutor-cli/pkg/client"
	"remote-task-excutor-cli/pkg/models"
	"remote-task-excutor-cli/pkg/service"
	"remote-task-excutor-cli/pkg/service/preparation"
	"time"
)

const (
	MainTaskID       = "1"
	DataReturnTaskID = "4"
	AITaskType       = "AI"
	DataReturnType   = "DataReturn"

	JobSetIDKey   = "jobSetID"
	LocalJobIDKey = "localJobID"
	ResponseOK    = "OK"
)

type taskService struct {
	httpClient *client.HTTPClient
}

func NewTaskServiceService(baseURL string, timeout time.Duration) service.TaskService {
	return &taskService{
		httpClient: client.NewHTTPClient(baseURL, timeout),
	}
}

func (s *taskService) SubmitTask(ctx context.Context, authData models.AuthData, config models.RunConfig,
	clusterID string, bindResultSet models.BindResultSet) (string, error) {
	submitTaskReq := s.buildSubmitTaskReq(ctx, authData, config, clusterID, bindResultSet)
	s.httpClient.Headers["Authorization"] = "Bearer " + authData.Token
	resp, err := s.httpClient.PostJSON("/jsm/jobSet/submit", submitTaskReq)
	if err != nil {
		fmt.Println("Submit task failed: ", err)
		return "", err
	}

	var submitTaskResp models.SubmitTaskResponse
	if err := json.Unmarshal(resp, &submitTaskResp); err != nil {
		fmt.Println("Submit task response unmarshal failed: ", err)
		return "", err
	}
	if submitTaskResp.Code != ResponseOK {
		fmt.Println("Submit task failed: ", submitTaskResp.Code)
		return "", fmt.Errorf("submit task failed: %s", submitTaskResp.Code)
	}
	fmt.Println("Submit task result: ", resp)
	return submitTaskResp.Data.JobSetID, nil
}

func (s *taskService) GetTaskStatus(ctx context.Context, authData models.AuthData, jobSetID string) (string, error) {
	s.httpClient.Headers["Authorization"] = "Bearer " + authData.Token
	params := map[string]string{
		JobSetIDKey:   jobSetID,
		LocalJobIDKey: MainTaskID,
	}
	resp, err := s.httpClient.Get("/jsm/jobMgr/detail", params)
	if err != nil {
		fmt.Println("Get task status failed: ", err)
		return "", err
	}

	var taskStatus models.TaskDetailResponse
	if err := json.Unmarshal(resp, &taskStatus); err != nil {
		fmt.Println("Get task status response unmarshal failed: ", err)
		return "", err
	}

	if taskStatus.Code != ResponseOK {
		fmt.Println("Get task status failed: ", taskStatus.Code)
		return "", fmt.Errorf("get task status failed: code: %s, message: %s", taskStatus.Code, taskStatus.Message)
	}
	fmt.Println("Get task status result: ", resp)
	return taskStatus.Data.SubTaskInfos[0].Status, nil
}

func (s *taskService) GetTaskLogs(ctx context.Context, authData models.AuthData, jobSetID string) (string, error) {
	s.httpClient.Headers["Authorization"] = "Bearer " + authData.Token
	params := map[string]string{
		JobSetIDKey:   jobSetID,
		LocalJobIDKey: MainTaskID,
	}
	resp, err := s.httpClient.Get("/jsm/jobMgr/log", params)
	if err != nil {
		fmt.Println("Get task logs failed: ", err)
		return "", err
	}

	var taskLogResponse models.TaskLogResponse
	if err := json.Unmarshal(resp, &taskLogResponse); err != nil {
		fmt.Println("Get task logs response unmarshal failed: ", err)
		return "", err
	}

	if taskLogResponse.Code != ResponseOK {
		fmt.Println("Get task logs failed: ", err)
		return "", fmt.Errorf("get task logs failed: %s", taskLogResponse.Code)
	}

	fmt.Println("Get task logs result: ", resp)
	return taskLogResponse.Data, nil
}

func (s *taskService) GetTaskResult(ctx context.Context, authData models.AuthData, path string, jobSetID string) error {
	result, err := s.getTaskResultDetail(ctx, authData, jobSetID)
	if err != nil {
		return err
	}

	filename, err := s.downloadTaskResult(ctx, authData, path, result.PcmJobData.ResultFiles[0].Objects[0].PackageID)
	if err != nil {
		return err
	}

	if err := archiver.Unarchive(filename, path); err != nil {
		fmt.Println("Unarchive task result failed: ", err)
		return err
	}
	// 删除filename文件
	if err := os.Remove(filename); err != nil {
		fmt.Println("Remove task result failed: ", err)
		return err
	}

	return nil
}

func (s *taskService) downloadTaskResult(ctx context.Context, authData models.AuthData, path string, packageID int) (string, error) {
	s.httpClient.Headers["Authorization"] = "Bearer " + authData.Token
	params := map[string]string{
		"packageId": fmt.Sprintf("%s", packageID),
		"userId":    fmt.Sprintf("%s", authData.JsmUserInfo.Data.UserID),
	}
	filename, err := s.httpClient.DownloadFile("/jcs/v1/package/download", path, params)
	if err != nil {
		fmt.Println("Download task result failed: ", err)
		return "", err
	}
	return filename, nil
}

func (s *taskService) getTaskResultDetail(ctx context.Context, authData models.AuthData,
	jobSetID string) (*models.TaskResultDetailData, error) {
	s.httpClient.Headers["Authorization"] = "Bearer " + authData.Token
	queryTaskResultReq := models.TaskQueryRequest{
		JobSetID:   jobSetID,
		LocalJobID: DataReturnTaskID,
	}

	resp, err := s.httpClient.PostJSON("/jsm/jobMgr/result", queryTaskResultReq)
	if err != nil {
		fmt.Println("Get task result details failed: ", err)
		return nil, err
	}

	var taskResultDetails models.TaskResultDetailResponse
	if err := json.Unmarshal(resp, &taskResultDetails); err != nil {
		fmt.Println("Get task result details response unmarshal failed: ", err)
		return nil, err
	}

	if taskResultDetails.Code != ResponseOK {
		fmt.Println("Get task result details failed: ", taskResultDetails.Code)
		return nil, fmt.Errorf("get task result details failed: %s", taskResultDetails.Code)
	}

	return &taskResultDetails.Data, nil
}

func (s *taskService) buildSubmitTaskReq(ctx context.Context, authData models.AuthData, config models.RunConfig,
	clusterID string, bindResultSet models.BindResultSet) models.SubtaskRequest {
	return models.SubtaskRequest{
		UserID: authData.JsmUserInfo.Data.UserID,
		JobSetInfo: models.JobSetInfo{
			Jobs: []models.Job{
				{
					LocalJobID:  MainTaskID,
					Name:        "remote-task" + preparation.GenerateUniqueID(),
					Description: "材料平台远程任务",
					Type:        AITaskType,
					Files: &models.JobFiles{
						Dataset: models.FileBinding{
							Type:      preparation.DatasetType,
							BindingID: bindResultSet.BindDatasetID,
						},
						Model: models.FileBinding{
							Type:      preparation.ModelType,
							BindingID: bindResultSet.BindModelID,
						},
						Image: models.ImageBinding{
							Type:    preparation.ImageType,
							ImageID: config.Image,
						},
					},
					JobResources: &models.JobResources{
						ScheduleStrategy: "dataLocality",
						Clusters: []models.Cluster{
							{
								ClusterID: clusterID,
								Runtime: models.Runtime{
									Envs:   nil,
									Params: nil,
								},
								Code: models.CodeInfo{
									Type:      preparation.CodeType,
									BindingID: bindResultSet.BindCodeID,
								},
								Resources: config.Resource.Resources,
							},
						},
					},
				},
				{
					LocalJobID: DataReturnTaskID,
					Type:       DataReturnType,
					TargetJob: []models.TargetJob{
						{
							TargetJobID: MainTaskID,
							InputParams: models.InputParams{
								PackageName: "name",
								ClusterID:   "ClusterID",
								Output:      "Output",
							},
						},
					},
				},
			},
		},
	}
}