aiforge/routers/repo/ai_model_manage.go

package repo

import (
	"archive/zip"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path"
	"path/filepath"
	"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"
)

func SaveModelByParameters(trainTaskId string, name string, version string, label string, description string, userId int64) {
	aiTask, err := models.GetCloudbrainByJobID(trainTaskId)
	if err != nil {
		log.Info("query task error." + err.Error())
		//ctx.Error(500, fmt.Sprintf("query cloud brain train task error. %v", err))
		return
	}

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

	log.Info("find task name:" + aiTask.JobName)
	aimodels := models.QueryModelByName(name, userId)
	if len(aimodels) > 0 {
		for _, model := range aimodels {
			if model.ID == model.Parent {
				parent = model.ID
			}
		}
	}
	cloudType = aiTask.Type
	//download model zip   //train type
	if cloudType == models.TypeCloudBrainTwo {
		modelPath, modelSize, err = downloadModelFromCloudBrainTwo(id, aiTask.JobName, "")
		if err == nil {

		} else {
			log.Info("download model from CloudBrainTwo faild." + err.Error())
			//ctx.Error(500, fmt.Sprintf("%v", err))
			return
		}
	}

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

	models.SaveModelToDb(model)

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

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

	aiTasks, _, err := models.Cloudbrains(&models.CloudbrainsOptions{
		JobID: trainTaskId,
	})
	if err != nil {
		log.Info("query task error." + err.Error())
		ctx.Error(500, fmt.Sprintf("query cloud brain train task error. %v", err))
		return
	}
	uuid := uuid.NewV4()
	id := uuid.String()
	modelPath := id
	parent := id
	var modelSize int64
	cloudType := models.TypeCloudBrainTwo

	if len(aiTasks) != 1 {
		log.Info("query task error. len=" + fmt.Sprint(len(aiTasks)))
		ctx.Error(500, fmt.Sprintf("query cloud brain train task error. %v", err))
		return
	}
	aiTask := aiTasks[0]
	log.Info("find task name:" + aiTask.JobName)
	aimodels := models.QueryModelByName(name, ctx.User.ID)
	if len(aimodels) > 0 {
		for _, model := range aimodels {
			if model.ID == model.Parent {
				parent = model.ID
			}
		}
	}
	cloudType = aiTask.Cloudbrain.Type
	//download model zip   //train type
	if cloudType == models.TypeCloudBrainTrainJob {
		modelPath, modelSize, err = downloadModelFromCloudBrainTwo(id, aiTask.JobName, "")
		if err == nil {

		} else {
			log.Info("download model from CloudBrainTwo faild." + err.Error())
			ctx.Error(500, fmt.Sprintf("%v", err))
			return
		}
	}

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

	models.SaveModelToDb(model)

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

func downloadModelFromCloudBrainTwo(modelUUID string, jobName string, parentDir string) (string, int64, error) {
	dataActualPath := setting.Bucket + "/" +
		"aimodels/" +
		models.AttachmentRelativePath(modelUUID) +
		"/"

	modelDbResult, err := storage.GetObsListObject(jobName, parentDir)
	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 := strings.TrimPrefix(path.Join(setting.TrainJobModelPath, jobName, setting.OutPutPath, parentDir), "/")
	for _, modelFile := range modelDbResult {
		destKeyName := "/aimodels/" + models.AttachmentRelativePath(modelUUID) + "/"
		log.Info("copy file, bucket=%s, src keyname=%s, dest keyname=%s,", setting.Bucket, prefix+modelFile.FileName, destKeyName)
		// err := storage.ObsCopyFile(setting.Bucket, modelFile.ParenDir+modelFile.FileName, setting.Bucket, dataActualPath+modelFile.FileName)
		// if err != nil {
		// 	log.Info("copy failed.")
		// }
	}

	return dataActualPath, 0, 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)
	return models.DeleteModelById(id)
}

func DownloadModel(ctx *context.Context) {
	log.Info("download model start.")

}

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,
	})
}

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

	page := ctx.QueryInt("page")
	if page <= 0 {
		page = 1
	}
	repoId := ctx.QueryInt64("repoId")
	Type := -1
	modelResult, count, err := models.QueryModel(&models.AiModelQueryOptions{
		ListOptions: models.ListOptions{
			Page:     page,
			PageSize: setting.UI.IssuePagingNum,
		},
		RepoID: repoId,
		Type:   Type,
	})
	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, "")
}

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

	modelActualPath := setting.Attachment.Minio.RealPath +
		setting.Attachment.Minio.Bucket + "/" +
		"aimodels/" +
		models.AttachmentRelativePath(modelUUID) +
		"/"
	os.MkdirAll(modelActualPath, 0755)
	zipFile := modelActualPath + "model.zip"

	modelDir := setting.JobPath + jobName + "/model/"

	dir, _ := ioutil.ReadDir(modelDir)
	if len(dir) == 0 {
		return "", 0, errors.New("cannot create model, as model is empty.")
	}

	err := zipDir(modelDir, zipFile)
	if err != nil {
		return "", 0, err
	}

	fi, err := os.Stat(zipFile)
	if err == nil {
		return modelActualPath, fi.Size(), nil
	} else {
		return "", 0, err
	}
}

func zipDir(dir, zipFile string) error {
	fz, err := os.Create(zipFile)
	if err != nil {
		log.Info("Create zip file failed: %s\n", err.Error())
		return err
	}
	defer fz.Close()

	w := zip.NewWriter(fz)
	defer w.Close()

	err = filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if !info.IsDir() {
			fDest, err := w.Create(path[len(dir)+1:])
			if err != nil {
				log.Info("Create failed: %s\n", err.Error())
				return err
			}
			fSrc, err := os.Open(path)
			if err != nil {
				log.Info("Open failed: %s\n", err.Error())
				return err
			}
			defer fSrc.Close()
			_, err = io.Copy(fDest, fSrc)
			if err != nil {
				log.Info("Copy failed: %s\n", err.Error())
				return err
			}
		}
		return nil
	})
	if err != nil {
		return err
	}
	return nil
}