提交代码。

Signed-off-by: zouap <zouap@pcl.ac.cn>

modules/aisafety/resty.go

@@ -206,7 +206,19 @@ func GetAlgorithmList() (map[string]interface{}, error) {
 
 }
 
-func GetTaskStatus(jobID string) (map[string]interface{}, error) {
+type ReturnMsg struct {
+	code  string
+	msg   string
+	data  ReturnData
+	times int64
+}
+type ReturnData struct {
+	no     string
+	code   int
+	status int
+}
+
+func GetTaskStatus(jobID string) (*ReturnMsg, error) {
 	checkSetting()
 	client := getRestyClient()
 	var taskResult string
@@ -229,9 +241,9 @@ func GetTaskStatus(jobID string) (map[string]interface{}, error) {
 		log.Info("error =" + err.Error())
 		return nil, fmt.Errorf("Get task status error: %v", err)
 	} else {
-		reMap := make(map[string]interface{})
-		err = json.Unmarshal(res.Body(), &reMap)
-		if err == nil && reMap["code"] == "0" {
+		var reMap *ReturnMsg
+		err = json.Unmarshal(res.Body(), reMap)
+		if err == nil && reMap.code == "0" {
 			return reMap, nil
 		} else {
 			return nil, fmt.Errorf("get error,code not 0")

routers/repo/aisafety.go

@@ -76,7 +76,9 @@ func GetAiSafetyTask(ctx *context.Context) {
 
 func syncAiSafetyTaskStatus(job *models.Cloudbrain) {
 	if job.Type == models.TypeCloudBrainTwo {
-
+		if isTaskNotFinished(job.Status) {
+			queryTaskStatusFromCloudbrainTwo(job)
+		}
 	} else if job.Type == models.TypeCloudBrainOne {
 		if isTaskNotFinished(job.Status) {
 			queryTaskStatusFromCloudbrain(job)
@@ -105,6 +107,72 @@ func TimerHandleModelSafetyTestTask() {
 	}
 }
 
+func queryTaskStatusFromCloudbrainTwo(job *models.Cloudbrain) {
+	log.Info("The task not finished,name=" + job.DisplayJobName)
+	result, err := modelarts.GetTrainJob(job.JobID, strconv.FormatInt(job.VersionID, 10))
+	if err != nil {
+		log.Info("query train job error." + err.Error())
+		return
+	}
+
+	job.Status = modelarts.TransTrainJobStatus(result.IntStatus)
+	job.Duration = result.Duration
+	job.TrainJobDuration = result.TrainJobDuration
+	if job.Status != string(models.ModelArtsTrainJobCompleted) {
+		err = models.UpdateJob(job)
+		if err != nil {
+			log.Error("UpdateJob failed:", err)
+		}
+	} else {
+		job.Status = string(models.ModelSafetyTesting)
+		err = models.UpdateJob(job)
+		if err != nil {
+			log.Error("UpdateJob failed:", err)
+		}
+		//send msg to beihang
+		sendNPUInferenceResultToTest(job)
+	}
+
+}
+
+func sendNPUInferenceResultToTest(job *models.Cloudbrain) {
+	datasetname := job.DatasetName
+	datasetnames := strings.Split(datasetname, ";")
+	indicator := job.LabelName
+
+	req := aisafety.TaskReq{
+		UnionId:     job.JobID,
+		EvalName:    job.DisplayJobName,
+		EvalContent: job.Description,
+		TLPath:      "test",
+		Indicators:  strings.Split(indicator, ";"),
+		CDName:      datasetnames[1],
+		BDName:      datasetnames[0],
+	}
+	jsonContent := ""
+	//TODO
+
+	if jsonContent != "" {
+		serialNo, err := aisafety.CreateSafetyTask(req, jsonContent)
+		if err == nil {
+			//update serial no to db
+			job.PreVersionName = serialNo
+			err = models.UpdateJob(job)
+			if err != nil {
+				log.Error("UpdateJob failed:", err)
+			}
+		}
+	} else {
+		log.Info("The json is null. so set it failed.")
+		//update task failed.
+		job.Status = string(models.ModelArtsTrainJobFailed)
+		err := models.UpdateJob(job)
+		if err != nil {
+			log.Error("UpdateJob failed:", err)
+		}
+	}
+}
+
 func queryTaskStatusFromCloudbrain(job *models.Cloudbrain) {
 
 	log.Info("The task not finished,name=" + job.DisplayJobName)
@@ -135,7 +203,7 @@ func queryTaskStatusFromCloudbrain(job *models.Cloudbrain) {
 			log.Error("UpdateJob failed:", err)
 		}
 		//send msg to beihang
-		sendGpuInferenceResultToTest(job)
+		sendGPUInferenceResultToTest(job)
 	}
 }
 
@@ -150,7 +218,7 @@ func queryTaskStatusFromModelSafetyTestServer(job *models.Cloudbrain) {
 	}
 }
 
-func sendGpuInferenceResultToTest(job *models.Cloudbrain) {
+func sendGPUInferenceResultToTest(job *models.Cloudbrain) {
 	datasetname := job.DatasetName
 	datasetnames := strings.Split(datasetname, ";")
 	indicator := job.LabelName
@@ -314,7 +382,7 @@ func AiSafetyCreateForPost(ctx *context.Context) {
 	} else if taskType == models.TypeCloudBrainOne {
 		createForGPU(ctx, jobName)
 	}
-	//ctx.Redirect(setting.AppSubURL + ctx.Repo.RepoLink + "/cloudbrain/benchmark")
+	ctx.Redirect(setting.AppSubURL + ctx.Repo.RepoLink + "/cloudbrain/benchmark")
 }
 
 func createForNPU(ctx *context.Context, jobName string) {