aiforge/routers/repo/ai_model_manage.go

package repo

import (
	"archive/zip"
	"errors"
	"fmt"
	"net/http"
	"path"
	"strings"

	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/context"
	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/setting"
	"code.gitea.io/gitea/modules/storage"
	uuid "github.com/satori/go.uuid"
)

const (
	Model_prefix           = "aimodels/"
	tplModelManageIndex    = "repo/modelmanage/index"
	tplModelManageDownload = "repo/modelmanage/download"
	MODEL_LATEST           = 1
	MODEL_NOT_LATEST       = 0
)

func SaveModelByParameters(jobId string, name string, version string, label string, description string, userId int64) error {
	aiTask, err := models.GetCloudbrainByJobID(jobId)
	if err != nil {
		log.Info("query task error." + err.Error())
		return err
	}

	uuid := uuid.NewV4()
	id := uuid.String()
	modelPath := id
	var lastNewModelId string
	var modelSize int64
	cloudType := models.TypeCloudBrainTwo

	log.Info("find task name:" + aiTask.JobName)
	aimodels := models.QueryModelByName(name, aiTask.RepoID)
	if len(aimodels) > 0 {
		for _, model := range aimodels {
			if model.New == MODEL_LATEST {
				lastNewModelId = model.ID
			}
		}
	}
	cloudType = aiTask.Type
	//download model zip   //train type
	if cloudType == models.TypeCloudBrainTwo {
		modelPath, modelSize, err = downloadModelFromCloudBrainTwo(id, aiTask.JobName, "")
		if err != nil {
			log.Info("download model from CloudBrainTwo faild." + err.Error())
			return err
		}
	}

	model := &models.AiModelManage{
		ID:           id,
		Version:      version,
		Label:        label,
		Name:         name,
		Description:  description,
		New:          MODEL_LATEST,
		Type:         cloudType,
		Path:         modelPath,
		Size:         modelSize,
		AttachmentId: aiTask.Uuid,
		RepoId:       aiTask.RepoID,
		UserId:       userId,
	}

	models.SaveModelToDb(model)

	if len(lastNewModelId) > 0 {
		//udpate status
		models.ModifyModelNewProperty(lastNewModelId, MODEL_NOT_LATEST)
	}

	log.Info("save model end.")

	return nil
}

func SaveModel(ctx *context.Context) {
	log.Info("save model start.")
	JobId := ctx.Query("JobId")
	name := ctx.Query("Name")
	version := ctx.Query("Version")
	label := ctx.Query("Label")
	description := ctx.Query("Description")

	err := SaveModelByParameters(JobId, name, version, label, description, ctx.User.ID)

	if err != nil {
		log.Info("save model error." + err.Error())
		ctx.Error(500, fmt.Sprintf("save model error. %v", err))
		return
	}

	log.Info("save model end.")
}

func downloadModelFromCloudBrainTwo(modelUUID string, jobName string, parentDir string) (string, int64, error) {

	objectkey := strings.TrimPrefix(path.Join(setting.TrainJobModelPath, jobName, setting.OutPutPath, parentDir), "/")
	modelDbResult, err := storage.GetOneLevelAllObjectUnderDir(setting.Bucket, objectkey, "")
	if err != nil {
		log.Info("get TrainJobListModel failed:", err)
		return "", 0, err
	}
	if len(modelDbResult) == 0 {
		return "", 0, errors.New("cannot create model, as model is empty.")
	}

	prefix := objectkey + "/"
	destKeyNamePrefix := Model_prefix + models.AttachmentRelativePath(modelUUID) + "/"

	size, err := storage.ObsCopyManyFile(setting.Bucket, prefix, setting.Bucket, destKeyNamePrefix)

	dataActualPath := setting.Bucket + "/" + destKeyNamePrefix
	return dataActualPath, size, nil
}

func DeleteModel(ctx *context.Context) {
	log.Info("delete model start.")
	id := ctx.Query("ID")
	err := DeleteModelByID(id)
	if err != nil {
		ctx.JSON(500, err.Error())
	} else {
		ctx.JSON(200, map[string]string{
			"result_code": "0",
		})
	}
}

func DeleteModelByID(id string) error {
	log.Info("delete model start. id=" + id)

	model, err := models.QueryModelById(id)
	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
			}
		}
		err = models.DeleteModelById(id)
		if err == nil { //find a model to change new
			if model.New == MODEL_LATEST {
				aimodels := models.QueryModelByName(model.Name, model.RepoId)
				if len(aimodels) > 0 {
					models.ModifyModelNewProperty(aimodels[0].ID, MODEL_LATEST)
				}
			}
		}
	}
	return err
}

func QueryModelByParameters(repoId int64, page int) ([]*models.AiModelManage, int64, error) {

	return models.QueryModel(&models.AiModelQueryOptions{
		ListOptions: models.ListOptions{
			Page:     page,
			PageSize: setting.UI.IssuePagingNum,
		},
		RepoID: repoId,
		Type:   -1,
		New:    MODEL_LATEST,
	})
}

