aiforge/modules/grampus/resty.go

package grampus

import (
	"crypto/tls"
	"encoding/json"
	"fmt"
	"net/http"
	"strconv"

	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/setting"
	"github.com/go-resty/resty/v2"
)

var (
	restyClient *resty.Client
	HOST        string
	TOKEN       string
)

const (
	urlOpenApiV1 = "/openapi/v1/"

	urlGetToken          = urlOpenApiV1 + "token"
	urlTrainJob          = urlOpenApiV1 + "trainjob"
	urlResourceSpecs     = "/job/resource-specs"
	urlTrainJobConfig    = "/training-job-configs"
	errorCodeExceedLimit = "ModelArts.0118"

	urlNotebook2 = ""

	errorIllegalToken = 1005
)

type GetTokenParams struct {
	UserName string `json:"username"`
	Password string `json:"password"`
}

type GetTokenResult struct {
	Token      string `json:"token"`
	Expiration int64  `json:"expiration"`
}

func getRestyClient() *resty.Client {
	if restyClient == nil {
		restyClient = resty.New()
		restyClient.SetTLSClientConfig(&tls.Config{InsecureSkipVerify: true})
	}
	return restyClient
}

func checkSetting() {
	if len(HOST) != 0 && len(TOKEN) != 0 && restyClient != nil {
		return
	}

	err := getToken()
	if err != nil {
		log.Error("getToken failed:%v", err)
	}
}

func getToken() error {
	HOST = setting.Grampus.Host

	client := getRestyClient()
	params := GetTokenParams{
		UserName: setting.Grampus.UserName,
		Password: setting.Grampus.Password,
	}

	var result GetTokenResult
	res, err := client.R().
		SetHeader("Content-Type", "application/json").
		SetBody(params).
		SetResult(&result).
		Post(HOST + urlGetToken)
	if err != nil {
		return fmt.Errorf("resty getToken: %v", err)
	}

	if res.StatusCode() != http.StatusOK {
		return fmt.Errorf("getToken failed:%s", res.String())
	}

	TOKEN = result.Token
	log.Info(TOKEN)

	return nil
}

func createJob(req models.CreateGrampusJobRequest) (*models.CreateGrampusJobResponse, error) {
	checkSetting()
	client := getRestyClient()
	var result models.CreateGrampusJobResponse

	retry := 0

sendjob:
	_, err := client.R().
		SetHeader("Content-Type", "application/json").
		SetAuthToken(TOKEN).
		SetBody(req).
		SetResult(&result).
		Post(HOST + urlTrainJob)

	if err != nil {
		return nil, fmt.Errorf("resty CreateJob: %s", err)
	}

	if result.ErrorCode == errorIllegalToken && retry < 1 {
		retry++
		_ = getToken()
		goto sendjob
	}

	if result.ErrorCode != 0 {
		log.Error("CreateJob failed(%d): %s", result.ErrorCode, result.ErrorMsg)
		return &result, fmt.Errorf("CreateJob failed(%d): %s", result.ErrorCode, result.ErrorMsg)
	}

	return &result, nil
}

func GetJob(jobID string) (*models.GetGrampusJobResponse, error) {
	checkSetting()
	client := getRestyClient()
	var result models.GetGrampusJobResponse

	retry := 0

sendjob:
	_, err := client.R().
		SetHeader("Content-Type", "application/json").
		SetAuthToken(TOKEN).
		SetResult(&result).
		Get(HOST + urlTrainJob + "/" + jobID)

	if err != nil {
		return nil, fmt.Errorf("resty GetJob: %v", err)
	}

	if result.ErrorCode == errorIllegalToken && retry < 1 {
		retry++
		log.Info("retry get token")
		_ = getToken()
		goto sendjob
	}

	if result.ErrorCode != 0 {
		log.Error("GetJob failed(%d): %s", result.ErrorCode, result.ErrorMsg)
		return &result, fmt.Errorf("GetJob failed(%d): %s", result.ErrorCode, result.ErrorMsg)
	}

	return &result, nil
}

