models/cloudbrain.go

@@ -591,6 +591,7 @@ type SpecialPool struct {
 	Pool         []*GpuInfo      `json:"pool"`
 	JobType      []string        `json:"jobType"`
 	ResourceSpec []*ResourceSpec `json:"resourceSpecs"`
+	Flavor       []*FlavorInfo   `json:"flavor"`
 }
 
 type ImageInfosModelArts struct {

models/org.go

@@ -470,6 +470,14 @@ func isOrganizationMember(e Engine, orgID, uid int64) (bool, error) {
 		Exist()
 }
 
+func IsOrganizationMemberByOrgName(orgName string, uid int64) (bool, error) {
+	org, _ := GetOrgByName(orgName)
+	if org != nil {
+		return IsOrganizationMember(org.ID, uid)
+	}
+	return false, nil
+}
+
 // IsPublicMembership returns true if given user public his/her membership.
 func IsPublicMembership(orgID, uid int64) (bool, error) {
 	return x.

modules/modelarts/modelarts.go

@@ -67,9 +67,11 @@ const (
 )
 
 var (
-	poolInfos   *models.PoolInfos
-	FlavorInfos *models.FlavorInfos
-	ImageInfos  *models.ImageInfosModelArts
+	poolInfos        *models.PoolInfos
+	FlavorInfos      *models.FlavorInfos
+	ImageInfos       *models.ImageInfosModelArts
+	TrainFlavorInfos *Flavor
+	SpecialPools     *models.SpecialPools
 )
 
 type GenerateTrainJobReq struct {
@@ -799,3 +801,9 @@ func GetNotebookImageName(imageId string) (string, error) {
 
 	return imageName, nil
 }
+
+func InitSpecialPool() {
+	if SpecialPools == nil && setting.ModelArtsSpecialPools != "" {
+		json.Unmarshal([]byte(setting.ModelArtsSpecialPools), &SpecialPools)
+	}
+}

modules/setting/setting.go

@@ -540,11 +540,12 @@ var (
 	ImageInfos        string
 	Capacity          int
 	//train-job
-	ResourcePools       string
-	Engines             string
-	EngineVersions      string
-	FlavorInfos         string
-	TrainJobFLAVORINFOS string
+	ResourcePools         string
+	Engines               string
+	EngineVersions        string
+	FlavorInfos           string
+	TrainJobFLAVORINFOS   string
+	ModelArtsSpecialPools string
 
 	//grampus config
 	Grampus = struct {
@@ -1418,6 +1419,7 @@ func NewContext() {
 	EngineVersions = sec.Key("Engine_Versions").MustString("")
 	FlavorInfos = sec.Key("FLAVOR_INFOS").MustString("")
 	TrainJobFLAVORINFOS = sec.Key("TrainJob_FLAVOR_INFOS").MustString("")
+	ModelArtsSpecialPools = sec.Key("SPECIAL_POOL").MustString("")
 
 	sec = Cfg.Section("elk")
 	ElkUrl = sec.Key("ELKURL").MustString("")

options/locale/locale_en-US.ini

@@ -3105,6 +3105,8 @@ select_dataset = select dataset
 specification = specification
 select_specification = select specification
 description = description
+wrong_specification=You cannot use this specification, please choose another item.
+
 
 job_name_rule = Please enter letters, numbers, _ and - up to 64 characters and cannot end with a dash (-).
 dataset_path_rule = The dataset location is stored in the environment variable data_url, and the training output path is stored in the environment variable train_url.

options/locale/locale_zh-CN.ini

@@ -3119,6 +3119,7 @@ select_dataset = 选择数据集
 specification = 规格
 select_specification = 选择资源规格
 description = 描述
+wrong_specification=您目前不能使用这个资源规格，请选择其他资源规格。
 
 job_name_rule = 请输入字母、数字、_和-，最长64个字符，且不能以中划线（-）结尾。
 dataset_path_rule = 数据集位置存储在环境变量data_url中，训练输出路径存储在环境变量train_url中。

routers/repo/cloudbrain.go

@@ -170,25 +170,22 @@ func cloudBrainNewDataPrepare(ctx *context.Context) error {
 		var trainGpuTypes []*models.GpuInfo
 
 		for _, pool := range cloudbrain.SpecialPools.Pools {
-			org, _ := models.GetOrgByName(pool.Org)
-			if org != nil {
-				isOrgMember, _ := models.IsOrganizationMember(org.ID, ctx.User.ID)
-				if isOrgMember {
-					for _, jobType := range pool.JobType {
-						if jobType == string(models.JobTypeDebug) {
-							debugGpuTypes = append(debugGpuTypes, pool.Pool...)
-							if pool.ResourceSpec != nil {
-								ctx.Data["resource_specs"] = pool.ResourceSpec
-							}
-						} else if jobType == string(models.JobTypeTrain) {
-							trainGpuTypes = append(trainGpuTypes, pool.Pool...)
-							if pool.ResourceSpec != nil {
-								ctx.Data["train_resource_specs"] = pool.ResourceSpec
-							}
+			isOrgMember, _ := models.IsOrganizationMemberByOrgName(pool.Org, ctx.User.ID)
+			if isOrgMember {
+				for _, jobType := range pool.JobType {
+					if jobType == string(models.JobTypeDebug) {
+						debugGpuTypes = append(debugGpuTypes, pool.Pool...)
+						if pool.ResourceSpec != nil {
+							ctx.Data["resource_specs"] = pool.ResourceSpec
+						}
+					} else if jobType == string(models.JobTypeTrain) {
+						trainGpuTypes = append(trainGpuTypes, pool.Pool...)
+						if pool.ResourceSpec != nil {
+							ctx.Data["train_resource_specs"] = pool.ResourceSpec
 						}
 					}
-					break
 				}
+				break
 			}
 
 		}

routers/repo/modelarts.go

@@ -146,6 +146,7 @@ func notebookNewDataPrepare(ctx *context.Context) error {
 		json.Unmarshal([]byte(setting.FlavorInfos), &modelarts.FlavorInfos)
 	}
 	ctx.Data["flavors"] = modelarts.FlavorInfos.FlavorInfo
+	setSpecBySpecialPoolConfig(ctx, string(models.JobTypeDebug))
 
 	ctx.Data["datasetType"] = models.TypeCloudBrainTwo
 
@@ -238,6 +239,13 @@ func Notebook2Create(ctx *context.Context, form auth.CreateModelArtsNotebookForm
 		}
 	}
 
+	errStr := checkModelArtsSpecialPool(ctx, flavor, string(models.JobTypeDebug))
+	if errStr != "" {
+		notebookNewDataPrepare(ctx)
+		ctx.RenderWithErr(ctx.Tr(errStr), tplModelArtsNotebookNew, &form)
+		return
+	}
+
 	err = modelarts.GenerateNotebook2(ctx, displayJobName, jobName, uuid, description, flavor, imageId)
 	if err != nil {
 		log.Error("GenerateNotebook2 failed, %v", err, ctx.Data["MsgID"])
@@ -326,15 +334,21 @@ func NotebookShow(ctx *context.Context) {
 	if modelarts.FlavorInfos == nil {
 		json.Unmarshal([]byte(setting.FlavorInfos), &modelarts.FlavorInfos)
 	}
+
+	findSpec := false
 	if modelarts.FlavorInfos != nil {
 		ctx.Data["resource_spec"] = modelarts.FlavorInfos.FlavorInfo[0].Desc
 		for _, f := range modelarts.FlavorInfos.FlavorInfo {
 			if fmt.Sprint(f.Value) == task.FlavorCode {
 				ctx.Data["resource_spec"] = f.Desc
+				findSpec = true
 				break
 			}
 		}
 	}
+
+	setShowSpecBySpecialPoolConfig(ctx, findSpec, task)
+
 	if task.TrainJobDuration == "" {
 		if task.Duration == 0 {
 			var duration int64
@@ -356,6 +370,20 @@ func NotebookShow(ctx *context.Context) {
 	ctx.HTML(200, tplModelArtsNotebookShow)
 }
 
+func setShowSpecBySpecialPoolConfig(ctx *context.Context, findSpec bool, task *models.Cloudbrain) {
+	modelarts.InitSpecialPool()
+	if modelarts.SpecialPools != nil && !findSpec {
+		for _, pool := range modelarts.SpecialPools.Pools {
+			for _, flavor := range pool.Flavor {
+				if flavor.Value == task.FlavorCode {
+					ctx.Data["resource_spec"] = flavor.Desc
+				}
+			}
+		}
+
+	}
+}
+
 func NotebookDebug(ctx *context.Context) {
 	var jobID = ctx.Params(":jobid")
 
@@ -690,6 +718,8 @@ func trainJobNewDataPrepare(ctx *context.Context) error {
 	}
 	ctx.Data["flavor_infos"] = flavorInfos.Info
 
+	setSpecBySpecialPoolConfig(ctx, string(models.JobTypeTrain))
+
 	ctx.Data["params"] = ""
 	ctx.Data["branchName"] = ctx.Repo.BranchName
 
@@ -704,6 +734,41 @@ func trainJobNewDataPrepare(ctx *context.Context) error {
 	return nil
 }
 
+func setSpecBySpecialPoolConfig(ctx *context.Context, jobType string) {
+	modelarts.InitSpecialPool()
+
+	if modelarts.SpecialPools != nil {
+		for _, specialPool := range modelarts.SpecialPools.Pools {
+			if cloudbrain.IsElementExist(specialPool.JobType, jobType) {
+
+				if isInOrg, _ := models.IsOrganizationMemberByOrgName(specialPool.Org, ctx.User.ID); isInOrg {
+					var specialFlavor []struct {
+						Code  string
+						Value string
+					}
+
+					if jobType == string(models.JobTypeDebug) {
+						ctx.Data["flavors"] = specialPool.Flavor
+					} else {
+
+						for _, tempFlavor := range specialPool.Flavor {
+							specialFlavor = append(specialFlavor, struct {
+								Code  string
+								Value string
+							}{Code: tempFlavor.Value, Value: tempFlavor.Desc})
+						}
+
+						ctx.Data["flavor_infos"] = specialFlavor
+					}
+
+				}
+
+			}
+		}
+
+	}
+}
+
 func trainJobErrorNewDataPrepare(ctx *context.Context, form auth.CreateModelArtsTrainJobForm) error {
 	ctx.Data["PageIsCloudBrain"] = true
 
@@ -757,6 +822,7 @@ func trainJobErrorNewDataPrepare(ctx *context.Context, form auth.CreateModelArts
 		return err
 	}
 	ctx.Data["flavor_infos"] = flavorInfos.Info
+	setSpecBySpecialPoolConfig(ctx, string(models.JobTypeTrain))
 
 	configList, err := getConfigList(modelarts.PerPage, 1, modelarts.SortByCreateTime, "desc", "", modelarts.ConfigTypeCustom)
 	if err != nil {
@@ -845,6 +911,8 @@ func trainJobNewVersionDataPrepare(ctx *context.Context) error {
 	}
 	ctx.Data["flavor_infos"] = flavorInfos.Info
 
+	setSpecBySpecialPoolConfig(ctx, string(models.JobTypeTrain))
+
 	var Parameters modelarts.Parameters
 	if err = json.Unmarshal([]byte(task.Parameters), &Parameters); err != nil {
 		ctx.ServerError("json.Unmarshal failed:", err)
@@ -931,6 +999,7 @@ func versionErrorDataPrepare(ctx *context.Context, form auth.CreateModelArtsTrai
 		return err
 	}
 	ctx.Data["flavor_infos"] = flavorInfos.Info
+	setSpecBySpecialPoolConfig(ctx, string(models.JobTypeTrain))
 
 	var Parameters modelarts.Parameters
 	if err = json.Unmarshal([]byte(form.Params), &Parameters); err != nil {
@@ -1018,6 +1087,13 @@ func TrainJobCreate(ctx *context.Context, form auth.CreateModelArtsTrainJobForm)
 		ctx.RenderWithErr(err.Error(), tplModelArtsTrainJobNew, &form)
 		return
 	}
+
+	errStr := checkModelArtsSpecialPool(ctx, flavorCode, string(models.JobTypeTrain))
+	if errStr != "" {
+		trainJobErrorNewDataPrepare(ctx, form)
+		ctx.RenderWithErr(ctx.Tr(errStr), tplModelArtsTrainJobNew, &form)
+		return
+	}
 	//Determine whether the task name of the task in the project is duplicated
 	tasks, err := models.GetCloudbrainsByDisplayJobName(repo.ID, string(models.JobTypeTrain), displayJobName)
 	if err == nil {
@@ -1297,6 +1373,7 @@ func TrainJobCreateVersion(ctx *context.Context, form auth.CreateModelArtsTrainJ
 
 	canNewJob, _ := canUserCreateTrainJobVersion(ctx, latestTask.UserID)
 	if !canNewJob {
+		versionErrorDataPrepare(ctx, form)
 		ctx.RenderWithErr("user cann't new trainjob", tplModelArtsTrainJobVersionNew, &form)
 		return
 	}
@@ -1308,6 +1385,13 @@ func TrainJobCreateVersion(ctx *context.Context, form auth.CreateModelArtsTrainJ
 		return
 	}
 
+	errStr := checkModelArtsSpecialPool(ctx, flavorCode, string(models.JobTypeTrain))
+	if errStr != "" {
+		versionErrorDataPrepare(ctx, form)
+		ctx.RenderWithErr(ctx.Tr(errStr), tplModelArtsTrainJobVersionNew, &form)
+		return
+	}
+
 	//todo: del the codeLocalPath
 	_, err = ioutil.ReadDir(codeLocalPath)
 	if err == nil {
@@ -1960,6 +2044,13 @@ func InferenceJobCreate(ctx *context.Context, form auth.CreateModelArtsInference
 		}
 	}
 
+	errStr := checkModelArtsSpecialPool(ctx, flavorCode, string(models.JobTypeInference))
+	if errStr != "" {
+		inferenceJobErrorNewDataPrepare(ctx, form)
+		ctx.RenderWithErr(ctx.Tr(errStr), tplModelArtsInferenceJobNew, &form)
+		return
+	}
+
 	//todo: del the codeLocalPath
 	_, err = ioutil.ReadDir(codeLocalPath)
 	if err == nil {
@@ -2076,6 +2167,65 @@ func InferenceJobCreate(ctx *context.Context, form auth.CreateModelArtsInference
 	}
 	ctx.Redirect(setting.AppSubURL + ctx.Repo.RepoLink + "/modelarts/inference-job")
 }
+
+func checkModelArtsSpecialPool(ctx *context.Context, flavorCode string, jobType string) string {
+	if modelarts.SpecialPools != nil {
+		isMatchPool := false
+
+		for _, specialPool := range modelarts.SpecialPools.Pools {
+			if cloudbrain.IsElementExist(specialPool.JobType, jobType) {
+				if isInOrg, _ := models.IsOrganizationMemberByOrgName(specialPool.Org, ctx.User.ID); isInOrg {
+					isMatchPool = true
+					isMatchSpec := false
+					for _, flavor := range specialPool.Flavor {
+						if flavor.Value == flavorCode {
+							isMatchSpec = true
+							break
+						}
+					}
+
+					if !isMatchSpec {
+						return "cloudbrain.wrong_specification"
+
+					}
+
+				}
+
+			}
+		}
+
+		if !isMatchPool {
+			isMatchSpec := false
+			if jobType == string(models.JobTypeDebug) {
+				for _, flavor := range modelarts.FlavorInfos.FlavorInfo {
+					if flavor.Value == flavorCode {
+						isMatchSpec = true
+						break
+					}
+				}
+			} else {
+
+				var flavorInfos modelarts.Flavor
+				json.Unmarshal([]byte(setting.TrainJobFLAVORINFOS), &flavorInfos)
+
+				for _, flavor := range flavorInfos.Info {
+					if flavor.Code == flavorCode {
+						isMatchSpec = true
+						break
+					}
+				}
+			}
+
+			if !isMatchSpec {
+
+				return "cloudbrain.wrong_specification"
+			}
+
+		}
+
+	}
+	return ""
+}
 func InferenceJobIndex(ctx *context.Context) {
 	MustEnableModelArts(ctx)
 
@@ -2193,6 +2343,8 @@ func inferenceJobNewDataPrepare(ctx *context.Context) error {
 	}
 
 	ctx.Data["flavor_infos"] = flavorInfos.Info
+	setSpecBySpecialPoolConfig(ctx, string(models.JobTypeInference))
+
 	ctx.Data["params"] = ""
 	ctx.Data["branchName"] = ctx.Repo.BranchName
 
@@ -2261,6 +2413,7 @@ func inferenceJobErrorNewDataPrepare(ctx *context.Context, form auth.CreateModel
 		return err
 	}
 	ctx.Data["flavor_infos"] = flavorInfos.Info
+	setSpecBySpecialPoolConfig(ctx, string(models.JobTypeInference))
 
 	configList, err := getConfigList(modelarts.PerPage, 1, modelarts.SortByCreateTime, "desc", "", modelarts.ConfigTypeCustom)
 	if err != nil {