aiforge/modules/grampus/resty.go

package grampus

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

	"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"
	urlGetResourceSpecs = urlOpenApiV1 + "resourcespec"
	urlGetAiCenter      = urlOpenApiV1 + "sharescreen/aicenter"
	urlGetImages        = urlOpenApiV1 + "image"
	urlNotebookJob      = urlOpenApiV1 + "notebook"

	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

	return nil
}

func createNotebookJob(req models.CreateGrampusNotebookRequest) (*models.GrampusNotebookResponse, error) {
	checkSetting()
	client := getRestyClient()
	var result models.GrampusNotebookResponse

	retry := 0

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

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

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

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

	return &result, 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 GetNotebookJob(jobID string) (*models.GrampusNotebookResponse, error) {
	checkSetting()
	client := getRestyClient()
	var result models.GrampusNotebookResponse

	retry := 0

sendjob:
	_, err := client.R().
		SetAuthToken(TOKEN).
		SetResult(&result).
		Get(HOST + urlNotebookJob + "/" + jobID)

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

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

	if result.ErrorCode != 0 {
		log.Error("GetNotebookJob failed(%d): %s", result.ErrorCode, result.ErrorMsg)
		return nil, fmt.Errorf("GetNotebookJob 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().
		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 nil, fmt.Errorf("GetJob failed(%d): %s", result.ErrorCode, result.ErrorMsg)
	}

	return &result, nil
}

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

	retry := 0

sendjob:
	_, err := client.R().
		SetAuthToken(TOKEN).
		SetResult(&result).
		Get(HOST + urlGetResourceSpecs + "?processorType=" + processorType)

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

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

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

	return &result, nil
}

func GetImages(processorType string, jobType string) (*models.GetGrampusImagesResult, error) {
	checkSetting()
	client := getRestyClient()
	var result models.GetGrampusImagesResult

	retry := 0

	queryType := "TrainJob"
	if jobType == string(models.JobTypeDebug) {
		queryType = "Notebook"
	}

sendjob:
	_, err := client.R().
		SetAuthToken(TOKEN).
		SetResult(&result).
		Get(HOST + urlGetImages + "?processorType=" + processorType + "&trainType=" + queryType)

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

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

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

	return &result, nil
}

func GetTrainJobLog(jobID string) (string, error) {
	checkSetting()
	client := getRestyClient()
	var logContent string

	res, err := client.R().
		SetAuthToken(TOKEN).
		SetResult(&logContent).
		Get(HOST + urlTrainJob + "/" + jobID + "/task/0/replica/0/log")

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

	if res.StatusCode() != http.StatusOK {
		var temp models.GrampusResult
		if err = json.Unmarshal([]byte(res.String()), &temp); err != nil {
			log.Error("json.Unmarshal failed(%s): %v", res.String(), err.Error())
			return logContent, 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 logContent, fmt.Errorf("GetTrainJobLog failed(%d):%d(%s)", res.StatusCode(), temp.ErrorCode, temp.ErrorMsg)
	}

	logContent = res.String()

	return logContent, nil
}

func GetGrampusMetrics(jobID string) (models.GetTrainJobMetricStatisticResult, error) {
	checkSetting()
	client := getRestyClient()
	var result models.GetTrainJobMetricStatisticResult
	res, err := client.R().
		SetAuthToken(TOKEN).
		Get(HOST + urlTrainJob + "/" + jobID + "/task/0/replica/0/metrics")

	if err != nil {
		return result, fmt.Errorf("resty GetTrainJobLog: %v", err)
	}
	if err = json.Unmarshal([]byte(res.String()), &result); err != nil {
		log.Error("GetGrampusMetrics json.Unmarshal failed(%s): %v", res.String(), err.Error())
		return result, fmt.Errorf("json.Unmarshal failed(%s): %v", res.String(), err.Error())
	}
	if res.StatusCode() != http.StatusOK {
		log.Error("Call GrampusMetrics failed(%d):%s(%s)", res.StatusCode(), result.ErrorCode, result.ErrorMsg)
		return result, fmt.Errorf("Call GrampusMetrics failed(%d):%d(%s)", res.StatusCode(), result.ErrorCode, result.ErrorMsg)
	}
	if !result.IsSuccess {
		log.Error("GetGrampusMetrics(%s) failed", jobID)
		return result, fmt.Errorf("GetGrampusMetrics failed:%s", result.ErrorMsg)
	}
	return result, nil
}

func StopJob(jobID string, jobType ...string) (*models.GrampusStopJobResponse, error) {
	checkSetting()
	client := getRestyClient()
	var result models.GrampusStopJobResponse

	retry := 0

	url := urlTrainJob
	if len(jobType) > 0 {
		if jobType[0] == string(models.JobTypeDebug) {
			url = urlNotebookJob
		}
	}

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

	if err != nil {
		return &result, fmt.Errorf("resty StopTrainJob: %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 GetAiCenters(pageIndex, pageSize int) (*models.GetGrampusAiCentersResult, error) {
	checkSetting()
	client := getRestyClient()
	var result models.GetGrampusAiCentersResult

	retry := 0

sendjob:
	_, err := client.R().
		SetAuthToken(TOKEN).
		SetResult(&result).
		Get(HOST + urlGetAiCenter + "?pageIndex=" + fmt.Sprint(pageIndex) + "&pageSize=" + fmt.Sprint(pageSize))

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

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

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

	return &result, nil
}

func RestartNotebookJob(jobID string) (*models.GrampusNotebookRestartResponse, error) {
	checkSetting()
	client := getRestyClient()
	var restartResponse *models.GrampusNotebookRestartResponse
	retry := 0

sendjob:
	res, err := client.R().
		SetAuthToken(TOKEN).
		SetResult(&restartResponse).
		Post(HOST + urlNotebookJob + "/" + jobID + "/start")

	if err != nil {
		return nil, fmt.Errorf("resty grampus restart note book job: %v", err)
	}
	if restartResponse.ErrorCode == errorIllegalToken && retry < 1 {
		retry++
		log.Info("retry get token")
		_ = getToken()
		goto sendjob
	}

	if res.StatusCode() != http.StatusOK {
		log.Error("resty grampus restart note book job failed(%s): %v", res.String(), err.Error())
		return nil, fmt.Errorf("resty grampus restart note book job failed: %v", err)
	}

	return restartResponse, nil
}