func DelJob(jobID string) (*models.NotebookDelResult, error) {
	checkSetting()
	client := getRestyClient()
	var result models.NotebookDelResult

	retry := 0

sendjob:
	res, err := client.R().
		SetHeader("Content-Type", "application/json").
		SetAuthToken(TOKEN).
		SetResult(&result).
		Delete(HOST + "/v1/" + setting.ProjectID + urlTrainJob + "/" + jobID)

	if err != nil {
		return &result, fmt.Errorf("resty DelJob: %v", err)
	}

	if res.StatusCode() == http.StatusUnauthorized && retry < 1 {
		retry++
		_ = getToken()
		goto sendjob
	}

	var response models.NotebookResult
	err = json.Unmarshal(res.Body(), &response)
	if err != nil {
		log.Error("json.Unmarshal failed: %s", err.Error())
		return &result, fmt.Errorf("son.Unmarshal failed: %s", err.Error())
	}

	if len(response.ErrorCode) != 0 {
		log.Error("DelJob failed(%s): %s", response.ErrorCode, response.ErrorMsg)
		return &result, fmt.Errorf("DelJob failed(%s): %s", response.ErrorCode, response.ErrorMsg)
	}

	return &result, nil
}

func createTrainJob(createJobParams models.CreateTrainJobParams) (*models.CreateTrainJobResult, error) {
	checkSetting()
	client := getRestyClient()
	var result models.CreateTrainJobResult

	retry := 0

sendjob:
	res, err := client.R().
		SetHeader("Content-Type", "application/json").
		SetAuthToken(TOKEN).
		SetBody(createJobParams).
		SetResult(&result).
		Post(HOST + "/v1/" + setting.ProjectID + urlTrainJob)

	if err != nil {
		return nil, fmt.Errorf("resty create train-job: %s", err)
	}

	req, _ := json.Marshal(createJobParams)
	log.Info("%s", req)

	if res.StatusCode() == http.StatusUnauthorized && retry < 1 {
		retry++
		_ = getToken()
		goto sendjob
	}

	if res.StatusCode() != http.StatusOK {
		var temp models.ErrorResult
		if err = json.Unmarshal([]byte(res.String()), &temp); err != nil {
			log.Error("json.Unmarshal failed(%s): %v", res.String(), err.Error())
			return &result, fmt.Errorf("json.Unmarshal failed(%s): %v", res.String(), err.Error())
		}
		log.Error("createTrainJob failed(%d):%s(%s)", res.StatusCode(), temp.ErrorCode, temp.ErrorMsg)
		BootFileErrorMsg := "Invalid OBS path '" + createJobParams.Config.BootFileUrl + "'."
		DataSetErrorMsg := "Invalid OBS path '" + createJobParams.Config.DataUrl + "'."
		if temp.ErrorMsg == BootFileErrorMsg {
			log.Error("启动文件错误！createTrainJob failed(%d):%s(%s)", res.StatusCode(), temp.ErrorCode, temp.ErrorMsg)
			return &result, fmt.Errorf("启动文件错误！")
		}
		if temp.ErrorMsg == DataSetErrorMsg {
			log.Error("数据集错误！createTrainJob failed(%d):%s(%s)", res.StatusCode(), temp.ErrorCode, temp.ErrorMsg)
			return &result, fmt.Errorf("数据集错误！")
		}
		return &result, fmt.Errorf("createTrainJob failed(%d):%s(%s)", res.StatusCode(), temp.ErrorCode, temp.ErrorMsg)
	}

	if !result.IsSuccess {
		log.Error("createTrainJob failed(%s): %s", result.ErrorCode, result.ErrorMsg)
		return &result, fmt.Errorf("createTrainJob failed(%s): %s", result.ErrorCode, result.ErrorMsg)
	}

	return &result, nil
}

func GetResourceSpecs() (*models.GetResourceSpecsResult, error) {
	checkSetting()
	client := getRestyClient()
	var result models.GetResourceSpecsResult

	retry := 0

sendjob:
	res, err := client.R().
		SetHeader("Content-Type", "application/json").
		SetAuthToken(TOKEN).
		SetResult(&result).
		Get(HOST + "/v1/" + setting.ProjectID + urlResourceSpecs)

	if err != nil {
		return nil, fmt.Errorf("resty GetResourceSpecs: %v", err)
	}

	if res.StatusCode() == http.StatusUnauthorized && retry < 1 {
		retry++
		_ = getToken()
		goto sendjob
	}

	if res.StatusCode() != http.StatusOK {
		var temp models.ErrorResult
		if err = json.Unmarshal([]byte(res.String()), &temp); err != nil {
			log.Error("json.Unmarshal failed(%s): %v", res.String(), err.Error())
			return &result, fmt.Errorf("json.Unmarshal failed(%s): %v", res.String(), err.Error())
		}
		log.Error("GetResourceSpecs failed(%d):%s(%s)", res.StatusCode(), temp.ErrorCode, temp.ErrorMsg)
		return &result, fmt.Errorf("GetResourceSpecs failed(%d):%s(%s)", res.StatusCode(), temp.ErrorCode, temp.ErrorMsg)
	}

	if !result.IsSuccess {
		log.Error("GetResourceSpecs failed(%s): %s", result.ErrorCode, result.ErrorMsg)
		return &result, fmt.Errorf("GetResourceSpecs failed(%s): %s", result.ErrorCode, result.ErrorMsg)
	}

	return &result, nil
}

