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

package hpcservice

import (
	"context"
	"github.com/go-resty/resty/v2"
	"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)
	if err != nil {
		return nil, err
	}
	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
	resp := types.CommitHpcTaskResp{}
	httpClient := resty.New().R()
	_, err := httpClient.SetHeader("Content-Type", "application/json").
		SetBody(req).
		SetResult(&resp).
		Post(reqUrl)
	if err != nil {
		return nil, err
	}
	return &resp, nil
}