From 311dbe09273f3cdb2acdcb3748772f5094c4cd15 Mon Sep 17 00:00:00 2001
From: lewis <747342561@qq.com>
Date: Mon, 27 Jun 2022 20:17:56 +0800
Subject: [PATCH] gpu multi dataset

---
 modules/cloudbrain/cloudbrain.go | 103 +++++++++++++++++--------------
 routers/repo/cloudbrain.go       |  14 ++---
 2 files changed, 65 insertions(+), 52 deletions(-)

diff --git a/modules/cloudbrain/cloudbrain.go b/modules/cloudbrain/cloudbrain.go
index d79b6ace6..0ac7b9937 100755
--- a/modules/cloudbrain/cloudbrain.go
+++ b/modules/cloudbrain/cloudbrain.go
@@ -249,7 +249,63 @@ func GenerateTask(req GenerateCloudBrainTaskReq) error {
 		return errors.New("no such resourceSpec")
 	}
 
-	log.Info(req.DatasetInfos[req.Uuids].DataLocalPath)
+	volumes := []models.Volume{
+		{
+			HostPath: models.StHostPath{
+				Path:      req.CodePath,
+				MountPath: CodeMountPath,
+				ReadOnly:  false,
+			},
+		},
+		{
+			HostPath: models.StHostPath{
+				Path:      req.ModelPath,
+				MountPath: ModelMountPath,
+				ReadOnly:  false,
+			},
+		},
+		{
+			HostPath: models.StHostPath{
+				Path:      req.BenchmarkPath,
+				MountPath: BenchMarkMountPath,
+				ReadOnly:  true,
+			},
+		},
+		{
+			HostPath: models.StHostPath{
+				Path:      req.Snn4ImageNetPath,
+				MountPath: Snn4imagenetMountPath,
+				ReadOnly:  true,
+			},
+		},
+		{
+			HostPath: models.StHostPath{
+				Path:      req.BrainScorePath,
+				MountPath: BrainScoreMountPath,
+				ReadOnly:  true,
+			},
+		},
+	}
+
+	if len(req.DatasetInfos) == 1 {
+		volumes = append(volumes, models.Volume{
+			HostPath: models.StHostPath{
+				Path:      req.DatasetInfos[req.Uuids].DataLocalPath,
+				MountPath: DataSetMountPath,
+				ReadOnly:  true,
+			},
+		})
+	} else {
+		for _, dataset := range req.DatasetInfos {
+			volumes = append(volumes, models.Volume{
+				HostPath: models.StHostPath{
+					Path:      dataset.DataLocalPath,
+					MountPath: DataSetMountPath + "/" + dataset.Name,
+					ReadOnly:  true,
+				},
+			})
+		}
+	}
 
 	createTime := timeutil.TimeStampNow()
 	jobResult, err := CreateJob(req.JobName, models.CreateJobParams{
@@ -273,50 +329,7 @@ func GenerateTask(req GenerateCloudBrainTaskReq) error {
 				UseNNI:                false,
 			},
 		},
-		Volumes: []models.Volume{
-			{
-				HostPath: models.StHostPath{
-					Path:      req.CodePath,
-					MountPath: CodeMountPath,
-					ReadOnly:  false,
-				},
-			},
-			{
-				HostPath: models.StHostPath{
-					Path:      req.DatasetInfos[req.Uuids].DataLocalPath,
-					MountPath: DataSetMountPath,
-					ReadOnly:  true,
-				},
-			},
-			{
-				HostPath: models.StHostPath{
-					Path:      req.ModelPath,
-					MountPath: ModelMountPath,
-					ReadOnly:  false,
-				},
-			},
-			{
-				HostPath: models.StHostPath{
-					Path:      req.BenchmarkPath,
-					MountPath: BenchMarkMountPath,
-					ReadOnly:  true,
-				},
-			},
-			{
-				HostPath: models.StHostPath{
-					Path:      req.Snn4ImageNetPath,
-					MountPath: Snn4imagenetMountPath,
-					ReadOnly:  true,
-				},
-			},
-			{
-				HostPath: models.StHostPath{
-					Path:      req.BrainScorePath,
-					MountPath: BrainScoreMountPath,
-					ReadOnly:  true,
-				},
-			},
-		},
+		Volumes: volumes,
 	})
 	if err != nil {
 		log.Error("CreateJob failed:", err.Error(), req.Ctx.Data["MsgID"])
diff --git a/routers/repo/cloudbrain.go b/routers/repo/cloudbrain.go
index 90c2146d1..a29c390ce 100755
--- a/routers/repo/cloudbrain.go
+++ b/routers/repo/cloudbrain.go
@@ -2290,7 +2290,7 @@ func GetBenchmarkTypes(ctx *context.Context) *models.BenchmarkTypes {
 func getDatasetInfo(uuidStr string) (map[string]cloudbrain.DatasetInfo, string, error) {
 	var datasetNames string
 	uuids := strings.Split(uuidStr, ";")
-	if len(uuids) > 5 {
+	if len(uuids) > setting.MaxDatasetNum {
 		log.Error("the dataset count(%d) exceed the limit", len(uuids))
 		return nil, datasetNames, errors.New("the dataset count exceed the limit")
 	}
@@ -2303,9 +2303,11 @@ func getDatasetInfo(uuidStr string) (map[string]cloudbrain.DatasetInfo, string,
 			return nil, datasetNames, err
 		}
 
-		if _, ok := datasetInfos[uuid]; ok {
-			log.Error("the dataset name is same: %v", attach.Name)
-			return nil, datasetNames, errors.New("the dataset name is same")
+		for _, datasetInfo := range datasetInfos {
+			if attach.Name == datasetInfo.Name {
+				log.Error("the dataset name is same: %v", attach.Name)
+				return nil, datasetNames, errors.New("the dataset name is same")
+			}
 		}
 
 		dataLocalPath := setting.Attachment.Minio.RealPath +
@@ -2316,7 +2318,7 @@ func getDatasetInfo(uuidStr string) (map[string]cloudbrain.DatasetInfo, string,
 
 		datasetInfos[uuid] = cloudbrain.DatasetInfo{
 			DataLocalPath: dataLocalPath,
-			Name:          attach.Name,
+			Name:          strings.TrimSuffix(strings.TrimSuffix(strings.TrimSuffix(attach.Name, ".zip"), ".tar.gz"), ".tgz"),
 		}
 		if i == 0 {
 			datasetNames = attach.Name
@@ -2325,7 +2327,5 @@ func getDatasetInfo(uuidStr string) (map[string]cloudbrain.DatasetInfo, string,
 		}
 	}
 
-	log.Info(datasetNames)
-
 	return datasetInfos, datasetNames, nil
 }