create

modules/cloudbrain/cloudbrain.go

@@ -32,6 +32,8 @@ const (
 	SubTaskName = "task1"
 
 	Success = "S000"
+
+	DefaultBranchName = "master"
 )
 
 var (

routers/repo/cloudbrain.go

@@ -183,15 +183,24 @@ func CloudBrainCreate(ctx *context.Context, form auth.CreateCloudBrainForm) {
 	image := form.Image
 	uuid := form.Attachment
 	jobType := form.JobType
-	command := cloudbrain.Command
 	gpuQueue := form.GpuType
 	codePath := setting.JobPath + jobName + cloudbrain.CodeMountPath
 	resourceSpecId := form.ResourceSpecId
+	branchName := form.BranchName
 	repo := ctx.Repo.Repository
 
 	tpl := tplCloudBrainNew
+	command := cloudbrain.Command
 	if jobType == string(models.JobTypeTrain) {
 		tpl = tplCloudBrainTrainJobNew
+		command, err := getTrainJobCommand(form)
+		if err != nil {
+			log.Error("getTrainJobCommand failed: %v", err)
+			ctx.RenderWithErr(err.Error(), tpl, &form)
+			return
+		}
+
+		log.Info("%s", command)
 	}
 
 	tasks, err := models.GetCloudbrainsByDisplayJobName(repo.ID, jobType, displayJobName)
@@ -216,7 +225,7 @@ func CloudBrainCreate(ctx *context.Context, form auth.CreateCloudBrainForm) {
 		return
 	}
 
-	if jobType != string(models.JobTypeBenchmark) && jobType != string(models.JobTypeDebug) && jobType != string(models.JobTypeSnn4imagenet) && jobType != string(models.JobTypeBrainScore) {
+	if jobType != string(models.JobTypeBenchmark) && jobType != string(models.JobTypeDebug) && jobType != string(models.JobTypeSnn4imagenet) && jobType != string(models.JobTypeBrainScore) && jobType != string(models.JobTypeTrain) {
 		log.Error("jobtype error:", jobType, ctx.Data["MsgID"])
 		cloudBrainNewDataPrepare(ctx)
 		ctx.RenderWithErr("jobtype error", tpl, &form)
@@ -238,7 +247,10 @@ func CloudBrainCreate(ctx *context.Context, form auth.CreateCloudBrainForm) {
 		}
 	}
 
-	downloadCode(repo, codePath)
+	if branchName == "" {
+		branchName = cloudbrain.DefaultBranchName
+	}
+	downloadCode(repo, codePath, branchName)
 	uploadCodeToMinio(codePath+"/", jobName, cloudbrain.CodeMountPath+"/")
 
 	modelPath := setting.JobPath + jobName + cloudbrain.ModelMountPath + "/"
@@ -764,8 +776,8 @@ func GetRate(ctx *context.Context) {
 	}
 }
 
-func downloadCode(repo *models.Repository, codePath string) error {
-	if err := git.Clone(repo.RepoPath(), codePath, git.CloneRepoOptions{}); err != nil {
+func downloadCode(repo *models.Repository, codePath, branchName string) error {
+	if err := git.Clone(repo.RepoPath(), codePath, git.CloneRepoOptions{Branch: branchName}); err != nil {
 		log.Error("Failed to clone repository: %s (%v)", repo.FullName(), err)
 		return err
 	}
@@ -1316,7 +1328,7 @@ func CloudBrainBenchmarkCreate(ctx *context.Context, form auth.CreateCloudBrainF
 	}
 
 	os.RemoveAll(codePath)
-	if err := downloadCode(repo, codePath); err != nil {
+	if err := downloadCode(repo, codePath, cloudbrain.DefaultBranchName); err != nil {
 		log.Error("downloadCode failed, %v", err, ctx.Data["MsgID"])
 		cloudBrainNewDataPrepare(ctx)
 		ctx.RenderWithErr("system error", tplCloudBrainBenchmarkNew, &form)
@@ -1415,3 +1427,15 @@ func CloudBrainTrainJobNew(ctx *context.Context) {
 	}
 	ctx.HTML(http.StatusOK, tplCloudBrainTrainJobNew)
 }
+
+func getTrainJobCommand(form auth.CreateCloudBrainForm) (string, error) {
+	var command string
+	bootFile := form.BootFile
+
+	if !strings.HasSuffix(bootFile, ".py") {
+		log.Error("bootFile(%s) format error", bootFile)
+		return command, errors.New("bootFile format error")
+	}
+
+	return command, nil
+}