定时任务及停止任务修改。

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

modules/cron/tasks_basic.go

@@ -5,12 +5,13 @@
 package cron
 
 import (
-	"code.gitea.io/gitea/services/reward"
-	"code.gitea.io/gitea/services/cloudbrain/resource"
-	"code.gitea.io/gitea/modules/modelarts"
 	"context"
 	"time"
 
+	"code.gitea.io/gitea/modules/modelarts"
+	"code.gitea.io/gitea/services/cloudbrain/resource"
+	"code.gitea.io/gitea/services/reward"
+
 	"code.gitea.io/gitea/models"
 	"code.gitea.io/gitea/modules/migrations"
 	repository_service "code.gitea.io/gitea/modules/repository"
@@ -122,6 +123,17 @@ func registerHandleUnDecompressAttachment() {
 	})
 }
 
+func registerHandleModelSafetyTask() {
+	RegisterTaskFatal("handle_modelsafety_task", &BaseConfig{
+		Enabled:    true,
+		RunAtStart: true,
+		Schedule:   "@every 5m",
+	}, func(ctx context.Context, _ *models.User, _ Config) error {
+		repo.TimerHandleModelSafetyTestTask()
+		return nil
+	})
+}
+
 func registerHandleBlockChainUnSuccessUsers() {
 	RegisterTaskFatal("handle_blockchain_unsuccess_users", &BaseConfig{
 		Enabled:    true,
@@ -279,4 +291,6 @@ func initBasicTasks() {
 
 	//registerRewardPeriodTask()
 	registerCloudbrainPointDeductTask()
+
+	registerHandleModelSafetyTask()
 }

modules/timer/timer.go

@@ -1,33 +0,0 @@
-package timer
-
-import (
-	"github.com/robfig/cron/v3"
-
-	"code.gitea.io/gitea/modules/log"
-	"code.gitea.io/gitea/routers/repo"
-)
-
-func LaunchCronJob() {
-	log.Trace("Run cron job")
-
-	c := cron.New()
-
-	spec := "*/10 * * * *"
-	c.AddFunc(spec, repo.HandleUnDecompressAttachment)
-
-	specCheckBlockChainUserSuccess := "*/10 * * * *"
-	c.AddFunc(specCheckBlockChainUserSuccess, repo.HandleBlockChainUnSuccessUsers)
-
-	specModelSafetySuccess := "*/2 * * * *"
-	c.AddFunc(specModelSafetySuccess, repo.TimerHandleModelSafetyTestTask)
-
-	specCheckRepoBlockChainSuccess := "*/1 * * * *"
-	c.AddFunc(specCheckRepoBlockChainSuccess, repo.HandleBlockChainUnSuccessRepos)
-
-	specCheckUnTransformedPRs := "*/1 * * * *"
-	c.AddFunc(specCheckUnTransformedPRs, repo.HandleBlockChainMergedPulls)
-
-	specCheckBlockChainCommitSuccess := "*/3 * * * *"
-	c.AddFunc(specCheckBlockChainCommitSuccess, repo.HandleBlockChainUnSuccessCommits)
-	c.Start()
-}

routers/repo/aisafety.go

@@ -99,11 +99,13 @@ func GetAiSafetyTask(ctx *context.Context) {
 }
 
 func StopAiSafetyTask(ctx *context.Context) {
+	log.Info("start to stop the task.")
 	var ID = ctx.Params(":jobid")
 	task, err := models.GetCloudbrainByIDWithDeleted(ID)
 	result := make(map[string]interface{})
 	result["code"] = -1
 	if err != nil {
+		log.Info("query task error.err=" + err.Error())
 		log.Error("GetCloudbrainByJobID failed:" + err.Error())
 		result["msg"] = "No such task."
 		ctx.JSON(200, result)
@@ -111,6 +113,23 @@ func StopAiSafetyTask(ctx *context.Context) {
 	}
 	if isTaskNotFinished(task.Status) {
 		if task.Type == models.TypeCloudBrainTwo {
+			log.Info("start to stop model arts task.")
+			_, err := modelarts.StopTrainJob(task.JobID, strconv.FormatInt(task.VersionID, 10))
+			if err != nil {
+				log.Info("stop failed.err=" + err.Error())
+			}
+			task.Status = string(models.JobStopped)
+			if task.EndTime == 0 {
+				task.EndTime = timeutil.TimeStampNow()
+			}
+			task.ComputeAndSetDuration()
+			err = models.UpdateJob(task)
+			if err != nil {
+				log.Error("UpdateJob(%s) failed:%v", task.JobName, err, ctx.Data["msgID"])
+				result["msg"] = "cloudbrain.Stopped_success_update_status_fail"
+				ctx.JSON(200, result)
+				return
+			}
 			//queryTaskStatusFromCloudbrainTwo(job)
 		} else if task.Type == models.TypeCloudBrainOne {
 			if task.Status == string(models.JobStopped) || task.Status == string(models.JobFailed) || task.Status == string(models.JobSucceeded) {