From d1c8e345e5b2de6f1f0e2aec970fd77952e5f507 Mon Sep 17 00:00:00 2001
From: lewis <747342561@qq.com>
Date: Tue, 28 Jun 2022 17:50:41 +0800
Subject: [PATCH] gpu add param

---
 routers/repo/cloudbrain.go | 30 ++++++++++++++++--------------
 1 file changed, 16 insertions(+), 14 deletions(-)

diff --git a/routers/repo/cloudbrain.go b/routers/repo/cloudbrain.go
index 812d96cdb..c08e864c6 100755
--- a/routers/repo/cloudbrain.go
+++ b/routers/repo/cloudbrain.go
@@ -215,20 +215,7 @@ func CloudBrainCreate(ctx *context.Context, form auth.CreateCloudBrainForm) {
 	resourceSpecId := form.ResourceSpecId
 	branchName := form.BranchName
 	repo := ctx.Repo.Repository
-
 	tpl := tplCloudBrainNew
-	command := cloudbrain.Command
-	if jobType == string(models.JobTypeTrain) {
-		tpl = tplCloudBrainTrainJobNew
-		commandTrain, err := getTrainJobCommand(form)
-		if err != nil {
-			log.Error("getTrainJobCommand failed: %v", err)
-			ctx.RenderWithErr(err.Error(), tpl, &form)
-			return
-		}
-
-		command = commandTrain
-	}
 
 	tasks, err := models.GetCloudbrainsByDisplayJobName(repo.ID, jobType, displayJobName)
 	if err == nil {
@@ -282,6 +269,19 @@ func CloudBrainCreate(ctx *context.Context, form auth.CreateCloudBrainForm) {
 		return
 	}
 
+	command := cloudbrain.Command
+	if jobType == string(models.JobTypeTrain) {
+		tpl = tplCloudBrainTrainJobNew
+		commandTrain, err := getTrainJobCommand(form, datasetNames)
+		if err != nil {
+			log.Error("getTrainJobCommand failed: %v", err)
+			ctx.RenderWithErr(err.Error(), tpl, &form)
+			return
+		}
+
+		command = commandTrain
+	}
+
 	if branchName == "" {
 		branchName = cloudbrain.DefaultBranchName
 	}
@@ -2246,7 +2246,7 @@ func CloudBrainTrainJobNew(ctx *context.Context) {
 	ctx.HTML(http.StatusOK, tplCloudBrainTrainJobNew)
 }
 
-func getTrainJobCommand(form auth.CreateCloudBrainForm) (string, error) {
+func getTrainJobCommand(form auth.CreateCloudBrainForm, datasetNames string) (string, error) {
 	var command string
 	bootFile := strings.TrimSpace(form.BootFile)
 	params := form.Params
@@ -2270,6 +2270,8 @@ func getTrainJobCommand(form auth.CreateCloudBrainForm) (string, error) {
 		}
 	}
 
+	param += " --dataset_list=" + datasetNames
+
 	command += "python /code/" + bootFile + param + " > " + cloudbrain.ModelMountPath + "/" + form.DisplayJobName + "-" + cloudbrain.LogFile
 
 	return command, nil