func GetTrainJob(jobID, versionID string) (*models.GetTrainJobResult, error) {
	checkSetting()
	client := getRestyClient()
	var result models.GetTrainJobResult

	retry := 0

sendjob:
	res, err := client.R().
		SetAuthToken(TOKEN).
		SetResult(&result).
		Get(HOST + "/v1/" + setting.ProjectID + urlTrainJob + "/" + jobID + "/versions/" + versionID)

	if err != nil {
		return nil, fmt.Errorf("resty GetTrainJob: %v", err)
	}

	if res.StatusCode() == http.StatusUnauthorized && retry < 1 {
		retry++
		_ = getToken()
		goto sendjob
	}

	if res.StatusCode() != http.StatusOK {
		var temp models.ErrorResult
		if err = json.Unmarshal([]byte(res.String()), &temp); err != nil {
			log.Error("json.Unmarshal failed(%s): %v", res.String(), err.Error())
			return &result, fmt.Errorf("json.Unmarshal failed(%s): %v", res.String(), err.Error())
		}
		log.Error("GetTrainJob failed(%d):%s(%s)", res.StatusCode(), temp.ErrorCode, temp.ErrorMsg)
		return &result, fmt.Errorf("获取作业详情失败(%d):%s(%s)", res.StatusCode(), temp.ErrorCode, temp.ErrorMsg)
	}

	if !result.IsSuccess {
		log.Error("GetTrainJob(%s) failed", jobID)
		return &result, fmt.Errorf("获取作业详情失败")
	}

	return &result, nil
}

func GetTrainJobLog(jobID, versionID, baseLine, logFile, order string, lines int) (*models.GetTrainJobLogResult, error) {
	checkSetting()
	client := getRestyClient()
	var result models.GetTrainJobLogResult

	retry := 0

sendjob:
	res, err := client.R().
		SetQueryParams(map[string]string{
			"base_line": baseLine,
			"lines":     strconv.Itoa(lines),
			"log_file":  logFile,
			"order":     order,
		}).
		SetAuthToken(TOKEN).
		SetResult(&result).
		Get(HOST + "/v1/" + setting.ProjectID + urlTrainJob + "/" + jobID + "/versions/" + versionID + "/aom-log")

	if err != nil {
		return nil, fmt.Errorf("resty GetTrainJobLog: %v", err)
	}

	if res.StatusCode() == http.StatusUnauthorized && retry < 1 {
		retry++
		_ = getToken()
		goto sendjob
	}

	if res.StatusCode() != http.StatusOK {
		var temp models.ErrorResult
		if err = json.Unmarshal([]byte(res.String()), &temp); err != nil {
			log.Error("json.Unmarshal failed(%s): %v", res.String(), err.Error())
			return &result, fmt.Errorf("json.Unmarshal failed(%s): %v", res.String(), err.Error())
		}
		log.Error("GetTrainJobLog failed(%d):%s(%s)", res.StatusCode(), temp.ErrorCode, temp.ErrorMsg)
		return &result, fmt.Errorf("获取作业日志失败(%d):%s(%s)", res.StatusCode(), temp.ErrorCode, temp.ErrorMsg)
	}

	if !result.IsSuccess {
		log.Error("GetTrainJobLog(%s) failed", jobID)
		return &result, fmt.Errorf("获取作业日志失败:%s", result.ErrorMsg)
	}

	return &result, nil
}

