From 7bd108696254cc960a3dbfef29cf64c8cb187684 Mon Sep 17 00:00:00 2001
From: zouap <zouap@pcl.ac.cn>
Date: Mon, 10 Oct 2022 10:18:17 +0800
Subject: [PATCH] =?UTF-8?q?=E5=AE=9A=E6=97=B6=E4=BB=BB=E5=8A=A1=E5=8F=8A?=
 =?UTF-8?q?=E5=81=9C=E6=AD=A2=E4=BB=BB=E5=8A=A1=E4=BF=AE=E6=94=B9=E3=80=82?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: zouap <zouap@pcl.ac.cn>
---
 modules/cron/tasks_basic.go | 20 +++++++++++++++++---
 modules/timer/timer.go      | 33 ---------------------------------
 routers/repo/aisafety.go    | 19 +++++++++++++++++++
 3 files changed, 36 insertions(+), 36 deletions(-)
 delete mode 100755 modules/timer/timer.go

diff --git a/modules/cron/tasks_basic.go b/modules/cron/tasks_basic.go
index 04cd7fe41..8dbc8d1ed 100755
--- a/modules/cron/tasks_basic.go
+++ b/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()
 }
diff --git a/modules/timer/timer.go b/modules/timer/timer.go
deleted file mode 100755
index 43535da96..000000000
--- a/modules/timer/timer.go
+++ /dev/null
@@ -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()
-}
diff --git a/routers/repo/aisafety.go b/routers/repo/aisafety.go
index e628bb88c..ccb24d7e1 100644
--- a/routers/repo/aisafety.go
+++ b/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) {