update

models/cloudbrain.go

@@ -2014,3 +2014,110 @@ func GetDatasetInfo(uuidStr string) (map[string]DatasetInfo, string, error) {
 
 	return datasetInfos, datasetNames, nil
 }
+
+var (
+	SpecsMapInitFlag                    = false
+	CloudbrainDebugResourceSpecsMap     map[int]*ResourceSpec
+	CloudbrainTrainResourceSpecsMap     map[int]*ResourceSpec
+	CloudbrainBenchmarkResourceSpecsMap map[int]*ResourceSpec
+	ModelArtsDebugResourceSpecsMap      map[string]*FlavorInfo
+	ModelArtsTrainResourceSpecsMap      map[string]*FlavorInfo
+)
+
+type ModelArtsFlavor struct {
+	Info []struct {
+		Code      string `json:"code"`
+		Value     string `json:"value"`
+		UnitPrice int64  `json:"unitPrice"`
+	} `json:"flavor"`
+}
+
+func InitResourceSpecMap() {
+	if CloudbrainDebugResourceSpecsMap == nil || len(CloudbrainDebugResourceSpecsMap) == 0 {
+		t := ResourceSpecs{}
+		json.Unmarshal([]byte(setting.ResourceSpecs), &t)
+		CloudbrainDebugResourceSpecsMap = make(map[int]*ResourceSpec, len(t.ResourceSpec))
+		for _, spec := range t.ResourceSpec {
+			CloudbrainDebugResourceSpecsMap[spec.Id] = spec
+		}
+	}
+	if CloudbrainTrainResourceSpecsMap == nil || len(CloudbrainTrainResourceSpecsMap) == 0 {
+		t := ResourceSpecs{}
+		json.Unmarshal([]byte(setting.TrainResourceSpecs), &t)
+		CloudbrainTrainResourceSpecsMap = make(map[int]*ResourceSpec, len(t.ResourceSpec))
+		for _, spec := range t.ResourceSpec {
+			CloudbrainTrainResourceSpecsMap[spec.Id] = spec
+		}
+	}
+	if CloudbrainBenchmarkResourceSpecsMap == nil || len(CloudbrainBenchmarkResourceSpecsMap) == 0 {
+		t := ResourceSpecs{}
+		json.Unmarshal([]byte(setting.BenchmarkResourceSpecs), &t)
+		CloudbrainBenchmarkResourceSpecsMap = make(map[int]*ResourceSpec, len(t.ResourceSpec))
+		for _, spec := range t.ResourceSpec {
+			CloudbrainBenchmarkResourceSpecsMap[spec.Id] = spec
+		}
+	}
+	if ModelArtsDebugResourceSpecsMap == nil || len(ModelArtsDebugResourceSpecsMap) == 0 {
+		t := FlavorInfos{}
+		json.Unmarshal([]byte(setting.FlavorInfos), &t)
+		ModelArtsDebugResourceSpecsMap = make(map[string]*FlavorInfo, len(t.FlavorInfo))
+		for _, spec := range t.FlavorInfo {
+			ModelArtsDebugResourceSpecsMap[spec.Value] = spec
+		}
+	}
+	if ModelArtsTrainResourceSpecsMap == nil || len(ModelArtsTrainResourceSpecsMap) == 0 {
+		t := ModelArtsFlavor{}
+		json.Unmarshal([]byte(setting.TrainJobFLAVORINFOS), &t)
+		ModelArtsTrainResourceSpecsMap = make(map[string]*FlavorInfo, len(t.Info))
+		for _, spec := range t.Info {
+			f := &FlavorInfo{
+				Value: spec.Code,
+				Desc:  spec.Value,
+			}
+			ModelArtsTrainResourceSpecsMap[spec.Value] = f
+		}
+	}
+	SpecsMapInitFlag = true
+}
+
+type ResourceAndFlavor struct {
+	ResourceSpec *ResourceSpec
+	FlavorInfo   *FlavorInfo
+}
+
+func NewResourceAndFlavor(resourceSpec *ResourceSpec, flavorInfo *FlavorInfo) *ResourceAndFlavor {
+	return &ResourceAndFlavor{
+		ResourceSpec: resourceSpec,
+		FlavorInfo:   flavorInfo,
+	}
+}
+
+func GetCloudbrainResourceSpec(jobType string, clusterType int, resourceSpecId int, flavorCode string) *ResourceAndFlavor {
+	if !SpecsMapInitFlag {
+		InitResourceSpecMap()
+	}
+	if clusterType == TypeCloudBrainOne {
+		switch jobType {
+		case string(JobTypeDebug):
+			return NewResourceAndFlavor(CloudbrainDebugResourceSpecsMap[resourceSpecId], nil)
+		case string(JobTypeTrain):
+			return NewResourceAndFlavor(CloudbrainTrainResourceSpecsMap[resourceSpecId], nil)
+		case string(JobTypeBenchmark):
+			return NewResourceAndFlavor(CloudbrainBenchmarkResourceSpecsMap[resourceSpecId], nil)
+
+		}
+	} else if clusterType == TypeCloudBrainTwo {
+		switch jobType {
+		case string(JobTypeDebug):
+			return NewResourceAndFlavor(nil, ModelArtsDebugResourceSpecsMap[flavorCode])
+		case string(JobTypeTrain):
+			return NewResourceAndFlavor(nil, ModelArtsTrainResourceSpecsMap[flavorCode])
+		case string(JobTypeInference):
+			return NewResourceAndFlavor(nil, ModelArtsTrainResourceSpecsMap[flavorCode])
+
+		}
+	}
+
+	return nil
+
+}