From bd6daf45590760e9ef616418740d304cf390b8f7 Mon Sep 17 00:00:00 2001
From: chenyifan01 <cyf_uestc@outlook.com>
Date: Tue, 13 Dec 2022 15:20:45 +0800
Subject: [PATCH] #3277 gcu:fix bug of npu dataset error

---
 models/cloudbrain.go       |  2 ++
 modules/grampus/grampus.go | 35 ++++++++++++++++++++++++++++++++++-
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/models/cloudbrain.go b/models/cloudbrain.go
index 4f8fb4c42..851ab2456 100755
--- a/models/cloudbrain.go
+++ b/models/cloudbrain.go
@@ -2579,6 +2579,7 @@ type DatasetInfo struct {
 	DataLocalPath string
 	Name          string
 	FullName      string
+	Type          int
 }
 
 func GetDatasetInfo(uuidStr string, grampusType ...string) (map[string]DatasetInfo, string, error) {
@@ -2640,6 +2641,7 @@ func GetDatasetInfo(uuidStr string, grampusType ...string) (map[string]DatasetIn
 			DataLocalPath: dataLocalPath,
 			Name:          fileName,
 			FullName:      attach.Name,
+			Type:          attach.Type,
 		}
 		if i == 0 {
 			datasetNames = attach.Name
diff --git a/modules/grampus/grampus.go b/modules/grampus/grampus.go
index f633674a3..11d4b253c 100755
--- a/modules/grampus/grampus.go
+++ b/modules/grampus/grampus.go
@@ -149,6 +149,35 @@ func getDatasetGPUGrampus(datasetInfos map[string]models.DatasetInfo) ([]models.
 	}
 	return datasetGrampus, command
 }
+func getDatasetGCUGrampus(datasetInfos map[string]models.DatasetInfo) ([]models.GrampusDataset, string) {
+	var datasetGrampus []models.GrampusDataset
+	var command = ""
+	epuEndPoint := getEndPoint()
+	for uuid, datasetInfo := range datasetInfos {
+		if datasetInfo.Type == models.TypeCloudBrainOne {
+			datasetGrampus = append(datasetGrampus, models.GrampusDataset{
+				Name:          datasetInfo.FullName,
+				Bucket:        setting.Attachment.Minio.Bucket,
+				EndPoint:      setting.Attachment.Minio.Endpoint,
+				ObjectKey:     datasetInfo.DataLocalPath,
+				ReadOnly:      true,
+				ContainerPath: "/dataset1/" + datasetInfo.Name,
+			})
+
+			command += "cp /dataset1/'" + datasetInfo.Name + "'/" + uuid + " /dataset/'" + datasetInfo.FullName + "';"
+
+		} else {
+			datasetGrampus = append(datasetGrampus, models.GrampusDataset{
+				Name:      datasetInfo.FullName,
+				Bucket:    setting.Bucket,
+				EndPoint:  epuEndPoint,
+				ObjectKey: datasetInfo.DataLocalPath + datasetInfo.FullName,
+			})
+		}
+
+	}
+	return datasetGrampus, command
+}
 
 func GenerateNotebookJob(ctx *context.Context, req *GenerateNotebookJobReq) (jobId string, err error) {
 	createTime := timeutil.TimeStampNow()
@@ -179,7 +208,11 @@ func GenerateNotebookJob(ctx *context.Context, req *GenerateNotebookJobReq) (job
 		imageUrl = ""
 		req.Command = ""
 	} else {
-		datasetGrampus, cpCommand = getDatasetGPUGrampus(req.DatasetInfos)
+		if ProcessorTypeGCU == req.ProcessType {
+			datasetGrampus, cpCommand = getDatasetGCUGrampus(req.DatasetInfos)
+		} else {
+			datasetGrampus, cpCommand = getDatasetGPUGrampus(req.DatasetInfos)
+		}
 		if len(req.ModelName) != 0 {
 			datasetGrampus = append(datasetGrampus, models.GrampusDataset{
 				Name:          req.ModelName,