提交代码

models/cloudbrain.go

@@ -570,11 +570,12 @@ type SpecialPools struct {
 	Pools []*SpecialPool `json:"pools"`
 }
 type SpecialPool struct {
-	Org         string     `json:"org"`
-	Type        string     `json:"type"`
-	IsExclusive bool       `json:"isExclusive"`
-	Pool        []*GpuInfo `json:"pool"`
-	JobType     []string   `json:"jobType"`
+	Org          string          `json:"org"`
+	Type         string          `json:"type"`
+	IsExclusive  bool            `json:"isExclusive"`
+	Pool         []*GpuInfo      `json:"pool"`
+	JobType      []string        `json:"support_type"`
+	ResourceSpec []*ResourceSpec `json:"resource_specs"`
 }
 
 type ImageInfosModelArts struct {

modules/cloudbrain/cloudbrain.go

@@ -42,6 +42,7 @@ const (
 var (
 	ResourceSpecs      *models.ResourceSpecs
 	TrainResourceSpecs *models.ResourceSpecs
+	SpecialPools       *models.SpecialPools
 )
 
 type GenerateCloudBrainTaskReq struct {
@@ -538,3 +539,9 @@ func RestartTask(ctx *context.Context, task *models.Cloudbrain, newID *string) e
 
 	return nil
 }
+
+func InitSpecialPool() {
+	if SpecialPools == nil && setting.SpecialPools != "" {
+		json.Unmarshal([]byte(setting.SpecialPools), &SpecialPools)
+	}
+}

modules/setting/setting.go

@@ -460,6 +460,7 @@ var (
 	CBCodePathPrefix   string
 	JobType            string
 	GpuTypes           string
+	SpecialPools       string
 	DebugServerHost    string
 	ResourceSpecs      string
 	MaxDuration        int64
@@ -1311,6 +1312,8 @@ func NewContext() {
 	MaxDuration = sec.Key("MAX_DURATION").MustInt64(14400)
 	TrainGpuTypes = sec.Key("TRAIN_GPU_TYPES").MustString("")
 	TrainResourceSpecs = sec.Key("TRAIN_RESOURCE_SPECS").MustString("")
+	SpecialPools = sec.Key("SPECIAL_POOL").MustString("")
+	
 	MaxDatasetNum = sec.Key("MAX_DATASET_NUM").MustInt(5)
 
 	sec = Cfg.Section("benchmark")

routers/repo/cloudbrain.go

@@ -149,6 +149,8 @@ func cloudBrainNewDataPrepare(ctx *context.Context) error {
 
 	ctx.Data["benchmark_types"] = GetBenchmarkTypes(ctx).BenchmarkType
 
+	cloudbrain.InitSpecialPool()
+
 	if gpuInfos == nil {
 		json.Unmarshal([]byte(setting.GpuTypes), &gpuInfos)
 	}
@@ -178,6 +180,44 @@ func cloudBrainNewDataPrepare(ctx *context.Context) error {
 		json.Unmarshal([]byte(setting.TrainResourceSpecs), &cloudbrain.TrainResourceSpecs)
 	}
 	ctx.Data["train_resource_specs"] = cloudbrain.TrainResourceSpecs.ResourceSpec
+
+	if cloudbrain.SpecialPools != nil {
+		var debugGpuTypes []*models.GpuInfo
+		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
+							}
+						}
+					}
+				}
+			}
+
+		}
+
+		if len(debugGpuTypes) > 0 {
+			ctx.Data["gpu_types"] = debugGpuTypes
+		}
+
+		if len(trainGpuTypes) > 0 {
+			ctx.Data["train_gpu_types"] = trainGpuInfos
+		}
+
+	}
+
 	ctx.Data["params"] = ""
 	ctx.Data["branchName"] = ctx.Repo.BranchName