aiforge/routers/repo/ai_model_manage.go

package repo

import (
	"archive/zip"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"

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

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
	if cloudType == models.TypeCloudBrainOne {
		modelPath, modelSize, err = downloadModelFromCloudBrainOne(id, aiTask.JobName, "")
		if err != nil {
			log.Info("download model from CloudBrainOne faild." + err.Error())
			ctx.Error(500, fmt.Sprintf("%v", err))
			return
		}
	} else if cloudType == models.TypeCloudBrainTwo {
		modelPath, err = downloadModelFromCloudBrainTwo(id)
		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 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
}

func downloadModelFromCloudBrainTwo(modelUUID string) (string, error) {
	dataActualPath := setting.Bucket + "/" +
		"aimodels/" +
		models.AttachmentRelativePath(modelUUID) +
		"/"
	return dataActualPath, nil
}

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

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

}

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

}