package modelarts import ( "code.gitea.io/gitea/modules/log" "crypto/tls" "encoding/json" "fmt" "net/http" "code.gitea.io/gitea/models" "code.gitea.io/gitea/modules/setting" "github.com/go-resty/resty/v2" ) var ( restyClient *resty.Client HOST string TOKEN string ) const ( methodPassword = "password" urlGetToken = "/v3/auth/tokens" urlNotebook = "/demanager/instances" urlTrainJob = "/training-jobs" urlResourceSpecs = "/job/resource-specs" urlTrainJobConfig = "/training-job-configs" errorCodeExceedLimit = "ModelArts.0118" ) 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.ModelArtsHost client := getRestyClient() params := models.GetTokenParams{ Auth: models.Auth{ Identity: models.Identity{ Methods: []string{methodPassword}, Password: models.Password{ User: models.NotebookUser{ Name: setting.ModelArtsUsername, Password: setting.ModelArtsPassword, Domain: models.Domain{ Name: setting.ModelArtsDomain, }, }, }, }, Scope: models.Scope{ Project: models.Project{ Name: setting.ProjectName, }, }, }, } res, err := client.R(). SetHeader("Content-Type", "application/json"). SetBody(params). Post(setting.IamHost + urlGetToken) if err != nil { return fmt.Errorf("resty getToken: %v", err) } if res.StatusCode() != http.StatusCreated { return fmt.Errorf("getToken failed:%s", res.String()) } TOKEN = res.Header().Get("X-Subject-Token") return nil } func createNotebook(createJobParams models.CreateNotebookParams) (*models.CreateNotebookResult, error) { checkSetting() client := getRestyClient() var result models.CreateNotebookResult retry := 0 sendjob: res, err := client.R(). SetHeader("Content-Type", "application/json"). SetAuthToken(TOKEN). SetBody(createJobParams). SetResult(&result). Post(HOST + "/v1/" + setting.ProjectID + urlNotebook) if err != nil { return nil, fmt.Errorf("resty create notebook: %s", 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("createNotebook failed(%s): %s", response.ErrorCode, response.ErrorMsg) if response.ErrorCode == errorCodeExceedLimit { response.ErrorMsg = "所选规格使用数量已超过最大配额限制。" } return &result, fmt.Errorf("createNotebook failed(%s): %s", response.ErrorCode, response.ErrorMsg) } return &result, nil } func GetJob(jobID string) (*models.GetNotebookResult, error) { checkSetting() client := getRestyClient() var result models.GetNotebookResult retry := 0 sendjob: res, err := client.R(). SetHeader("Content-Type", "application/json"). SetAuthToken(TOKEN). SetResult(&result). Get(HOST + "/v1/" + setting.ProjectID + urlNotebook + "/" + jobID) if err != nil { return nil, fmt.Errorf("resty GetJob: %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("GetJob failed(%s): %s", response.ErrorCode, response.ErrorMsg) return &result, fmt.Errorf("GetJob failed(%s): %s", response.ErrorCode, response.ErrorMsg) } return &result, nil } func StopJob(jobID string, param models.NotebookAction) (*models.NotebookActionResult, error) { checkSetting() client := getRestyClient() var result models.NotebookActionResult retry := 0 sendjob: res, err := client.R(). SetHeader("Content-Type", "application/json"). SetBody(param). SetAuthToken(TOKEN). SetResult(&result). Post(HOST + "/v1/" + setting.ProjectID + urlNotebook + "/" + jobID + "/action") if err != nil { return &result, fmt.Errorf("resty StopJob: %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("StopJob failed(%s): %s", response.ErrorCode, response.ErrorMsg) return &result, fmt.Errorf("StopJob failed(%s): %s", response.ErrorCode, response.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 + urlNotebook + "/" + 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 GetJobToken(jobID string) (*models.NotebookGetJobTokenResult, error) { checkSetting() client := getRestyClient() var result models.NotebookGetJobTokenResult retry := 0 sendjob: res, err := client.R(). SetHeader("Content-Type", "application/json"). SetAuthToken(TOKEN). SetResult(&result). Get(HOST + "/v1/" + setting.ProjectID + urlNotebook + "/" + jobID + "/token") if err != nil { return &result, fmt.Errorf("resty GetJobToken: %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("GetJobToken failed(%s): %s", response.ErrorCode, response.ErrorMsg) return &result, fmt.Errorf("GetJobToken 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) } //log.Info("%d", res.StatusCode()) //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) 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 CreateTrainJobConfig(req models.CreateConfigParams) (*models.CreateTrainJobConfigResult, error) { checkSetting() client := getRestyClient() var result models.CreateTrainJobConfigResult retry := 0 sendjob: res, err := client.R(). SetHeader("Content-Type", "application/json"). SetAuthToken(TOKEN). SetBody(req). SetResult(&result). Post(HOST + "/v1/" + setting.ProjectID + urlTrainJobConfig) if err != nil { return nil, fmt.Errorf("resty CreateTrainJobConfig: %s", err) } if res.StatusCode() == http.StatusUnauthorized && retry < 1 { retry++ _ = getToken() goto sendjob } temp, _ := json.Marshal(req) log.Info("%s", temp) 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("CreateTrainJobConfig failed(%d):%s(%s)", res.StatusCode(), temp.ErrorCode, temp.ErrorMsg) return &result, fmt.Errorf("CreateTrainJobConfig failed(%d):%s(%s)", res.StatusCode(), temp.ErrorCode, temp.ErrorMsg) } if !result.IsSuccess { log.Error("CreateTrainJobConfig failed(%s): %s", result.ErrorCode, result.ErrorMsg) return &result, fmt.Errorf("CreateTrainJobConfig failed(%s): %s", result.ErrorCode, result.ErrorMsg) } return &result, nil }