From f6b550aad2bc890ab194e494a6238a981ab9cc6c Mon Sep 17 00:00:00 2001
From: ychao_1983 <ychao_1983@sina.com>
Date: Fri, 2 Dec 2022 15:41:28 +0800
Subject: [PATCH] =?UTF-8?q?=E6=8F=90=E4=BA=A4=E4=BB=A3=E7=A0=81?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 modules/cloudbrain/cloudbrain.go  | 27 ++++++++++++++++++---------
 routers/api/v1/repo/cloudbrain.go |  2 +-
 routers/api/v1/repo/modelarts.go  |  2 +-
 3 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/modules/cloudbrain/cloudbrain.go b/modules/cloudbrain/cloudbrain.go
index b70ca2ada..6714c57ae 100755
--- a/modules/cloudbrain/cloudbrain.go
+++ b/modules/cloudbrain/cloudbrain.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"os"
 	"strconv"
+	"strings"
 
 	"code.gitea.io/gitea/modules/timeutil"
 
@@ -145,7 +146,7 @@ func isAdminOrImageCreater(ctx *context.Context, image *models.Image, err error)
 func AdminOrOwnerOrJobCreaterRight(ctx *context.Context) {
 
 	var id = ctx.Params(":id")
-	job, err := GetCloudBrainByIdOrJobId(id)
+	job, err := GetCloudBrainByIdOrJobId(id, "id")
 	if err != nil {
 		log.Error("GetCloudbrainByID failed:%v", err.Error())
 		ctx.NotFound(ctx.Req.URL.RequestURI(), nil)
@@ -161,7 +162,7 @@ func AdminOrOwnerOrJobCreaterRight(ctx *context.Context) {
 func AdminOrJobCreaterRight(ctx *context.Context) {
 
 	var id = ctx.Params(":id")
-	job, err := GetCloudBrainByIdOrJobId(id)
+	job, err := GetCloudBrainByIdOrJobId(id, "id")
 	if err != nil {
 		log.Error("GetCloudbrainByID failed:%v", err.Error())
 		ctx.NotFound(ctx.Req.URL.RequestURI(), nil)
@@ -177,7 +178,7 @@ func AdminOrJobCreaterRight(ctx *context.Context) {
 func AdminOrOwnerOrJobCreaterRightForTrain(ctx *context.Context) {
 
 	var jobID = ctx.Params(":jobid")
-	job, err := GetCloudBrainByIdOrJobId(jobID)
+	job, err := GetCloudBrainByIdOrJobId(jobID, "jobid")
 	if err != nil {
 		log.Error("GetCloudbrainByJobID failed:%v", err.Error())
 		ctx.NotFound(ctx.Req.URL.RequestURI(), nil)
@@ -193,7 +194,7 @@ func AdminOrOwnerOrJobCreaterRightForTrain(ctx *context.Context) {
 func AdminOrJobCreaterRightForTrain(ctx *context.Context) {
 
 	var jobID = ctx.Params(":jobid")
-	job, err := GetCloudBrainByIdOrJobId(jobID)
+	job, err := GetCloudBrainByIdOrJobId(jobID, "jobid")
 	if err != nil {
 		log.Error("GetCloudbrainByJobID failed:%v", err.Error())
 		ctx.NotFound(ctx.Req.URL.RequestURI(), nil)
@@ -499,7 +500,7 @@ func RestartTask(ctx *context.Context, task *models.Cloudbrain, newID *string) e
 			volumes = append(volumes, models.Volume{
 				HostPath: models.StHostPath{
 					Path:      realPath,
-					MountPath: ModelMountPath,
+					MountPath: PretrainModelMountPath,
 					ReadOnly:  true,
 				},
 			})
@@ -669,18 +670,26 @@ func IsElementExist(s []string, str string) bool {
 	return false
 }
 
-func GetCloudBrainByIdOrJobId(id string) (*models.Cloudbrain,error) {
+func GetCloudBrainByIdOrJobId(id string, initialQuery string) (*models.Cloudbrain, error) {
 	_, err := strconv.ParseInt(id, 10, 64)
 	var job *models.Cloudbrain
 	if err != nil {
 
 		job, err = models.GetCloudbrainByJobID(id)
 	} else {
-		job, err = models.GetCloudbrainByID(id)
-		if err!=nil{
+
+		if strings.EqualFold(initialQuery, "id") {
+			job, err = models.GetCloudbrainByID(id)
+			if err != nil {
+				job, err = models.GetCloudbrainByJobID(id)
+			}
+		} else {
 			job, err = models.GetCloudbrainByJobID(id)
+			if err != nil {
+				job, err = models.GetCloudbrainByID(id)
+			}
 		}
 
 	}
-	return job,err
+	return job, err
 }
diff --git a/routers/api/v1/repo/cloudbrain.go b/routers/api/v1/repo/cloudbrain.go
index cd8340c41..7bcaa7b26 100755
--- a/routers/api/v1/repo/cloudbrain.go
+++ b/routers/api/v1/repo/cloudbrain.go
@@ -211,7 +211,7 @@ func GetCloudbrainTask(ctx *context.APIContext) {
 
 	ID := ctx.Params(":id")
 
-	job, err := cloudbrain.GetCloudBrainByIdOrJobId(ID)
+	job, err := cloudbrain.GetCloudBrainByIdOrJobId(ID, "id")
 
 	if err != nil {
 		ctx.NotFound(err)
diff --git a/routers/api/v1/repo/modelarts.go b/routers/api/v1/repo/modelarts.go
index 127ddd835..a0abab38b 100755
--- a/routers/api/v1/repo/modelarts.go
+++ b/routers/api/v1/repo/modelarts.go
@@ -39,7 +39,7 @@ func GetModelArtsNotebook2(ctx *context.APIContext) {
 
 	ID := ctx.Params(":id")
 
-	job,err := cloudbrain.GetCloudBrainByIdOrJobId(ID)
+	job, err := cloudbrain.GetCloudBrainByIdOrJobId(ID, "id")
 
 	if err != nil {
 		ctx.NotFound(err)