func GetTrainJobLogFileNames(jobID, versionID string) (*models.GetTrainJobLogFileNamesResult, error) {
	checkSetting()
	client := getRestyClient()
	var result models.GetTrainJobLogFileNamesResult

	retry := 0

sendjob:
	res, err := client.R().
		SetAuthToken(TOKEN).
		SetResult(&result).
		Get(HOST + "/v1/" + setting.ProjectID + urlTrainJob + "/" + jobID + "/versions/" + versionID + "/log/file-names")

	if err != nil {
		return nil, fmt.Errorf("resty GetTrainJobLogFileNames: %v", err)
	}

	if res.StatusCode() == http.StatusUnauthorized && retry < 1 {
		retry++
		_ = getToken()
		goto sendjob
	}

	if res.StatusCode() != http.StatusOK {
		var temp models.ErrorResult
		if err = json.Unmarshal([]byte(res.String()), &temp); err != nil {
			log.Error("json.Unmarshal failed(%s): %v", res.String(), err.Error())
			return &result, fmt.Errorf("json.Unmarshal failed(%s): %v", res.String(), err.Error())
		}
		log.Error("GetTrainJobLogFileNames failed(%d):%s(%s)", res.StatusCode(), temp.ErrorCode, temp.ErrorMsg)
		return &result, fmt.Errorf("GetTrainJobLogFileNames failed(%d):%s(%s)", res.StatusCode(), temp.ErrorCode, temp.ErrorMsg)
	}

	if !result.IsSuccess {
		log.Error("GetTrainJobLogFileNames(%s) failed", jobID)
		return &result, fmt.Errorf("获取作业日志文件失败:%s", result.ErrorMsg)
	}

	return &result, nil
}

func DelTrainJob(jobID string) (*models.TrainJobResult, error) {
	checkSetting()
	client := getRestyClient()
	var result models.TrainJobResult

	retry := 0

sendjob:
	res, err := client.R().
		SetAuthToken(TOKEN).
		SetResult(&result).
		Delete(HOST + "/v1/" + setting.ProjectID + urlTrainJob + "/" + jobID)

	if err != nil {
		return &result, fmt.Errorf("resty DelTrainJob: %v", err)
	}

	if res.StatusCode() == http.StatusUnauthorized && retry < 1 {
		retry++
		_ = getToken()
		goto sendjob
	}

	if res.StatusCode() != http.StatusOK {
		var temp models.ErrorResult
		if err = json.Unmarshal([]byte(res.String()), &temp); err != nil {
			log.Error("json.Unmarshal failed(%s): %v", res.String(), err.Error())
			return &result, fmt.Errorf("json.Unmarshal failed(%s): %v", res.String(), err.Error())
		}
		log.Error("DelTrainJob failed(%d):%s(%s)", res.StatusCode(), temp.ErrorCode, temp.ErrorMsg)
		return &result, fmt.Errorf("删除训练作业失败(%d):%s(%s)", res.StatusCode(), temp.ErrorCode, temp.ErrorMsg)
	}

	if !result.IsSuccess {
		log.Error("DelTrainJob(%s) failed", jobID)
		return &result, fmt.Errorf("删除训练作业失败:%s", result.ErrorMsg)
	}

	return &result, nil
}

func StopTrainJob(jobID, versionID string) (*models.TrainJobResult, error) {
	checkSetting()
	client := getRestyClient()
	var result models.TrainJobResult

	retry := 0

sendjob:
	res, err := client.R().
		SetAuthToken(TOKEN).
		SetResult(&result).
		Post(HOST + "/v1/" + setting.ProjectID + urlTrainJob + "/" + jobID + "/versions/" + versionID + "/stop")

	if err != nil {
		return &result, fmt.Errorf("resty StopTrainJob: %v", err)
	}

	if res.StatusCode() == http.StatusUnauthorized && retry < 1 {
		retry++
		_ = getToken()
		goto sendjob
	}

	if res.StatusCode() != http.StatusOK {
		var temp models.ErrorResult
		if err = json.Unmarshal([]byte(res.String()), &temp); err != nil {
			log.Error("json.Unmarshal failed(%s): %v", res.String(), err.Error())
			return &result, fmt.Errorf("json.Unmarshal failed(%s): %v", res.String(), err.Error())
		}
		log.Error("StopTrainJob failed(%d):%s(%s)", res.StatusCode(), temp.ErrorCode, temp.ErrorMsg)
		return &result, fmt.Errorf("停止训练作业失败(%d):%s(%s)", res.StatusCode(), temp.ErrorCode, temp.ErrorMsg)
	}

	if !result.IsSuccess {
		log.Error("StopTrainJob(%s) failed", jobID)
		return &result, fmt.Errorf("停止训练作业失败:%s", result.ErrorMsg)
	}

	return &result, nil
}

