From ec9643347c522f58f63ab90aa77c17adb65eee02 Mon Sep 17 00:00:00 2001 From: zouap Date: Wed, 22 Jun 2022 14:58:00 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E6=94=B9GPU=E5=88=9B=E5=BB=BA?= =?UTF-8?q?=E6=A8=A1=E5=9E=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: zouap --- modules/setting/setting.go | 2 ++ modules/storage/minio_ext.go | 19 +++++++++++++++++++ modules/storage/obs.go | 17 +++++++++++++++++ routers/repo/ai_model_manage.go | 24 +++++++++++++++++++++--- 4 files changed, 59 insertions(+), 3 deletions(-) diff --git a/modules/setting/setting.go b/modules/setting/setting.go index 5c87b68c5..df2e96928 100755 --- a/modules/setting/setting.go +++ b/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) diff --git a/modules/storage/minio_ext.go b/modules/storage/minio_ext.go index 89787307b..4ad83da82 100755 --- a/modules/storage/minio_ext.go +++ b/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 diff --git a/modules/storage/obs.go b/modules/storage/obs.go index cedef3a46..bbd40fe65 100755 --- a/modules/storage/obs.go +++ b/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 diff --git a/routers/repo/ai_model_manage.go b/routers/repo/ai_model_manage.go index 424f91517..1dd28194b 100644 --- a/routers/repo/ai_model_manage.go +++ b/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