代码提交

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

models/cloudbrain.go

@@ -52,6 +52,8 @@ const (
 	JobFailed    CloudbrainStatus = "FAILED"
 	JobRunning   CloudbrainStatus = "RUNNING"
 
+	ModelSafetyTesting CloudbrainStatus = "TESTING"
+
 	JobTypeDebug        JobType = "DEBUG"
 	JobTypeBenchmark    JobType = "BENCHMARK"
 	JobTypeModelSafety  JobType = "MODELSAFETY"
@@ -1935,6 +1937,13 @@ func GetStoppedJobWithNoStartTimeEndTime() ([]*Cloudbrain, error) {
 	return cloudbrains, x.SQL("select * from cloudbrain where status in (?,?,?,?,?,?,?) and (start_time is null or end_time is null) limit 100", ModelArtsTrainJobCompleted, ModelArtsTrainJobFailed, ModelArtsTrainJobKilled, ModelArtsStopped, JobStopped, JobFailed, JobSucceeded).Find(&cloudbrains)
 }
 
+func GetRunningModelSafetyTestTask() ([]*Cloudbrain, error) {
+	cloudbrains := make([]*Cloudbrain, 0)
+	sess := x.Where("job_type = ? and status=?", string(JobTypeModelSafety), JobRunning)
+	err := sess.Find(&cloudbrains)
+	return cloudbrains, err
+}
+
 func GetCloudbrainCountByUserID(userID int64, jobType string) (int, error) {
 	count, err := x.In("status", JobWaiting, JobRunning).And("job_type = ? and user_id = ? and type = ?", jobType, userID, TypeCloudBrainOne).Count(new(Cloudbrain))
 	return int(count), err

routers/repo/aisafety.go

@@ -108,9 +108,36 @@ func CloudBrainAiSafetyCreate(ctx *context.Context) {
 	} else if taskType == models.TypeCloudBrainOne {
 		createForGPU(ctx, jobName)
 	}
+	//ctx.Redirect(setting.AppSubURL + ctx.Repo.RepoLink + "/cloudbrain/benchmark")
+}
+
+func TimerHandleModelSafetyTestTask() {
+	tasks, err := models.GetRunningModelSafetyTestTask()
+	if err == nil {
+		if tasks != nil && len(tasks) > 0 {
+			for _, task := range tasks {
+				if task.Type == models.TypeCloudBrainOne {
+					if task.Status == string(models.JobRunning) {
+						//调用云脑任务状态接口
+
+					} else if task.Status == string(models.ModelSafetyTesting) {
+						//调用模型评测查询状态接口
+
+					}
+				} else if task.Type == models.TypeCloudBrainTwo {
+					if task.Status == string(models.ModelArtsRunning) {
 
-	ctx.Redirect(setting.AppSubURL + ctx.Repo.RepoLink + "/cloudbrain/benchmark")
+					} else if task.Status == string(models.ModelArtsTrainJobCompleted) {
 
+					}
+				}
+			}
+		} else {
+			log.Info("query running model safety test task 0.")
+		}
+	} else {
+		log.Info("query running model safety test task err." + err.Error())
+	}
 }
 
 func createForNPU(ctx *context.Context, jobName string) {