diff --git a/models/cloudbrain.go b/models/cloudbrain.go index 48ecc3998..997955baa 100755 --- a/models/cloudbrain.go +++ b/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 diff --git a/routers/repo/aisafety.go b/routers/repo/aisafety.go index 29ed58197..dc656033b 100644 --- a/routers/repo/aisafety.go +++ b/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) {