提交代码

modules/cloudbrain/cloudbrain.go

@@ -28,6 +28,8 @@ const (
 	Snn4imagenetMountPath = "/snn4imagenet"
 	BrainScoreMountPath   = "/brainscore"
 	TaskInfoName          = "/taskInfo"
+	Snn4imagenetCommand   = `/opt/conda/bin/python /snn4imagenet/testSNN_script.py --modelname '%s' --modelpath '/dataset' --modeldescription '%s'`
+	BrainScoreCommand     = `/opt/conda/bin/python /brainscore/score_a_model.py --brainRegion '%s' --modelname '%s' --modelpath '/dataset' --modeldescription '%s'`
 
 	SubTaskName = "task1"
 

routers/repo/cloudbrain.go

@@ -2066,19 +2066,23 @@ func ModelBenchmarkCreate(ctx *context.Context, form auth.CreateCloudBrainForm)
 	if setting.IsSnn4imagenetEnabled && jobType == string(models.JobTypeSnn4imagenet) {
 		downloadRateCode(repo, jobName, setting.Snn4imagenetOwner, setting.Snn4imagenetName, snn4imagenetPath, "", "")
 		uploadCodeToMinio(snn4imagenetPath+"/", jobName, cloudbrain.Snn4imagenetMountPath+"/")
-	}
+		command = fmt.Sprintf(cloudbrain.Snn4imagenetCommand, displayJobName, form.Description)
 
+	}
+	benchmarkChildTypeID := 0
 	brainScorePath := setting.JobPath + jobName + cloudbrain.BrainScoreMountPath
 	if setting.IsBrainScoreEnabled && jobType == string(models.JobTypeBrainScore) {
 		downloadRateCode(repo, jobName, setting.BrainScoreOwner, setting.BrainScoreName, brainScorePath, "", "")
 		uploadCodeToMinio(brainScorePath+"/", jobName, cloudbrain.BrainScoreMountPath+"/")
+		benchmarkChildTypeID = form.BenchmarkChildTypeID
+		command = fmt.Sprintf(cloudbrain.BrainScoreCommand, getBrainRegion(benchmarkChildTypeID), displayJobName, form.Description)
 	}
 
 	err = cloudbrain.GenerateTask(ctx, displayJobName, jobName, image, command, uuid, storage.GetMinioPath(jobName, cloudbrain.CodeMountPath+"/"),
 		storage.GetMinioPath(jobName, cloudbrain.ModelMountPath+"/"),
 		storage.GetMinioPath(jobName, cloudbrain.BenchMarkMountPath+"/"), storage.GetMinioPath(jobName, cloudbrain.Snn4imagenetMountPath+"/"),
 		storage.GetMinioPath(jobName, cloudbrain.BrainScoreMountPath+"/"), jobType, gpuQueue, form.Description, branchName, form.BootFile, form.Params,
-		0, 0, resourceSpecId)
+		0, benchmarkChildTypeID, resourceSpecId)
 	if err != nil {
 		cloudBrainNewDataPrepare(ctx)
 		ctx.RenderWithErr(err.Error(), tpl, &form)
@@ -2088,6 +2092,11 @@ func ModelBenchmarkCreate(ctx *context.Context, form auth.CreateCloudBrainForm)
 	ctx.Redirect(setting.AppSubURL + ctx.Repo.RepoLink + "/cloudbrain/benchmark")
 }
 
+func getBrainRegion(benchmarkChildTypeID int) string {
+	values := []string{"V1", "V2", "V4", "IT"}
+	return values[benchmarkChildTypeID]
+}
+
 func BenchmarkDel(ctx *context.Context) {
 	if err := deleteCloudbrainJob(ctx); err != nil {
 		log.Error("deleteCloudbrainJob failed: %v", err, ctx.Data["msgID"])