修改GPU创建模型

Signed-off-by: zouap <zouap@pcl.ac.cn>

modules/setting/setting.go

@@ -465,6 +465,7 @@ var (
 	MaxDuration        int64
 	TrainGpuTypes      string
 	TrainResourceSpecs string
+	MaxModelSize       int64
 
 	//benchmark config
 	IsBenchmarkEnabled     bool
@@ -1294,6 +1295,7 @@ func NewContext() {
 	MaxDuration = sec.Key("MAX_DURATION").MustInt64(14400)
 	TrainGpuTypes = sec.Key("TRAIN_GPU_TYPES").MustString("")
 	TrainResourceSpecs = sec.Key("TRAIN_RESOURCE_SPECS").MustString("")
+	MaxModelSize = sec.Key("MAX_MODEL_SIZE").MustInt64(500)
 
 	sec = Cfg.Section("benchmark")
 	IsBenchmarkEnabled = sec.Key("ENABLED").MustBool(false)

modules/storage/minio_ext.go

@@ -217,6 +217,25 @@ func GetOneLevelAllObjectUnderDirMinio(bucket string, prefixRootPath string, rel
 
 }
 
+func MinioGetFilesSize(bucketName string, Files []string) int64 {
+	_, core, err := getClients()
+	var fileTotalSize int64
+	fileTotalSize = 0
+	if err != nil {
+		log.Error("getClients failed:", err.Error())
+		return fileTotalSize
+	}
+	for _, file := range Files {
+		log.Info("file=" + file)
+		meta, err := core.StatObject(bucketName, file, miniov6.StatObjectOptions{})
+		if err != nil {
+			log.Info("Get file error:" + err.Error())
+		}
+		fileTotalSize += meta.Size
+	}
+	return fileTotalSize
+}
+
 func MinioCopyFiles(bucketName string, srcPath string, destPath string, Files []string) (int64, error) {
 	_, core, err := getClients()
 	var fileTotalSize int64

modules/storage/obs.go

@@ -264,6 +264,23 @@ func ObsModelDownload(JobName string, fileName string) (io.ReadCloser, error) {
 	}
 }
 
+func ObsGetFilesSize(srcBucket string, Files []string) int64 {
+	var fileTotalSize int64
+	for _, file := range Files {
+		log.Info("file=" + file)
+		out, err := ObsCli.GetObjectMetadata(&obs.GetObjectMetadataInput{
+			Bucket: srcBucket,
+			Key:    file,
+		})
+		if err != nil {
+			log.Info("Get File error, error=" + err.Error())
+			continue
+		}
+		fileTotalSize += out.ContentLength
+	}
+	return fileTotalSize
+}
+
 func ObsCopyManyFile(srcBucket string, srcPath string, destBucket string, destPath string, Files []string) (int64, error) {
 
 	var fileTotalSize int64

routers/repo/ai_model_manage.go

@@ -26,6 +26,7 @@ const (
 	tplModelInfo           = "repo/modelmanage/showinfo"
 	MODEL_LATEST           = 1
 	MODEL_NOT_LATEST       = 0
+	MODEL_MAX_SIZE         = 1 //1024 * 1024 * 1024
 )
 
 func saveModelByParameters(jobId string, versionName string, name string, version string, label string, description string, engine int, ctx *context.Context) error {
@@ -204,6 +205,17 @@ func downloadModelFromCloudBrainTwo(modelUUID string, jobName string, parentDir
 	if trainUrl != "" {
 		objectkey = strings.Trim(trainUrl[len(setting.Bucket)+1:], "/")
 	}
+	prefix := objectkey + "/"
+
+	filterFiles := strings.Split(modelSelectedFile, ";")
+	Files := make([]string, 0)
+	for _, shortFile := range filterFiles {
+		Files = append(Files, prefix+shortFile)
+	}
+	totalSize := storage.ObsGetFilesSize(setting.Bucket, Files)
+	if totalSize > setting.MaxModelSize*MODEL_MAX_SIZE {
+		return "", 0, errors.New("Cannot create model, as model is exceed " + fmt.Sprint(setting.MaxModelSize) + "G.")
+	}
 
 	modelDbResult, err := storage.GetOneLevelAllObjectUnderDir(setting.Bucket, objectkey, "")
 	log.Info("bucket=" + setting.Bucket + "  objectkey=" + objectkey)
@@ -212,12 +224,10 @@ func downloadModelFromCloudBrainTwo(modelUUID string, jobName string, parentDir
 		return "", 0, err
 	}
 	if len(modelDbResult) == 0 {
-		return "", 0, errors.New("cannot create model, as model is empty.")
+		return "", 0, errors.New("Cannot create model, as model is empty.")
 	}
 
-	prefix := objectkey + "/"
 	destKeyNamePrefix := Model_prefix + models.AttachmentRelativePath(modelUUID) + "/"
-	filterFiles := strings.Split(modelSelectedFile, ";")
 	size, err := storage.ObsCopyManyFile(setting.Bucket, prefix, setting.Bucket, destKeyNamePrefix, filterFiles)
 
 	dataActualPath := setting.Bucket + "/" + destKeyNamePrefix
@@ -232,6 +242,14 @@ func downloadModelFromCloudBrainOne(modelUUID string, jobName string, parentDir
 	bucketName := setting.Attachment.Minio.Bucket
 	log.Info("destKeyNamePrefix=" + destKeyNamePrefix + "   modelSrcPrefix=" + modelSrcPrefix + "  bucket=" + bucketName)
 	filterFiles := strings.Split(modelSelectedFile, ";")
+	Files := make([]string, 0)
+	for _, shortFile := range filterFiles {
+		Files = append(Files, modelSrcPrefix+shortFile)
+	}
+	totalSize := storage.MinioGetFilesSize(bucketName, Files)
+	if totalSize > setting.MaxModelSize*MODEL_MAX_SIZE {
+		return "", 0, errors.New("Cannot create model, as model is exceed " + fmt.Sprint(setting.MaxModelSize) + "G.")
+	}
 	size, err := storage.MinioCopyFiles(bucketName, modelSrcPrefix, destKeyNamePrefix, filterFiles)
 	if err == nil {
 		dataActualPath := bucketName + "/" + destKeyNamePrefix