func DownloadMultiModelFile(ctx *context.Context) {
	log.Info("DownloadMultiModelFile start.")
	id := ctx.Query("ID")
	log.Info("id=" + id)
	task, err := models.QueryModelById(id)
	if err != nil {
		log.Error("no such model!", err.Error())
		ctx.ServerError("no such model:", err)
		return
	}
	path := Model_prefix + models.AttachmentRelativePath(id) + "/"

	allFile, err := storage.GetAllObjectByBucketAndPrefix(setting.Bucket, path)
	if err == nil {
		returnFileName := task.Name + "_" + task.Version + ".zip"
		ctx.Resp.Header().Set("Content-Disposition", "attachment; filename="+returnFileName)
		ctx.Resp.Header().Set("Content-Type", "application/octet-stream")
		w := zip.NewWriter(ctx.Resp)
		defer w.Close()
		for _, oneFile := range allFile {
			if oneFile.IsDir {
				log.Info("zip dir name:" + oneFile.FileName)
			} else {
				log.Info("zip file name:" + oneFile.FileName)
				fDest, err := w.Create(oneFile.FileName)
				if err != nil {
					log.Info("create zip entry error, download file failed: %s\n", err.Error())
					ctx.ServerError("download file failed:", err)
					return
				}
				body, err := storage.ObsDownloadAFile(setting.Bucket, path+oneFile.FileName)
				if err != nil {
					log.Info("download file failed: %s\n", err.Error())
					ctx.ServerError("download file failed:", err)
					return
				} else {
					defer body.Close()
					p := make([]byte, 1024)
					var readErr error
					var readCount int
					// 读取对象内容
					for {
						readCount, readErr = body.Read(p)
						if readCount > 0 {
							fDest.Write(p[:readCount])
						}
						if readErr != nil {
							break
						}
					}
				}
			}
		}
	} else {
		log.Info("error,msg=" + err.Error())
		ctx.ServerError("no file to download.", err)
	}
}

func DownloadSingleModelFile(ctx *context.Context) {
	log.Info("DownloadSingleModelFile start.")
	id := ctx.Params(":ID")
	parentDir := ctx.Query("parentDir")
	fileName := ctx.Query("fileName")
	path := Model_prefix + models.AttachmentRelativePath(id) + "/" + parentDir + fileName

	if setting.PROXYURL != "" {
		body, err := storage.ObsDownloadAFile(setting.Bucket, path)
		if err != nil {
			log.Info("download error.")
		} else {
			defer body.Close()
			ctx.Resp.Header().Set("Content-Disposition", "attachment; filename="+fileName)
			ctx.Resp.Header().Set("Content-Type", "application/octet-stream")
			p := make([]byte, 1024)
			var readErr error
			var readCount int
			// 读取对象内容
			for {
				readCount, readErr = body.Read(p)
				if readCount > 0 {
					ctx.Resp.Write(p[:readCount])
					//fmt.Printf("%s", p[:readCount])
				}
				if readErr != nil {
					break
				}
			}
		}
	} else {
		url, err := storage.GetObsCreateSignedUrlByBucketAndKey(setting.Bucket, path)
		if err != nil {
			log.Error("GetObsCreateSignedUrl failed: %v", err.Error(), ctx.Data["msgID"])
			ctx.ServerError("GetObsCreateSignedUrl", err)
			return
		}
		http.Redirect(ctx.Resp, ctx.Req.Request, url, http.StatusMovedPermanently)
	}
}

func ShowSingleModel(ctx *context.Context) {
	id := ctx.Params(":ID")
	parentDir := ctx.Query("parentDir")
	log.Info("Show single ModelInfo start.id=" + id)
	task, err := models.QueryModelById(id)
	if err != nil {
		log.Error("no such model!", err.Error())
		ctx.ServerError("no such model:", err)
		return
	}
	log.Info("bucket=" + setting.Bucket + " key=" + task.Path[len(setting.Bucket)+1:])
	models, err := storage.GetOneLevelAllObjectUnderDir(setting.Bucket, task.Path[len(setting.Bucket)+1:], parentDir)
	if err != nil {
		log.Info("get model list failed:", err)
		ctx.ServerError("GetObsListObject:", err)
		return
	} else {
		log.Info("get model file,size=" + fmt.Sprint(len(models)))
	}

	ctx.Data["Dirs"] = models
	ctx.Data["task"] = task
	ctx.Data["ID"] = id
	ctx.HTML(200, tplModelManageDownload)
}

func ShowOneVersionOtherModel(ctx *context.Context) {
	repoId := ctx.Repo.Repository.ID
	name := ctx.Query("name")
	aimodels := models.QueryModelByName(name, repoId)
	if len(aimodels) > 0 {
		ctx.JSON(200, aimodels[1:])
	} else {
		ctx.JSON(200, aimodels)
	}
}

func ShowModelPageInfo(ctx *context.Context) {
	log.Info("ShowModelInfo start.")

	page := ctx.QueryInt("page")
	if page <= 0 {
		page = 1
	}
	repoId := ctx.Repo.Repository.ID
	Type := -1
	modelResult, count, err := models.QueryModel(&models.AiModelQueryOptions{
		ListOptions: models.ListOptions{
			Page:     page,
			PageSize: setting.UI.IssuePagingNum,
		},
		RepoID: repoId,
		Type:   Type,
		New:    MODEL_LATEST,
	})
	if err != nil {
		ctx.ServerError("Cloudbrain", err)
		return
	}

	pager := context.NewPagination(int(count), setting.UI.IssuePagingNum, page, 5)
	pager.SetDefaultParams(ctx)
	ctx.Data["Page"] = pager
	ctx.Data["PageIsCloudBrain"] = true
	ctx.Data["Tasks"] = modelResult

	ctx.HTML(200, tplModelManageIndex)
}

func ModifyModel(id string, description string) error {
	err := models.ModifyModelDescription(id, description)
	if err == nil {
		log.Info("modify success.")
	} else {
		log.Info("Failed to modify.id=" + id + "  desc=" + description + " error:" + err.Error())
	}
	return err
}

func ModifyModelInfo(ctx *context.Context) {
	log.Info("delete model start.")
	id := ctx.Query("ID")
	description := ctx.Query("Description")

	err := ModifyModel(id, description)

	if err == nil {
		ctx.HTML(200, "success")
	} else {
		ctx.HTML(500, "Failed.")
	}

}