func DelTrainJobVersion(jobID string, versionID string) (*models.TrainJobResult, error) {
	checkSetting()
	client := getRestyClient()
	var result models.TrainJobResult

	retry := 0

sendjob:
	res, err := client.R().
		SetAuthToken(TOKEN).
		SetResult(&result).
		Delete(HOST + "/v1/" + setting.ProjectID + urlTrainJob + "/" + jobID + "/versions/" + versionID)

	if err != nil {
		return &result, fmt.Errorf("resty DelTrainJobVersion: %v", err)
	}

	if res.StatusCode() == http.StatusUnauthorized && retry < 1 {
		retry++
		_ = getToken()
		goto sendjob
	}

	if res.StatusCode() != http.StatusOK {
		var temp models.ErrorResult
		if err = json.Unmarshal([]byte(res.String()), &temp); err != nil {
			log.Error("json.Unmarshal failed(%s): %v", res.String(), err.Error())
			return &result, fmt.Errorf("json.Unmarshal failed(%s): %v", res.String(), err.Error())
		}
		log.Error("DelTrainJob failed(%d):%s(%s)", res.StatusCode(), temp.ErrorCode, temp.ErrorMsg)
		return &result, fmt.Errorf("删除训练作业版本失败(%d):%s(%s)", res.StatusCode(), temp.ErrorCode, temp.ErrorMsg)
	}

	if !result.IsSuccess {
		log.Error("DelTrainJob(%s) failed", jobID)
		return &result, fmt.Errorf("删除训练作业版本失败:%s", result.ErrorMsg)
	}

	return &result, nil
}

func createInferenceJob(createJobParams models.CreateInferenceJobParams) (*models.CreateTrainJobResult, error) {
	checkSetting()
	client := getRestyClient()
	var result models.CreateTrainJobResult

	retry := 0

sendjob:
	res, err := client.R().
		SetHeader("Content-Type", "application/json").
		SetAuthToken(TOKEN).
		SetBody(createJobParams).
		SetResult(&result).
		Post(HOST + "/v1/" + setting.ProjectID + urlTrainJob)

	if err != nil {
		return nil, fmt.Errorf("resty create inference-job: %s", err)
	}

	req, _ := json.Marshal(createJobParams)
	log.Info("%s", req)

	if res.StatusCode() == http.StatusUnauthorized && retry < 1 {
		retry++
		_ = getToken()
		goto sendjob
	}

	if res.StatusCode() != http.StatusOK {
		var temp models.ErrorResult
		if err = json.Unmarshal([]byte(res.String()), &temp); err != nil {
			log.Error("json.Unmarshal failed(%s): %v", res.String(), err.Error())
			return &result, fmt.Errorf("json.Unmarshal failed(%s): %v", res.String(), err.Error())
		}
		log.Error("createInferenceJob failed(%d):%s(%s)", res.StatusCode(), temp.ErrorCode, temp.ErrorMsg)
		BootFileErrorMsg := "Invalid OBS path '" + createJobParams.InfConfig.BootFileUrl + "'."
		DataSetErrorMsg := "Invalid OBS path '" + createJobParams.InfConfig.DataUrl + "'."
		if temp.ErrorMsg == BootFileErrorMsg {
			log.Error("启动文件错误！createInferenceJob failed(%d):%s(%s)", res.StatusCode(), temp.ErrorCode, temp.ErrorMsg)
			return &result, fmt.Errorf("启动文件错误！")
		}
		if temp.ErrorMsg == DataSetErrorMsg {
			log.Error("数据集错误！createInferenceJob failed(%d):%s(%s)", res.StatusCode(), temp.ErrorCode, temp.ErrorMsg)
			return &result, fmt.Errorf("数据集错误！")
		}
		return &result, fmt.Errorf("createInferenceJob failed(%d):%s(%s)", res.StatusCode(), temp.ErrorCode, temp.ErrorMsg)
	}

	if !result.IsSuccess {
		log.Error("createInferenceJob failed(%s): %s", result.ErrorCode, result.ErrorMsg)
		return &result, fmt.Errorf("createInferenceJob failed(%s): %s", result.ErrorCode, result.ErrorMsg)
	}

	return &result, nil
}