From 546226167e7db67476a1de5a3ec4e9a5a5877a27 Mon Sep 17 00:00:00 2001
From: lewis <747342561@qq.com>
Date: Tue, 7 Dec 2021 14:50:22 +0800
Subject: [PATCH] download model

---
 routers/repo/cloudbrain.go | 46 ++++++++++++++++++++++++++++++--------
 1 file changed, 37 insertions(+), 9 deletions(-)

diff --git a/routers/repo/cloudbrain.go b/routers/repo/cloudbrain.go
index a3af588fb..0689741d0 100755
--- a/routers/repo/cloudbrain.go
+++ b/routers/repo/cloudbrain.go
@@ -251,13 +251,9 @@ func CloudBrainCreate(ctx *context.Context, form auth.CreateCloudBrainForm) {
 	downloadCode(repo, codePath)
 	uploadCodeToMinio(codePath + "/", jobName, "/code/")
 
-	modelPath := setting.JobPath + jobName + cloudbrain.ModelMountPath
-	err = os.MkdirAll(modelPath, os.ModePerm)
-	if err != nil {
-		cloudBrainNewDataPrepare(ctx)
-		ctx.RenderWithErr(err.Error(), tplCloudBrainNew, &form)
-		return
-	}
+	modelPath := setting.JobPath + jobName + cloudbrain.ModelMountPath + "/"
+	mkModelPath(modelPath)
+	uploadCodeToMinio(modelPath, jobName, cloudbrain.ModelMountPath + "/")
 
 	benchmarkPath := setting.JobPath + jobName + cloudbrain.BenchMarkMountPath
 	if setting.IsBenchmarkEnabled && jobType == string(models.JobTypeBenchmark) {
@@ -283,7 +279,9 @@ func CloudBrainCreate(ctx *context.Context, form auth.CreateCloudBrainForm) {
 		uploadCodeToMinio(brainScorePath + "/", jobName, cloudbrain.BrainScoreMountPath + "/")
 	}
 
-	err = cloudbrain.GenerateTask(ctx, jobName, image, command, uuid, codePath, modelPath, benchmarkPath, snn4imagenetPath, brainScorePath, jobType, gpuQueue, resourceSpecId)
+	err = cloudbrain.GenerateTask(ctx, jobName, image, command, uuid, codePath, getMinioPath(jobName, cloudbrain.ModelMountPath + "/"),
+		getMinioPath(jobName, cloudbrain.BenchMarkMountPath + "/"), getMinioPath(jobName, cloudbrain.Snn4imagenetMountPath + "/"),
+		getMinioPath(jobName, cloudbrain.BrainScoreMountPath + "/"), jobType, gpuQueue, resourceSpecId)
 	if err != nil {
 		cloudBrainNewDataPrepare(ctx)
 		ctx.RenderWithErr(err.Error(), tplCloudBrainNew, &form)
@@ -607,7 +605,7 @@ func getImages(ctx *context.Context, imageType string) {
 
 func getModelDirs(jobName string, parentDir string) (string, error) {
 	var req string
-	modelActualPath := setting.JobPath + jobName + "/model/"
+	modelActualPath := getMinioPath(jobName, cloudbrain.ModelMountPath + "/")
 	if parentDir == "" {
 		req = "baseDir=" + modelActualPath
 	} else {
@@ -617,6 +615,10 @@ func getModelDirs(jobName string, parentDir string) (string, error) {
 	return getDirs(req)
 }
 
+func getMinioPath(jobName, suffixPath string) string {
+	return setting.Attachment.Minio.RealPath + setting.Attachment.Minio.Bucket + "/" + setting.CBCodePathPrefix + jobName + suffixPath
+}
+
 func CloudBrainDownloadModel(ctx *context.Context) {
 	parentDir := ctx.Query("parentDir")
 	fileName := ctx.Query("fileName")
@@ -772,6 +774,32 @@ func uploadCodeToMinio(codePath, jobName, parentDir string) error {
 	return nil
 }
 
+func mkModelPath(modelPath string) error {
+	err := os.MkdirAll(modelPath, os.ModePerm)
+	if err != nil {
+		log.Error("MkdirAll(%s) failed:%v", modelPath, err)
+		return err
+	}
+
+	fileName := modelPath + "README"
+	f, err := os.OpenFile(fileName, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, os.ModePerm)
+	if err != nil {
+		log.Error("OpenFile failed", err.Error())
+		return err
+	}
+
+	defer f.Close()
+
+	_, err = f.WriteString("You can put the model file into this directory and download it by the web page.")
+	if err != nil {
+		log.Error("WriteString failed", err.Error())
+		return err
+	}
+
+	return nil
+}
+
+
 func SyncCloudbrainStatus() {
 	cloudBrains, err := models.GetCloudBrainUnStoppedJob()
 	if err != nil {