From 9d23ea7a4b97303c2fd3fb0409f10cd7e05b698b Mon Sep 17 00:00:00 2001
From: ychao_1983 <ychao_1983@sina.com>
Date: Mon, 14 Nov 2022 15:28:59 +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

---
 routers/api/v1/api.go             |  2 +-
 routers/api/v1/repo/cloudbrain.go | 59 +++++++++++++++++++++++++++++--
 2 files changed, 57 insertions(+), 4 deletions(-)

diff --git a/routers/api/v1/api.go b/routers/api/v1/api.go
index b7e3c6415..96de00c55 100755
--- a/routers/api/v1/api.go
+++ b/routers/api/v1/api.go
@@ -737,7 +737,7 @@ func RegisterRoutes(m *macaron.Macaron) {
 		}, reqToken(), repoAssignment())
 
 		m.Group("/file_notebook", func() {
-
+			m.Get("", reqToken(), repo.GetFileNoteBookInfo)
 			m.Post("/create", reqToken(), reqWeChat(), bind(api.CreateFileNotebookJobOption{}), repo.CreateFileNoteBook)
 
 		})
diff --git a/routers/api/v1/repo/cloudbrain.go b/routers/api/v1/repo/cloudbrain.go
index 5ea797e97..22b533baa 100755
--- a/routers/api/v1/repo/cloudbrain.go
+++ b/routers/api/v1/repo/cloudbrain.go
@@ -83,6 +83,56 @@ func CreateFileNoteBook(ctx *context.APIContext, option api.CreateFileNotebookJo
 	cloudbrainTask.FileNotebookCreate(ctx.Context, option)
 }
 
+func GetFileNoteBookInfo(ctx *context.APIContext) {
+	//image description  spec description  waiting count
+
+	specs, err := models.GetResourceSpecificationByIds([]int64{setting.FileNoteBook.SpecIdCPU, setting.FileNoteBook.SpecIdGPU, setting.FileNoteBook.SpecIdNPU, setting.FileNoteBook.SpecIdNPUCD})
+	if err != nil {
+		log.Error("Fail to query specifications", err)
+		ctx.JSON(http.StatusOK, models.BaseErrorMessageApi(ctx.Tr("repo.notebook_query_fail")))
+		return
+	}
+
+	var specCPU, specGpu, specNPU, specNPUCD *api.SpecificationShow
+	var specGpuQueueCode string
+	for _, spec := range specs {
+		if spec.ID == setting.FileNoteBook.SpecIdCPU {
+			specCPU = convert.ToSpecification(spec)
+		} else if spec.ID == setting.FileNoteBook.SpecIdGPU {
+			specGpu = convert.ToSpecification(spec)
+			specGpuQueueCode = spec.QueueCode
+		} else if spec.ID == setting.FileNoteBook.SpecIdNPU {
+			specNPU = convert.ToSpecification(spec)
+		} else if spec.ID == setting.FileNoteBook.SpecIdNPUCD {
+			specNPUCD = convert.ToSpecification(spec)
+		}
+	}
+
+	waitCountNpu := cloudbrain.GetWaitingCloudbrainCount(models.TypeCloudBrainTwo, "")
+
+	queuesMap, err := cloudbrain.GetQueuesDetail()
+	if err != nil {
+		log.Error("Fail to query gpu queues waiting count", err)
+		ctx.JSON(http.StatusOK, models.BaseErrorMessageApi(ctx.Tr("repo.notebook_query_fail")))
+		return
+	}
+	waitCountGPU := (*queuesMap)[specGpuQueueCode]
+	ctx.JSON(http.StatusOK, map[string]interface{}{
+		"code":                  0,
+		"specCpu":               specCPU,
+		"specGpu":               specGpu,
+		"specNpu":               specNPU,
+		"specNpuCd":             specNPUCD,
+		"waitCountGpu":          waitCountGPU,
+		"waitCountNpu":          waitCountNpu,
+		"imageCpuDescription":   setting.FileNoteBook.ImageCPUDescription,
+		"imageGpuDescription":   setting.FileNoteBook.ImageGPUDescription,
+		"imageNpuDescription":   setting.FileNoteBook.ImageNPUDescription,
+		"imageNpuCDDescription": setting.FileNoteBook.ImageNPUCDDescription,
+	})
+
+}
+
 func CreateCloudBrain(ctx *context.APIContext, option api.CreateTrainJobOption) {
 	if option.Type == cloudbrainTask.TaskTypeCloudbrainOne {
 		cloudbrainTask.CloudbrainOneTrainJobCreate(ctx.Context, option)
@@ -146,9 +196,12 @@ func GetCloudbrainTask(ctx *context.APIContext) {
 	ID := ctx.Params(":id")
 	job, err := models.GetCloudbrainByID(ID)
 	if err != nil {
-		ctx.NotFound(err)
-		log.Error("GetCloudbrainByID failed:", err)
-		return
+		job, err = models.GetCloudbrainByJobID(ID)
+		if err != nil {
+			ctx.NotFound(err)
+			log.Error("GetCloudbrainByID failed:", err)
+			return
+		}
 	}
 	if job.JobType == string(models.JobTypeModelSafety) {
 		routerRepo.GetAiSafetyTaskByJob(job)