提交代码。

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

models/ai_model_manage.go

@@ -14,6 +14,7 @@ import (
 type AiModelManage struct {
 	ID                string             `xorm:"pk" json:"id"`
 	Name              string             `xorm:"INDEX NOT NULL" json:"name"`
+	ModelType         int                `xorm:"NULL"`
 	Version           string             `xorm:"NOT NULL" json:"version"`
 	VersionCount      int                `xorm:"NOT NULL DEFAULT 0" json:"versionCount"`
 	New               int                `xorm:"NOT NULL" json:"new"`
@@ -287,6 +288,20 @@ func ModifyModelDescription(id string, description string) error {
 	return nil
 }
 
+func ModifyModelSize(id string, size int64) error {
+	var sess *xorm.Session
+	sess = x.ID(id)
+	defer sess.Close()
+	re, err := sess.Cols("size").Update(&AiModelManage{
+		Size: size,
+	})
+	if err != nil {
+		return err
+	}
+	log.Info("success to update size from  db.re=" + fmt.Sprint((re)))
+	return nil
+}
+
 func ModifyModelStatus(id string, modelSize int64, status int, modelPath string, statusDesc string) error {
 	var sess *xorm.Session
 	sess = x.ID(id)

routers/repo/ai_model_manage.go

@@ -244,6 +244,7 @@ func SaveLocalModel(ctx *context.Context) {
 	model := &models.AiModelManage{
 		ID:            id,
 		Version:       version,
+		ModelType:     1,
 		VersionCount:  len(aimodels) + 1,
 		Label:         label,
 		Name:          name,
@@ -291,6 +292,39 @@ func SaveLocalModel(ctx *context.Context) {
 	ctx.JSON(200, re)
 }
 
+func getSize(files []storage.FileInfo) int64 {
+	var size int64
+	for _, file := range files {
+		size += file.Size
+	}
+	return size
+}
+
+func UpdateModelSize(modeluuid string) {
+	model, err := models.QueryModelById(modeluuid)
+	if err == nil {
+		if model.Type == models.TypeCloudBrainOne {
+			if strings.HasPrefix(model.Path, setting.Bucket+"/"+Model_prefix) {
+				files, err := storage.GetAllObjectByBucketAndPrefixMinio(setting.Bucket, model.Path[len(setting.Bucket)+1:])
+				if err != nil {
+					log.Info("Failed to query model size from minio. id=" + modeluuid)
+				}
+				size := getSize(files)
+				models.ModifyModelSize(modeluuid, size)
+			}
+		} else if model.Type == models.TypeCloudBrainTwo {
+			if strings.HasPrefix(model.Path, setting.Bucket+"/"+Model_prefix) {
+				files, err := storage.GetAllObjectByBucketAndPrefix(setting.Bucket, model.Path[len(setting.Bucket)+1:])
+				if err != nil {
+					log.Info("Failed to query model size from obs. id=" + modeluuid)
+				}
+				size := getSize(files)
+				models.ModifyModelSize(modeluuid, size)
+			}
+		}
+	}
+}
+
 func SaveModel(ctx *context.Context) {
 	if !ctx.Repo.CanWrite(models.UnitTypeModelManage) {
 		ctx.Error(403, ctx.Tr("repo.model_noright"))
@@ -420,13 +454,24 @@ func deleteModelByID(ctx *context.Context, id string) error {
 	}
 	if err == nil {
 		log.Info("bucket=" + setting.Bucket + " path=" + model.Path)
-		if strings.HasPrefix(model.Path, setting.Bucket+"/"+Model_prefix) {
-			err := storage.ObsRemoveObject(setting.Bucket, model.Path[len(setting.Bucket)+1:])
-			if err != nil {
-				log.Info("Failed to delete model. id=" + id)
-				return err
+		if model.Type == models.TypeCloudBrainOne {
+			if strings.HasPrefix(model.Path, setting.Bucket+"/"+Model_prefix) {
+				err := storage.Attachments.DeleteDir(model.Path[len(setting.Bucket)+1:])
+				if err != nil {
+					log.Info("Failed to delete model. id=" + id)
+					return err
+				}
+			}
+		} else if model.Type == models.TypeCloudBrainTwo {
+			if strings.HasPrefix(model.Path, setting.Bucket+"/"+Model_prefix) {
+				err := storage.ObsRemoveObject(setting.Bucket, model.Path[len(setting.Bucket)+1:])
+				if err != nil {
+					log.Info("Failed to delete model. id=" + id)
+					return err
+				}
 			}
 		}
+
 		err = models.DeleteModelById(id)
 		if err == nil { //find a model to change new
 			aimodels := models.QueryModelByName(model.Name, model.RepoId)

routers/repo/attachment.go

@@ -951,6 +951,12 @@ func CompleteMultipart(ctx *context.Context) {
 	}
 
 	if scene == Attachment_model {
+		//更新模型大小信息
+		UpdateModelSize(modeluuid)
+
+		ctx.JSON(200, map[string]string{
+			"result_code": "0",
+		})
 
 	} else {
 		dataset, _ := models.GetDatasetByID(ctx.QueryInt64("dataset_id"))