From 39aca553f7497ef072fa0e51be472fdbad5c518f Mon Sep 17 00:00:00 2001
From: liuzx <liuzx@pcl.ac.cn>
Date: Tue, 8 Nov 2022 18:03:48 +0800
Subject: [PATCH] fix-3173

---
 models/cloudbrain.go        |  7 +++++++
 routers/private/internal.go |  4 +++-
 routers/repo/cloudbrain.go  |  3 ++-
 routers/repo/grampus.go     | 36 ++++++++++++++++++++++++++++++++++++
 4 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/models/cloudbrain.go b/models/cloudbrain.go
index dacb1b03a..a9182b7b8 100755
--- a/models/cloudbrain.go
+++ b/models/cloudbrain.go
@@ -2030,6 +2030,13 @@ func GetStoppedJobWithNoStartTimeEndTime() ([]*Cloudbrain, error) {
 	cloudbrains := make([]*Cloudbrain, 0)
 	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 GetC2NetWithAiCenterWrongJob() ([]*Cloudbrain, error) {
+	cloudbrains := make([]*Cloudbrain, 0)
+	return cloudbrains, x.
+		In("status", ModelArtsTrainJobFailed, ModelArtsTrainJobKilled, ModelArtsStopped, JobStopped, JobFailed).
+		Where("type = 2 ").
+		Find(&cloudbrains)
+}
 
 func GetModelSafetyTestTask() ([]*Cloudbrain, error) {
 	cloudbrains := make([]*Cloudbrain, 0)
diff --git a/routers/private/internal.go b/routers/private/internal.go
index 3e2eeab31..3d67afe8f 100755
--- a/routers/private/internal.go
+++ b/routers/private/internal.go
@@ -6,9 +6,10 @@
 package private
 
 import (
-	"code.gitea.io/gitea/routers/admin"
 	"strings"
 
+	"code.gitea.io/gitea/routers/admin"
+
 	"code.gitea.io/gitea/routers/repo"
 
 	"code.gitea.io/gitea/modules/log"
@@ -52,6 +53,7 @@ func RegisterRoutes(m *macaron.Macaron) {
 		m.Get("/tool/org_stat", OrgStatisticManually)
 		m.Post("/tool/update_repo_visit/:date", UpdateRepoVisit)
 		m.Post("/task/history_handle/duration", repo.HandleTaskWithNoDuration)
+		m.Post("/task/history_handle/aicenter", repo.HandleTaskWithAiCenter)
 		m.Post("/resources/specification/handle_historical_task", admin.RefreshHistorySpec)
 
 	}, CheckInternalToken)
diff --git a/routers/repo/cloudbrain.go b/routers/repo/cloudbrain.go
index 61c2925b8..7d96d1b58 100755
--- a/routers/repo/cloudbrain.go
+++ b/routers/repo/cloudbrain.go
@@ -2,7 +2,6 @@ package repo
 
 import (
 	"bufio"
-	"code.gitea.io/gitea/modules/urfs_client/urchin"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -16,6 +15,8 @@ import (
 	"time"
 	"unicode/utf8"
 
+	"code.gitea.io/gitea/modules/urfs_client/urchin"
+
 	"code.gitea.io/gitea/modules/dataset"
 
 	"code.gitea.io/gitea/services/cloudbrain/cloudbrainTask"
diff --git a/routers/repo/grampus.go b/routers/repo/grampus.go
index bd10d167a..ae4ad8a76 100755
--- a/routers/repo/grampus.go
+++ b/routers/repo/grampus.go
@@ -13,6 +13,7 @@ import (
 	"time"
 
 	"code.gitea.io/gitea/modules/urfs_client/urchin"
+	"code.gitea.io/gitea/routers/response"
 
 	"code.gitea.io/gitea/services/cloudbrain/cloudbrainTask"
 
@@ -1103,3 +1104,38 @@ func downloadZipCode(ctx *context.Context, codePath, branchName string) error {
 
 	return nil
 }
+func HandleTaskWithAiCenter(ctx *context.Context) {
+	log.Info("HandleTaskWithAiCenter start")
+	updateCounts := 0
+	cloudBrains, err := models.GetC2NetWithAiCenterWrongJob()
+	if err != nil {
+		log.Error("GetC2NetWithAiCenterWrongJob failed:" + err.Error())
+		return
+	}
+	if len(cloudBrains) == 0 {
+		log.Info("HandleC2NetWithAiCenterWrongJob:no task need handle")
+		return
+	}
+	cloudBrainCounts := len(cloudBrains)
+	for _, task := range cloudBrains {
+		result, err := grampus.GetJob(task.JobID)
+		if err != nil {
+			log.Error("GetJob failed:" + err.Error())
+			continue
+		}
+		if result != nil {
+			if len(result.JobInfo.Tasks[0].CenterID) == 1 && len(result.JobInfo.Tasks[0].CenterName) == 1 {
+				task.AiCenter = result.JobInfo.Tasks[0].CenterID[0] + "+" + result.JobInfo.Tasks[0].CenterName[0]
+			}
+			err = models.UpdateJob(task)
+			if err != nil {
+				log.Error("UpdateJob failed:" + err.Error())
+			}
+			updateCounts++
+		}
+	}
+	r := make(map[string]interface{}, 0)
+	r["cloudBrainCounts"] = cloudBrainCounts
+	r["updateCounts"] = updateCounts
+	ctx.JSON(http.StatusOK, response.SuccessWithData(r))
+}