From 15f1289b6bb013a9789f15fb77ef5d151f9b992f Mon Sep 17 00:00:00 2001
From: lewis <747342561@qq.com>
Date: Wed, 5 Jan 2022 11:30:25 +0800
Subject: [PATCH] restart opt

---
 models/cloudbrain.go             | 25 +++++++++++++++++++++++++
 modules/cloudbrain/cloudbrain.go | 27 +++++++++++++++++++++------
 routers/repo/cloudbrain.go       |  7 ++-----
 3 files changed, 48 insertions(+), 11 deletions(-)

diff --git a/models/cloudbrain.go b/models/cloudbrain.go
index c3707c9e2..efaa9ffeb 100755
--- a/models/cloudbrain.go
+++ b/models/cloudbrain.go
@@ -1210,3 +1210,28 @@ func GetCloudbrainTrainJobCountByUserID(userID int64) (int, error) {
 		And("job_type = ? and user_id = ? and type = ?", JobTypeTrain, userID, TypeCloudBrainTwo).Count(new(Cloudbrain))
 	return int(count), err
 }
+
+func RestartCloudbrain(old *Cloudbrain, new *Cloudbrain) (err error) {
+	sess := x.NewSession()
+	defer sess.Close()
+
+	if err = sess.Begin(); err != nil {
+		return err
+	}
+
+	if _, err = sess.Delete(old); err != nil {
+		sess.Rollback()
+		return err
+	}
+
+	if _, err = sess.Insert(new); err != nil {
+		sess.Rollback()
+		return err
+	}
+
+	if err = sess.Commit(); err != nil {
+		return err
+	}
+
+	return nil
+}
diff --git a/modules/cloudbrain/cloudbrain.go b/modules/cloudbrain/cloudbrain.go
index 40da8a8c7..74dcbe7b0 100755
--- a/modules/cloudbrain/cloudbrain.go
+++ b/modules/cloudbrain/cloudbrain.go
@@ -223,7 +223,7 @@ func GenerateTask(ctx *context.Context, jobName, image, command, uuid, codePath,
 	return nil
 }
 
-func RestartTask(ctx *context.Context, task *models.Cloudbrain) error {
+func RestartTask(ctx *context.Context, task *models.Cloudbrain, newJobID *string) error {
 	dataActualPath := setting.Attachment.Minio.RealPath +
 		setting.Attachment.Minio.Bucket + "/" +
 		setting.Attachment.Minio.BasePath +
@@ -313,7 +313,7 @@ func RestartTask(ctx *context.Context, task *models.Cloudbrain) error {
 		},
 	})
 	if err != nil {
-		log.Error("CreateJob failed:", err.Error(), ctx.Data["MsgID"])
+		log.Error("CreateJob failed:%v", err.Error(), ctx.Data["MsgID"])
 		return err
 	}
 	if jobResult.Code != Success {
@@ -322,14 +322,29 @@ func RestartTask(ctx *context.Context, task *models.Cloudbrain) error {
 	}
 
 	var jobID = jobResult.Payload["jobId"].(string)
-	task.JobID = jobID
-	task.Status = string(models.JobWaiting)
-	err =  models.UpdateJob(task)
+	newTask := &models.Cloudbrain{
+		Status:         string(models.JobWaiting),
+		UserID:         task.UserID,
+		RepoID:         task.RepoID,
+		JobID:          jobID,
+		JobName:        task.JobName,
+		SubTaskName:    task.SubTaskName,
+		JobType:        task.JobType,
+		Type:           task.Type,
+		Uuid:           task.Uuid,
+		Image:          task.Image,
+		GpuQueue:       task.GpuQueue,
+		ResourceSpecId: task.ResourceSpecId,
+		ComputeResource: task.ComputeResource,
+	}
 
+	err = models.RestartCloudbrain(task, newTask)
 	if err != nil {
-		log.Error("UpdateJob(%s) failed:%v", jobName, err.Error(), ctx.Data["MsgID"])
+		log.Error("RestartCloudbrain(%s) failed:%v", jobName, err.Error(), ctx.Data["MsgID"])
 		return err
 	}
 
+	*newJobID = jobID
+
 	return nil
 }
diff --git a/routers/repo/cloudbrain.go b/routers/repo/cloudbrain.go
index 000b6c83c..c8818d9ba 100755
--- a/routers/repo/cloudbrain.go
+++ b/routers/repo/cloudbrain.go
@@ -251,7 +251,7 @@ func CloudBrainRestart(ctx *context.Context) {
 	var jobID = ctx.Params(":jobid")
 	var resultCode = "0"
 	var errorMsg = ""
-	var status = ""
+	var status = string(models.JobWaiting)
 
 	task := ctx.Cloudbrain
 	for {
@@ -291,7 +291,7 @@ func CloudBrainRestart(ctx *context.Context) {
 			}
 		}
 
-		err = cloudbrain.RestartTask(ctx, task)
+		err = cloudbrain.RestartTask(ctx, task, &jobID)
 		if err != nil {
 			log.Error("RestartTask failed:%v", err.Error(), ctx.Data["MsgID"])
 			resultCode = "-1"
@@ -299,9 +299,6 @@ func CloudBrainRestart(ctx *context.Context) {
 			break
 		}
 
-		status = task.Status
-		jobID = task.JobID
-
 		break
 	}