pcm-hpc/routers/v1/job.go

package v1

import (
	"encoding/json"
	"fmt"
	"github.com/gin-gonic/gin"
	"github.com/go-resty/resty/v2"
	"github.com/pkg/errors"
	"github.com/zeromicro/go-zero/core/logx"
	"gitlink.org.cn/JointCloud/pcm-hpc/slurm"
	"io/ioutil"
	"net/http"
)

type SubmitJobReq struct {
	Server     string
	Version    string
	Username   string
	Token      string
	JobOptions JobOptions
}

type JobInfoReq struct {
	Server   string `json:"server" form:"server""`
	Version  string `json:"version" form:"version""`
	Username string `json:"username" form:"username"`
	Token    string `json:"token" form:"token"`
	JobId    string `json:"jobId" form:"jobId"`
}

type JobCancelReq struct {
	Server   string `json:"server" form:"server""`
	Version  string `json:"version" form:"version""`
	Username string `json:"username" form:"username"`
	Token    string `json:"token" form:"token"`
	JobId    string `json:"jobId" form:"jobId"`
}

type SubmitJobResp struct {
	Meta struct {
		Plugin struct {
			Type string `json:"type"`
			Name string `json:"name"`
		} `json:"plugin"`
		Slurm struct {
			Version struct {
				Major int `json:"major"`
				Micro int `json:"micro"`
				Minor int `json:"minor"`
			} `json:"version"`
			Release string `json:"release"`
		} `json:"Slurm"`
	} `json:"meta"`
	Errors           []interface{} `json:"errors"`
	JobId            int           `json:"job_id"`
	StepId           string        `json:"step_id"`
	JobSubmitUserMsg string        `json:"job_submit_user_msg"`
}

type Job struct {
	Meta struct {
		Plugin struct {
			Type string `json:"type"`
			Name string `json:"name"`
		} `json:"plugin"`
		Slurm struct {
			Version struct {
				Major int `json:"major"`
				Micro int `json:"micro"`
				Minor int `json:"minor"`
			} `json:"version"`
			Release string `json:"release"`
		} `json:"Slurm"`
	} `json:"meta"`
	Errors []interface{} `json:"errors"`
}

type JobOptions struct {
	Script string         `json:"script"`
	Job    *JobProperties `json:"job"`
}

type JobProperties struct {
	Account                 string            `json:"account"`
	CurrentWorkingDirectory string            `json:"current_working_directory"`
	Environment             map[string]string `json:"environment"`
	Name                    string            `json:"name"`
	Partition               string            `json:"partition"`
	StandardError           string            `json:"standard_error"`
	StandardOutput          string            `json:"standard_output"`
	NTasks                  uint32            `json:"ntasks"`
}

func SubmitJob(ctx *gin.Context) {
	var req SubmitJobReq
	var resp SubmitJobResp
	if err := ctx.BindJSON(&req); err != nil {
		Response(ctx, http.StatusBadRequest, "invalid request params.", "")
		return
	}
	url := req.Server + "/slurm/" + req.Version + "/job/submit"
	logx.Info("调用rest api接口开始")
	httpClient := resty.New().R()

	httpResp, err := httpClient.SetHeader("Content-Type", "application/json").
		SetHeader("Accept", "application/json").
		SetHeader("X-SLURM-USER-NAME", req.Username).
		SetHeader("X-SLURM-USER-TOKEN", req.Token).
		SetBody(&req.JobOptions).
		SetResult(&resp).
		Post(url)
	logx.Info("调用rest api接口完成")
	if err != nil {
		Response(ctx, http.StatusInternalServerError, "rest api error.", err)
	}
	if httpResp.IsError() {
		Response(ctx, http.StatusInternalServerError, "rest api error.", errors.Errorf(httpResp.String()))
	}

	if len(resp.Errors) > 0 {
		strErr, _ := resp.Errors[0].(string)
		Response(ctx, http.StatusInternalServerError, "rest api error.", strErr)
	}
	Response(ctx, http.StatusOK, "success", resp)
}

func JobInfo(ctx *gin.Context) {
	var req JobInfoReq
	if err := ctx.ShouldBindQuery(&req); err != nil {
		Response(ctx, http.StatusBadRequest, "invalid request params.", "")
		return
	}

	url := req.Server + "/slurm/" + req.Version + "/job/" + req.JobId
	client := &http.Client{}
	request, err := http.NewRequest("GET", url, nil)

	if err != nil {
		fmt.Println(err)
		return
	}
	request.Header.Add("Accept", "*/*")
	request.Header.Add("Host", req.Server)
	request.Header.Add("X-SLURM-USER-NAME", req.Username)
	request.Header.Add("X-SLURM-USER-TOKEN", req.Token)
	request.Header.Add("Cookie", "X-SLURM-USER-NAME=;X-SLURM-USER-TOKEN=")
	request.Header.Add("Content-Type", "application/json")

	res, err := client.Do(request)
	if err != nil {
		fmt.Println(err)
		return
	}
	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		fmt.Println(err)
		return
	}
	var resp slurm.GetJobResp
	json.Unmarshal(body, &resp)
	ResponseData(ctx, &resp)
}

func CancelJob(ctx *gin.Context) {
	var req JobCancelReq
	if err := ctx.ShouldBindQuery(&req); err != nil {
		Response(ctx, http.StatusBadRequest, "invalid request params.", "")
		return
	}

	url := req.Server + "/slurm/" + req.Version + "/job/" + req.JobId
	client := &http.Client{}
	request, err := http.NewRequest("DELETE", url, nil)

	if err != nil {
		fmt.Println(err)
		return
	}
	request.Header.Add("Accept", "*/*")
	request.Header.Add("Host", req.Server)
	request.Header.Add("X-SLURM-USER-NAME", req.Username)
	request.Header.Add("X-SLURM-USER-TOKEN", req.Token)
	request.Header.Add("Cookie", "X-SLURM-USER-NAME=;X-SLURM-USER-TOKEN=")
	request.Header.Add("Content-Type", "application/json")

	res, err := client.Do(request)
	if err != nil {
		fmt.Println(err)
		return
	}
	defer res.Body.Close()

	body, err := ioutil.ReadAll(res.Body)
	if err != nil {
		Response(ctx, http.StatusInternalServerError, "rest api error.", err.Error())
		return
	}
	var job Job
	json.Unmarshal(body, &job)
	ResponseData(ctx, &job)
}