From 57b1a05ab5e953e8623ddf43ca443017803022a7 Mon Sep 17 00:00:00 2001
From: liuzx <liuzx@pcl.ac.cn>
Date: Thu, 11 Nov 2021 11:01:40 +0800
Subject: [PATCH] =?UTF-8?q?=E5=88=B7=E6=96=B0=E8=AE=AD=E7=BB=83=E7=89=88?=
 =?UTF-8?q?=E6=9C=AC=E7=9A=84=E7=8A=B6=E6=80=81=E5=92=8C=E6=97=B6=E9=95=BF?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 models/cloudbrain.go             |  5 +++
 routers/api/v1/api.go            |  3 ++
 routers/api/v1/repo/modelarts.go | 62 +++++++++++++++++++++++++-------
 3 files changed, 57 insertions(+), 13 deletions(-)

diff --git a/models/cloudbrain.go b/models/cloudbrain.go
index a061857c4..366ba74e4 100755
--- a/models/cloudbrain.go
+++ b/models/cloudbrain.go
@@ -1066,6 +1066,11 @@ func GetRepoCloudBrainByJobID(repoID int64, jobID string) (*Cloudbrain, error) {
 	return getRepoCloudBrain(cb)
 }
 
+func GetRepoCloudBrainByJobIDAndVersionName(repoID int64, jobID string, versionName string) (*Cloudbrain, error) {
+	cb := &Cloudbrain{JobID: jobID, RepoID: repoID, VersionName: versionName}
+	return getRepoCloudBrain(cb)
+}
+
 func GetCloudbrainByJobID(jobID string) (*Cloudbrain, error) {
 	cb := &Cloudbrain{JobID: jobID}
 	return getRepoCloudBrain(cb)
diff --git a/routers/api/v1/api.go b/routers/api/v1/api.go
index 8b068d612..3842907cf 100755
--- a/routers/api/v1/api.go
+++ b/routers/api/v1/api.go
@@ -875,6 +875,9 @@ func RegisterRoutes(m *macaron.Macaron) {
 						m.Group("/:jobid", func() {
 							m.Get("", repo.GetModelArtsTrainJob)
 							m.Get("/log", repo.TrainJobGetLog)
+							m.Group("/:version-name", func() {
+								m.Get("", repo.GetModelArtsTrainJobVersion)
+							})
 						})
 					})
 				}, reqRepoReader(models.UnitTypeCloudBrain))
diff --git a/routers/api/v1/repo/modelarts.go b/routers/api/v1/repo/modelarts.go
index 2e825d8cc..6d9bba64b 100755
--- a/routers/api/v1/repo/modelarts.go
+++ b/routers/api/v1/repo/modelarts.go
@@ -6,12 +6,13 @@
 package repo
 
 import (
+	"net/http"
+	"strconv"
+
 	"code.gitea.io/gitea/models"
 	"code.gitea.io/gitea/modules/context"
 	"code.gitea.io/gitea/modules/log"
 	"code.gitea.io/gitea/modules/modelarts"
-	"net/http"
-	"strconv"
 )
 
 func GetModelArtsNotebook(ctx *context.APIContext) {
@@ -72,9 +73,44 @@ func GetModelArtsTrainJob(ctx *context.APIContext) {
 	}
 
 	ctx.JSON(http.StatusOK, map[string]interface{}{
-		"JobID":         jobID,
-		"JobStatus":     job.Status,
-		"JobDuration":   job.Duration,
+		"JobID":       jobID,
+		"JobStatus":   job.Status,
+		"JobDuration": job.Duration,
+	})
+
+}
+
+func GetModelArtsTrainJobVersion(ctx *context.APIContext) {
+	var (
+		err error
+	)
+
+	jobID := ctx.Params(":jobid")
+	versionName := ctx.Params(":version-name")
+	repoID := ctx.Repo.Repository.ID
+	job, err := models.GetRepoCloudBrainByJobIDAndVersionName(repoID, jobID, versionName)
+	if err != nil {
+		ctx.NotFound(err)
+		return
+	}
+	result, err := modelarts.GetTrainJob(jobID, strconv.FormatInt(job.VersionID, 10))
+	if err != nil {
+		ctx.NotFound(err)
+		return
+	}
+
+	job.Status = modelarts.TransTrainJobStatus(result.IntStatus)
+	job.Duration = result.Duration
+	job.TrainJobDuration = result.TrainJobDuration
+	err = models.UpdateJob(job)
+	if err != nil {
+		log.Error("UpdateJob failed:", err)
+	}
+
+	ctx.JSON(http.StatusOK, map[string]interface{}{
+		"JobID":       jobID,
+		"JobStatus":   job.Status,
+		"JobDuration": job.Duration,
 	})
 
 }
@@ -94,7 +130,7 @@ func TrainJobGetLog(ctx *context.APIContext) {
 	if order != modelarts.OrderDesc && order != modelarts.OrderAsc {
 		log.Error("order(%s) check failed", order)
 		ctx.JSON(http.StatusBadRequest, map[string]interface{}{
-			"err_msg":         "order check failed",
+			"err_msg": "order check failed",
 		})
 		return
 	}
@@ -103,7 +139,7 @@ func TrainJobGetLog(ctx *context.APIContext) {
 	if err != nil {
 		log.Error("GetCloudbrainByJobID(%s) failed:%v", jobID, err.Error())
 		ctx.JSON(http.StatusInternalServerError, map[string]interface{}{
-			"err_msg":         "GetCloudbrainByJobID failed",
+			"err_msg": "GetCloudbrainByJobID failed",
 		})
 		return
 	}
@@ -112,16 +148,16 @@ func TrainJobGetLog(ctx *context.APIContext) {
 	if err != nil {
 		log.Error("GetTrainJobLog(%s) failed:%v", jobID, err.Error())
 		ctx.JSON(http.StatusInternalServerError, map[string]interface{}{
-			"err_msg":         "GetTrainJobLog failed",
+			"err_msg": "GetTrainJobLog failed",
 		})
 		return
 	}
 
 	ctx.JSON(http.StatusOK, map[string]interface{}{
-		"JobID":         jobID,
-		"StartLine":     result.StartLine,
-		"EndLine":     	 result.EndLine,
-		"Content":     	 result.Content,
-		"Lines":     	 result.Lines,
+		"JobID":     jobID,
+		"StartLine": result.StartLine,
+		"EndLine":   result.EndLine,
+		"Content":   result.Content,
+		"Lines":     result.Lines,
 	})
 }