From b2e2d81f0df512fe4e25fd340853d1d2093d8257 Mon Sep 17 00:00:00 2001
From: ychao_1983 <ychao_1983@sina.com>
Date: Wed, 13 Jul 2022 12:57:09 +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 | 20 ++++++++++++++----
 routers/repo/cloudbrain.go       | 36 ++++++++++++++++++++++++++++++--
 2 files changed, 50 insertions(+), 6 deletions(-)

diff --git a/modules/cloudbrain/cloudbrain.go b/modules/cloudbrain/cloudbrain.go
index 73dcb150c..f9c7c5b5c 100755
--- a/modules/cloudbrain/cloudbrain.go
+++ b/modules/cloudbrain/cloudbrain.go
@@ -43,9 +43,10 @@ const (
 )
 
 var (
-	ResourceSpecs      *models.ResourceSpecs
-	TrainResourceSpecs *models.ResourceSpecs
-	SpecialPools       *models.SpecialPools
+	ResourceSpecs          *models.ResourceSpecs
+	TrainResourceSpecs     *models.ResourceSpecs
+	InferenceResourceSpecs *models.ResourceSpecs
+	SpecialPools           *models.SpecialPools
 )
 
 type GenerateCloudBrainTaskReq struct {
@@ -222,7 +223,7 @@ func AdminOrImageCreaterRight(ctx *context.Context) {
 func GenerateTask(req GenerateCloudBrainTaskReq) error {
 	var resourceSpec *models.ResourceSpec
 	var versionCount int
-	if req.JobType == string(models.JobTypeTrain) || req.JobType == string(models.JobTypeInference) {
+	if req.JobType == string(models.JobTypeTrain) {
 		versionCount = 1
 		if TrainResourceSpecs == nil {
 			json.Unmarshal([]byte(setting.TrainResourceSpecs), &TrainResourceSpecs)
@@ -233,6 +234,17 @@ func GenerateTask(req GenerateCloudBrainTaskReq) error {
 				break
 			}
 		}
+	} else if req.JobType == string(models.JobTypeInference) {
+		if InferenceResourceSpecs == nil {
+			json.Unmarshal([]byte(setting.InferenceResourceSpecs), &InferenceResourceSpecs)
+		}
+		for _, spec := range InferenceResourceSpecs.ResourceSpec {
+			if req.ResourceSpecId == spec.Id {
+				resourceSpec = spec
+				break
+			}
+		}
+
 	} else {
 		if ResourceSpecs == nil {
 			json.Unmarshal([]byte(setting.ResourceSpecs), &ResourceSpecs)
diff --git a/routers/repo/cloudbrain.go b/routers/repo/cloudbrain.go
index 6af613649..e8ea71ab7 100755
--- a/routers/repo/cloudbrain.go
+++ b/routers/repo/cloudbrain.go
@@ -59,6 +59,7 @@ var (
 	benchmarkGpuInfos      *models.GpuInfos
 	benchmarkResourceSpecs *models.ResourceSpecs
 	trainGpuInfos          *models.GpuInfos
+	inferenceGpuInfos      *models.GpuInfos
 )
 
 const BENCHMARK_TYPE_CODE = "repo.cloudbrain.benchmark.types"
@@ -130,6 +131,11 @@ func cloudBrainNewDataPrepare(ctx *context.Context) error {
 	}
 	ctx.Data["train_gpu_types"] = trainGpuInfos.GpuInfo
 
+	if inferenceGpuInfos == nil {
+		json.Unmarshal([]byte(setting.InferenceGpuTypes), &inferenceGpuInfos)
+	}
+	ctx.Data["train_gpu_types"] = trainGpuInfos.GpuInfo
+
 	if benchmarkGpuInfos == nil {
 		json.Unmarshal([]byte(setting.BenchmarkGpuTypes), &benchmarkGpuInfos)
 	}
@@ -150,6 +156,11 @@ func cloudBrainNewDataPrepare(ctx *context.Context) error {
 	}
 	ctx.Data["train_resource_specs"] = cloudbrain.TrainResourceSpecs.ResourceSpec
 
+	if cloudbrain.InferenceResourceSpecs == nil {
+		json.Unmarshal([]byte(setting.InferenceResourceSpecs), &cloudbrain.InferenceResourceSpecs)
+	}
+	ctx.Data["inference_resource_specs"] = cloudbrain.InferenceResourceSpecs.ResourceSpec
+
 	if cloudbrain.SpecialPools != nil {
 		var debugGpuTypes []*models.GpuInfo
 		var trainGpuTypes []*models.GpuInfo
@@ -610,7 +621,7 @@ func cloudBrainShow(ctx *context.Context, tpName base.TplName, jobType models.Jo
 
 	var task *models.Cloudbrain
 	var err error
-	if jobType == models.JobTypeTrain {
+	if jobType == models.JobTypeTrain || jobType == models.JobTypeInference {
 		task, err = models.GetCloudbrainByJobID(ctx.Params(":jobid"))
 	} else {
 		task, err = models.GetCloudbrainByIDWithDeleted(ctx.Params(":id"))
@@ -641,6 +652,18 @@ func cloudBrainShow(ctx *context.Context, tpName base.TplName, jobType models.Jo
 				ctx.Data["ShareMemMiB"] = tmp.ShareMemMiB
 			}
 		}
+	} else if task.JobType == string(models.JobTypeInference) {
+		if cloudbrain.InferenceResourceSpecs == nil {
+			json.Unmarshal([]byte(setting.InferenceResourceSpecs), &cloudbrain.InferenceResourceSpecs)
+		}
+		for _, tmp := range cloudbrain.InferenceResourceSpecs.ResourceSpec {
+			if tmp.Id == task.ResourceSpecId {
+				ctx.Data["GpuNum"] = tmp.GpuNum
+				ctx.Data["CpuNum"] = tmp.CpuNum
+				ctx.Data["MemMiB"] = tmp.MemMiB
+				ctx.Data["ShareMemMiB"] = tmp.ShareMemMiB
+			}
+		}
 	} else {
 		if cloudbrain.ResourceSpecs == nil {
 			json.Unmarshal([]byte(setting.ResourceSpecs), &cloudbrain.ResourceSpecs)
@@ -669,6 +692,15 @@ func cloudBrainShow(ctx *context.Context, tpName base.TplName, jobType models.Jo
 					ctx.Data["resource_type"] = resourceType.Value
 				}
 			}
+		} else if task.JobType == string(models.JobTypeInference) {
+			if inferenceGpuInfos == nil {
+				json.Unmarshal([]byte(setting.InferenceGpuTypes), &inferenceGpuInfos)
+			}
+			for _, resourceType := range inferenceGpuInfos.GpuInfo {
+				if resourceType.Queue == jobRes.Config.GpuType {
+					ctx.Data["resource_type"] = resourceType.Value
+				}
+			}
 		} else if cloudbrain.IsBenchmarkJob(task.JobType) {
 			if benchmarkGpuInfos == nil {
 				json.Unmarshal([]byte(setting.BenchmarkGpuTypes), &benchmarkGpuInfos)
@@ -2472,7 +2504,7 @@ func InferenceCloudBrainJobNew(ctx *context.Context) {
 }
 
 func InferenceCloudBrainJobShow(ctx *context.Context) {
-	cloudBrainShow(ctx, tplCloudBrainInferenceJobShow, models.JobTypeTrain)
+	cloudBrainShow(ctx, tplCloudBrainInferenceJobShow, models.JobTypeInference)
 }
 
 func DownloadInferenceResultFile(ctx *context.Context) {