pcm-coordinator/internal/scheduler/service/hpc/slurm.go

package hpcservice

import (
	"context"
	"github.com/go-resty/resty/v2"
	"github.com/zeromicro/go-zero/core/logx"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/scheduler/service/collector"
	"gitlink.org.cn/JointCloud/pcm-coordinator/internal/types"
	"gitlink.org.cn/JointCloud/pcm-coordinator/pkg/constants"
)

type ParticipantHpc struct {
	participantId int64
	platform      string
	host          string
	userName      string
	accessToken   string
}

const (
	JobDetailUrl  = "/api/v1/jobs/detail/{backend}/{jobId}"
	SubmitTaskUrl = "/api/v1/jobs"
)

func NewHpc(host string, id int64, platform string) *ParticipantHpc {
	return &ParticipantHpc{
		host:          host,
		participantId: id,
		platform:      platform,
	}
}

func (c *ParticipantHpc) GetTask(ctx context.Context, taskId string) (*collector.Task, error) {
	reqUrl := c.host + JobDetailUrl
	hpcResp := &collector.HpcJobDetailResp{}
	httpClient := resty.New().R()
	_, err := httpClient.SetHeader("Content-Type", "application/json").
		SetPathParam("jobId", taskId).
		SetPathParam("backend", "slurm").
		SetResult(&hpcResp).
		Get(reqUrl)
	logx.Info("远程调用p端接口开始")
	if err != nil {
		return nil, err
	}
	logx.Info("远程调用p端接口完成")
	var resp collector.Task
	resp.Id = hpcResp.Data.ID
	if !hpcResp.Data.StartTime.IsZero() {
		resp.Start = hpcResp.Data.StartTime.Format(constants.Layout)
	}
	if !hpcResp.Data.EndTime.IsZero() {
		resp.End = hpcResp.Data.EndTime.Format(constants.Layout)
	}
	switch hpcResp.Data.StatusText {
	case "COMPLETED":
		resp.Status = constants.Completed
	case "FAILED":
		resp.Status = constants.Failed
	case "CREATED_FAILED":
		resp.Status = constants.Failed
	case "RUNNING":
		resp.Status = constants.Running
	case "STOPPED":
		resp.Status = constants.Stopped
	case "PENDING":
		resp.Status = constants.Pending
	case "WAITING":
		resp.Status = constants.Waiting
	default:
		resp.Status = "undefined"
	}

	return &resp, nil
}

func (c *ParticipantHpc) SubmitTask(ctx context.Context, req types.CommitHpcTaskReq) (*types.CommitHpcTaskResp, error) {
	reqUrl := c.host + SubmitTaskUrl
	req.Parameters["jobName"] = req.Name + "_" + req.OperateType
	resp := types.CommitHpcTaskResp{}
	httpClient := resty.New().R()
	_, err := httpClient.SetHeader("Content-Type", "application/json").
		SetBody(map[string]interface{}{
			"name":         req.Name,         // 应用名称:  BWA/lammps
			"backend":      req.Backend,      // 后端类型：slurm/sugonac
			"app":          req.App,          // 超算应用:  bwa/lammps
			"operateType":  req.OperateType,  // 应用内操作类型: bwa:构建索引/对比序列
			"parameters":   req.Parameters,   // 通用参数
			"customParams": req.CustomParams, // 各平台自定义参数
		}).
		SetResult(&resp).
		Post(reqUrl)
	if err != nil {
		return nil, err
	}
	return &resp, nil
}