Merge pull request 'benchmark-add-task-info' (#95) from benchmark-add-task-info into develop

Reviewed-on: https://git.openi.org.cn/OpenI/aiforge/pulls/95
Reviewed-by: avadesian <xuchx@pcl.ac.cn>

models/cloudbrain.go

@@ -167,6 +167,7 @@ type TaskInfo struct {
 	CodeName          string   `json:"code_name"`
 	BenchmarkCategory []string `json:"selected_category"`
 	CodeLink          string   `json:"code_link"`
+	GpuType           string   `json:"gpu_type"`
 }
 
 func ConvertToTaskPod(input map[string]interface{}) (TaskPod, error) {
@@ -276,11 +277,11 @@ type Category struct {
 	Value string `json:"value"`
 }
 
-type GpuTypes struct {
-	GpuType []*GpuType `json:"gpu_type"`
+type GpuInfos struct {
+	GpuInfo []*GpuInfo `json:"gpu_type"`
 }
 
-type GpuType struct {
+type GpuInfo struct {
 	Id    int    `json:"id"`
 	Value string `json:"value"`
 	Queue string `json:"queue"`

modules/cloudbrain/cloudbrain.go

@@ -23,7 +23,7 @@ const (
 	Success = "S000"
 )
 
-func GenerateTask(ctx *context.Context, jobName, image, command, uuid, codePath, modelPath, benchmarkPath, snn4imagenetPath, jobType, gpuType string) error {
+func GenerateTask(ctx *context.Context, jobName, image, command, uuid, codePath, modelPath, benchmarkPath, snn4imagenetPath, jobType, gpuQueue string) error {
 	dataActualPath := setting.Attachment.Minio.RealPath +
 		setting.Attachment.Minio.Bucket + "/" +
 		setting.Attachment.Minio.BasePath +
@@ -32,7 +32,7 @@ func GenerateTask(ctx *context.Context, jobName, image, command, uuid, codePath,
 	jobResult, err := CreateJob(jobName, models.CreateJobParams{
 		JobName:    jobName,
 		RetryCount: 1,
-		GpuType:    gpuType,
+		GpuType:    gpuQueue,
 		Image:      image,
 		TaskRoles: []models.TaskRole{
 			{

routers/repo/cloudbrain.go

@@ -25,6 +25,11 @@ const (
 	tplCloudBrainShow  base.TplName = "repo/cloudbrain/show"
 )
 
+var (
+	gpuInfos   *models.GpuInfos
+	categories *models.Categories
+)
+
 // MustEnableDataset check if repository enable internal cb
 func MustEnableCloudbrain(ctx *context.Context) {
 	if !ctx.Repo.CanRead(models.UnitTypeCloudBrain) {
@@ -131,12 +136,15 @@ func CloudBrainNew(ctx *context.Context) {
 	ctx.Data["benchmark_path"] = cloudbrain.BenchMarkMountPath
 	ctx.Data["is_benchmark_enabled"] = setting.IsBenchmarkEnabled
 
-	var categories *models.Categories
-	json.Unmarshal([]byte(setting.BenchmarkCategory), &categories)
+	if categories == nil {
+		json.Unmarshal([]byte(setting.BenchmarkCategory), &categories)
+	}
 	ctx.Data["benchmark_categories"] = categories.Category
-	var gpuTypes *models.GpuTypes
-	json.Unmarshal([]byte(setting.GpuTypes), &gpuTypes)
-	ctx.Data["gpu_types"] = gpuTypes.GpuType
+
+	if gpuInfos == nil {
+		json.Unmarshal([]byte(setting.GpuTypes), &gpuInfos)
+	}
+	ctx.Data["gpu_types"] = gpuInfos.GpuInfo
 	ctx.Data["snn4imagenet_path"] = cloudbrain.Snn4imagenetMountPath
 	ctx.Data["is_snn4imagenet_enabled"] = setting.IsSnn4imagenetEnabled
 	ctx.HTML(200, tplCloudBrainNew)
@@ -149,7 +157,7 @@ func CloudBrainCreate(ctx *context.Context, form auth.CreateCloudBrainForm) {
 	command := form.Command
 	uuid := form.Attachment
 	jobType := form.JobType
-	gpuType := setting.JobType
+	gpuQueue := setting.JobType
 	codePath := setting.JobPath + jobName + cloudbrain.CodeMountPath
 
 	if jobType != string(models.JobTypeBenchmark) && jobType != string(models.JobTypeDebug) && jobType != string(models.JobTypeSnn4imagenet) {
@@ -169,16 +177,22 @@ func CloudBrainCreate(ctx *context.Context, form auth.CreateCloudBrainForm) {
 
 	benchmarkPath := setting.JobPath + jobName + cloudbrain.BenchMarkMountPath
 	if setting.IsBenchmarkEnabled && jobType == string(models.JobTypeBenchmark) {
-		gpuType = form.GpuType
-		downloadRateCode(repo, jobName, setting.BenchmarkCode, benchmarkPath, form.BenchmarkCategory)
+		gpuQueue = form.GpuType
+		var gpuType string
+		for _, gpuInfo := range gpuInfos.GpuInfo {
+			if gpuInfo.Queue == gpuQueue {
+				gpuType = gpuInfo.Value
+			}
+		}
+		downloadRateCode(repo, jobName, setting.BenchmarkCode, benchmarkPath, form.BenchmarkCategory, gpuType)
 	}
 
 	snn4imagenetPath := setting.JobPath + jobName + cloudbrain.Snn4imagenetMountPath
 	if setting.IsSnn4imagenetEnabled && jobType == string(models.JobTypeSnn4imagenet) {
-		downloadRateCode(repo, jobName, setting.Snn4imagenetCode, snn4imagenetPath, "")
+		downloadRateCode(repo, jobName, setting.Snn4imagenetCode, snn4imagenetPath, "", "")
 	}
 
-	err = cloudbrain.GenerateTask(ctx, jobName, image, command, uuid, codePath, modelPath, benchmarkPath, snn4imagenetPath, jobType, gpuType)
+	err = cloudbrain.GenerateTask(ctx, jobName, image, command, uuid, codePath, modelPath, benchmarkPath, snn4imagenetPath, jobType, gpuQueue)
 	if err != nil {
 		ctx.RenderWithErr(err.Error(), tplCloudBrainNew, &form)
 		return
@@ -344,7 +358,7 @@ func downloadCode(repo *models.Repository, codePath string) error {
 	return nil
 }
 
-func downloadRateCode(repo *models.Repository, taskName, gitPath, codePath, benchmarkCategory string) error {
+func downloadRateCode(repo *models.Repository, taskName, gitPath, codePath, benchmarkCategory, gpuType string) error {
 	err := os.MkdirAll(codePath, os.ModePerm)
 	if err != nil {
 		log.Error("mkdir codePath failed", err.Error())
@@ -375,6 +389,7 @@ func downloadRateCode(repo *models.Repository, taskName, gitPath, codePath, benc
 		CodeName:          repo.Name,
 		BenchmarkCategory: strings.Split(benchmarkCategory, ","),
 		CodeLink:          strings.TrimSuffix(repo.CloneLink().HTTPS, ".git"),
+		GpuType:           gpuType,
 	})
 	if err != nil {
 		log.Error("json.Marshal failed", err.Error())